cryprot_core/
aes_rng.rs

1//! RNG based on AES in CTR mode.
2//!
3//! This implementation is based on the implementation given in the
4//! [scuttlebutt](https://github.com/GaloisInc/swanky/blob/4455754abadee07f168079ac45ef33535b0df27d/scuttlebutt/src/rand_aes.rs)
5//! crate. Instead of using an own AES implementation, [`AesRng`](`AesRng`) uses
6//! the [aes](`aes`) crate.
7//!
8//! On platforms wwith hardware accelerated AES instructions, the [`AesRng`] can
9//! generate multiple GiB of random data per second. Make sure to compile with
10//! the `aes` target feature enabled to have optimal performance without runtime
11//! detection of the feature.
12use std::mem;
13
14use aes::{
15    Aes128,
16    cipher::{BlockCipherEncrypt, KeyInit},
17};
18use rand::{CryptoRng, Rng, RngCore, SeedableRng};
19use rand_core::block::{BlockRng, BlockRngCore, CryptoBlockRng};
20
21use crate::{AES_PAR_BLOCKS, Block};
22
23// TODO i think softspoken ot has some implementation performance optimizations
24// see sect 7 https://eprint.iacr.org/2022/192.pdf
25
26/// This uses AES in a counter-mode to implement a PRG. TODO: Citation for
27/// why/when this is secure.
28#[derive(Clone, Debug)]
29pub struct AesRng(BlockRng<AesRngCore>);
30
31impl RngCore for AesRng {
32    #[inline]
33    fn next_u32(&mut self) -> u32 {
34        self.0.next_u32()
35    }
36
37    #[inline]
38    fn next_u64(&mut self) -> u64 {
39        self.0.next_u64()
40    }
41
42    #[inline]
43    fn fill_bytes(&mut self, dest: &mut [u8]) {
44        let block_size = mem::size_of::<aes::Block>();
45        let block_len = dest.len() / block_size * block_size;
46        let (block_bytes, rest_bytes) = dest.split_at_mut(block_len);
47        // fast path so we don't unnecessarily copy u32 from BlockRngCore::generate into
48        // dest
49        let blocks = bytemuck::cast_slice_mut::<_, aes::Block>(block_bytes);
50        for chunk in blocks.chunks_mut(AES_PAR_BLOCKS) {
51            for block in chunk.iter_mut() {
52                *block = aes::cipher::Array(self.0.core.state.to_le_bytes());
53                self.0.core.state += 1;
54            }
55            self.0.core.aes.encrypt_blocks(chunk);
56        }
57        // handle the tail
58        self.0.fill_bytes(rest_bytes)
59    }
60}
61
62impl SeedableRng for AesRng {
63    type Seed = Block;
64
65    #[inline]
66    fn from_seed(seed: Self::Seed) -> Self {
67        AesRng(BlockRng::<AesRngCore>::from_seed(seed))
68    }
69}
70
71impl CryptoRng for AesRng {}
72
73impl AesRng {
74    /// Create a new random number generator using a random seed from
75    /// `rand::random`.
76    #[inline]
77    pub fn new() -> Self {
78        let seed = rand::random::<Block>();
79        AesRng::from_seed(seed)
80    }
81
82    /// Create a new RNG using a random seed from this one.
83    #[inline]
84    pub fn fork(&mut self) -> Self {
85        let seed = self.random::<Block>();
86        AesRng::from_seed(seed)
87    }
88}
89
90impl Default for AesRng {
91    #[inline]
92    fn default() -> Self {
93        Self::new()
94    }
95}
96
97/// The core of `AesRng`, used with `BlockRng`.
98#[derive(Clone)]
99pub struct AesRngCore {
100    aes: Aes128,
101    state: u128,
102}
103
104impl std::fmt::Debug for AesRngCore {
105    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
106        write!(f, "AesRngCore {{}}")
107    }
108}
109
110impl BlockRngCore for AesRngCore {
111    type Item = u32;
112    // This is equivalent to `[Block; 9]`
113    type Results = hidden::ParBlockWrapper;
114
115    // Compute `E(state)` nine times, where `state` is a counter.
116    #[inline]
117    fn generate(&mut self, results: &mut Self::Results) {
118        let blocks = bytemuck::cast_slice_mut::<_, aes::Block>(results.as_mut());
119        blocks.iter_mut().for_each(|blk| {
120            // aes::Block is a type alias to Array, but type aliases can't be used as
121            // constructors
122            *blk = aes::cipher::Array(self.state.to_le_bytes());
123            self.state += 1;
124        });
125        self.aes.encrypt_blocks(blocks);
126    }
127}
128
129mod hidden {
130    /// Equivalent to [aes::Block; 9] (which is the parralel block size for the
131    /// aes-ni backend). Since size 36 arrays don't impl Default we write a
132    /// wrapper.
133    #[derive(Copy, Clone)]
134    pub struct ParBlockWrapper([u32; 36]);
135
136    impl Default for ParBlockWrapper {
137        fn default() -> Self {
138            Self([0; 36])
139        }
140    }
141
142    impl AsMut<[u32]> for ParBlockWrapper {
143        fn as_mut(&mut self) -> &mut [u32] {
144            &mut self.0
145        }
146    }
147
148    impl AsRef<[u32]> for ParBlockWrapper {
149        fn as_ref(&self) -> &[u32] {
150            &self.0
151        }
152    }
153}
154
155impl SeedableRng for AesRngCore {
156    type Seed = Block;
157
158    #[inline]
159    fn from_seed(seed: Self::Seed) -> Self {
160        let aes = Aes128::new(&seed.into());
161        AesRngCore {
162            aes,
163            state: Default::default(),
164        }
165    }
166}
167
168impl CryptoBlockRng for AesRngCore {}
169
170impl From<AesRngCore> for AesRng {
171    #[inline]
172    fn from(core: AesRngCore) -> Self {
173        AesRng(BlockRng::new(core))
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180
181    #[test]
182    fn test_generate() {
183        let mut rng = AesRng::new();
184        let a = rng.random::<[Block; 8]>();
185        let b = rng.random::<[Block; 8]>();
186        assert_ne!(a, b);
187    }
188}