Skip to main content

tfhe_csprng/generators/
mod.rs

1//! A module containing random generators objects.
2//!
3//! See [crate-level](`crate`) explanations.
4use crate::generators::aes_ctr::{AesCtrParams, TableIndex};
5use std::error::Error;
6use std::fmt::{Display, Formatter};
7
8/// The number of children created when a generator is forked.
9#[derive(Debug, Copy, Clone)]
10pub struct ChildrenCount(pub u64);
11
12/// The number of bytes each child can generate, when a generator is forked.
13#[derive(Debug, Copy, Clone)]
14pub struct BytesPerChild(pub u64);
15
16/// A structure representing the number of bytes between two table indices.
17#[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq)]
18pub struct ByteCount(pub u128);
19
20/// Multiplies two u64 values without overflow, returning the full 128-bit product.
21pub(crate) fn widening_mul(a: u64, b: u64) -> u128 {
22    (a as u128) * (b as u128)
23}
24
25/// An error occurring during a generator fork.
26#[derive(Debug)]
27pub enum ForkError {
28    ForkTooLarge,
29    ZeroChildrenCount,
30    ZeroBytesPerChild,
31}
32
33impl Display for ForkError {
34    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
35        match self {
36            ForkError::ForkTooLarge => {
37                write!(
38                    f,
39                    "The children generators would output bytes after the parent bound. "
40                )
41            }
42            ForkError::ZeroChildrenCount => {
43                write!(
44                    f,
45                    "The number of children in the fork must be greater than zero."
46                )
47            }
48            ForkError::ZeroBytesPerChild => {
49                write!(
50                    f,
51                    "The number of bytes per child must be greater than zero."
52                )
53            }
54        }
55    }
56}
57impl Error for ForkError {}
58
59/// A trait for cryptographically secure pseudo-random generators.
60///
61/// See the [crate-level](#crate) documentation for details.
62pub trait RandomGenerator: Iterator<Item = u8> {
63    /// The iterator over children generators, returned by `try_fork` in case of success.
64    type ChildrenIter: Iterator<Item = Self>;
65
66    /// Creates a new generator from parameters (seed + optional state).
67    ///
68    /// This operation is usually costly to perform, as the aes round keys need to be generated from
69    /// the seed.
70    fn new(params: impl Into<AesCtrParams>) -> Self;
71
72    /// Returns the number of bytes that can still be outputted by the generator before reaching its
73    /// bound.
74    ///
75    /// Note:
76    /// -----
77    ///
78    /// A fresh generator can generate 2¹³² bytes. Unfortunately, no rust integer type in is able
79    /// to encode such a large number. Consequently [`ByteCount`] uses the largest integer type
80    /// available to encode this value: the `u128` type. For this reason, this method does not
81    /// effectively return the number of remaining bytes, but instead
82    /// `min(2¹²⁸-1, remaining_bytes)`.
83    fn remaining_bytes(&self) -> ByteCount;
84
85    /// Returns the table index of the next byte to be generated
86    fn next_table_index(&self) -> Option<TableIndex>;
87
88    /// Returns the next byte of the stream, if the generator did not yet reach its bound.
89    fn next_byte(&mut self) -> Option<u8> {
90        self.next()
91    }
92
93    /// Tries to fork the generator into an iterator of `n_children` new generators, each able to
94    /// output `n_bytes` bytes.
95    ///
96    /// Note:
97    /// -----
98    ///
99    /// To be successful, the number of remaining bytes for the parent generator must be larger than
100    /// `n_children*n_bytes`.
101    fn try_fork(
102        &mut self,
103        n_children: ChildrenCount,
104        n_bytes: BytesPerChild,
105    ) -> Result<Self::ChildrenIter, ForkError>;
106}
107
108/// A trait extending [`RandomGenerator`] to the parallel iterators of `rayon`.
109#[cfg(feature = "parallel")]
110pub trait ParallelRandomGenerator: RandomGenerator + Send {
111    /// The iterator over children generators, returned by `par_try_fork` in case of success.
112    type ParChildrenIter: rayon::prelude::IndexedParallelIterator<Item = Self>;
113
114    /// Tries to fork the generator into a parallel iterator of `n_children` new generators, each
115    /// able to output `n_bytes` bytes.
116    ///
117    /// Note:
118    /// -----
119    ///
120    /// To be successful, the number of remaining bytes for the parent generator must be larger than
121    /// `n_children*n_bytes`.
122    fn par_try_fork(
123        &mut self,
124        n_children: ChildrenCount,
125        n_bytes: BytesPerChild,
126    ) -> Result<Self::ParChildrenIter, ForkError>;
127}
128
129pub mod aes_ctr;
130pub mod backward_compatibility;
131
132mod implem;
133pub use implem::*;
134
135pub mod default;
136/// Convenience alias for the most efficient CSPRNG implementation available.
137pub use default::DefaultRandomGenerator;
138
139#[cfg(test)]
140#[allow(unused)] // to please clippy when tests are not activated
141pub mod generator_generic_test {
142    use super::*;
143    use crate::seeders::{Seed, XofSeed};
144    use rand::Rng;
145
146    const REPEATS: usize = 1_000;
147
148    fn any_seed() -> impl Iterator<Item = Seed> {
149        std::iter::repeat_with(|| Seed(rand::thread_rng().gen()))
150    }
151
152    fn some_children_count() -> impl Iterator<Item = ChildrenCount> {
153        std::iter::repeat_with(|| ChildrenCount(rand::thread_rng().gen::<u64>() % 16 + 1))
154    }
155
156    fn some_bytes_per_child() -> impl Iterator<Item = BytesPerChild> {
157        std::iter::repeat_with(|| BytesPerChild(rand::thread_rng().gen::<u64>() % 128 + 1))
158    }
159
160    /// Checks that the PRNG roughly generates uniform numbers.
161    ///
162    /// To do that, we perform an histogram of the occurrences of each byte value, over a fixed
163    /// number of samples and check that the empirical probabilities of the bins are close to
164    /// the theoretical probabilities.
165    pub fn test_roughly_uniform<G: RandomGenerator>() {
166        // Number of bins to use for the histogram.
167        const N_BINS: usize = u8::MAX as usize + 1;
168        // Number of samples to use for the histogram.
169        let n_samples = 10_000_000_usize;
170        // Theoretical probability of a each bins.
171        let expected_prob: f64 = 1. / N_BINS as f64;
172        // Absolute error allowed on the empirical probabilities.
173        // This value was tuned to make the test pass on an arguably correct state of
174        // implementation. 10^-4 precision is arguably pretty fine for this rough test, but it would
175        // be interesting to improve this test.
176        let precision = 10f64.powi(-3);
177
178        for _ in 0..REPEATS {
179            // We instantiate a new generator.
180            let seed = any_seed().next().unwrap();
181            let mut generator = G::new(seed);
182            // We create a new histogram
183            let mut counts = [0usize; N_BINS];
184            // We fill the histogram.
185            for _ in 0..n_samples {
186                counts[generator.next_byte().unwrap() as usize] += 1;
187            }
188            // We check that the empirical probabilities are close enough to the theoretical one.
189            counts
190                .iter()
191                .map(|a| (*a as f64) / (n_samples as f64))
192                .for_each(|a| assert!((a - expected_prob).abs() < precision))
193        }
194    }
195
196    /// Checks that given a state and a key, the PRNG is determinist.
197    pub fn test_generator_determinism<G: RandomGenerator>() {
198        for _ in 0..REPEATS {
199            let seed = any_seed().next().unwrap();
200            let mut first_generator = G::new(seed);
201            let mut second_generator = G::new(seed);
202            for _ in 0..1024 {
203                assert_eq!(first_generator.next(), second_generator.next());
204            }
205        }
206    }
207
208    /// Checks that forks returns a bounded child, and that the proper number of bytes can be
209    /// generated.
210    pub fn test_fork_children<G: RandomGenerator>() {
211        for _ in 0..REPEATS {
212            let ((seed, n_children), n_bytes) = any_seed()
213                .zip(some_children_count())
214                .zip(some_bytes_per_child())
215                .next()
216                .unwrap();
217            let mut gen = G::new(seed);
218            let mut bounded = gen.try_fork(n_children, n_bytes).unwrap().next().unwrap();
219            assert_eq!(bounded.remaining_bytes(), ByteCount(n_bytes.0 as u128));
220            for _ in 0..n_bytes.0 {
221                bounded.next().unwrap();
222            }
223
224            // Assert we are at the bound
225            assert!(bounded.next().is_none());
226        }
227    }
228
229    /// Checks that a bounded prng returns none when exceeding the allowed number of bytes.
230    ///
231    /// To properly check for panic use `#[should_panic(expected = "expected test panic")]` as an
232    /// attribute on the test function.
233    pub fn test_bounded_none_should_panic<G: RandomGenerator>() {
234        let ((seed, n_children), n_bytes) = any_seed()
235            .zip(some_children_count())
236            .zip(some_bytes_per_child())
237            .next()
238            .unwrap();
239        let mut gen = G::new(seed);
240        let mut bounded = gen.try_fork(n_children, n_bytes).unwrap().next().unwrap();
241        assert_eq!(bounded.remaining_bytes(), ByteCount(n_bytes.0 as u128));
242        for _ in 0..n_bytes.0 {
243            assert!(bounded.next().is_some());
244        }
245
246        // One call too many, should panic
247        bounded.next().ok_or("expected test panic").unwrap();
248    }
249
250    pub fn test_vectors<G: RandomGenerator>() {
251        // Number of random bytes to generate,
252        // this should be 2 batch worth of aes calls (where a batch is 8 aes)
253        const N_BYTES: usize = 16 * 2 * 8;
254
255        const EXPECTED_BYTE: [u8; N_BYTES] = [
256            220, 14, 216, 93, 249, 97, 26, 187, 114, 73, 205, 209, 104, 197, 70, 126, 250, 235, 1,
257            136, 141, 46, 146, 174, 231, 14, 204, 28, 99, 139, 246, 214, 112, 253, 151, 34, 114,
258            235, 7, 76, 37, 36, 154, 226, 148, 68, 238, 117, 87, 212, 183, 174, 200, 222, 153, 62,
259            48, 166, 134, 27, 97, 230, 206, 78, 128, 151, 166, 15, 156, 120, 158, 35, 41, 121, 55,
260            180, 184, 108, 160, 33, 208, 255, 147, 246, 159, 10, 239, 6, 103, 124, 123, 83, 72,
261            189, 237, 225, 36, 30, 151, 134, 94, 211, 181, 108, 239, 137, 18, 246, 237, 233, 59,
262            61, 24, 111, 198, 76, 92, 86, 129, 171, 50, 124, 2, 72, 143, 160, 223, 32, 187, 175,
263            239, 111, 51, 85, 110, 134, 45, 193, 113, 247, 249, 78, 230, 103, 123, 66, 48, 31, 169,
264            228, 140, 202, 168, 202, 199, 147, 89, 135, 104, 254, 198, 72, 31, 103, 236, 207, 138,
265            24, 100, 230, 168, 233, 214, 130, 195, 0, 25, 220, 136, 128, 173, 40, 154, 116, 87,
266            114, 187, 170, 150, 131, 163, 155, 98, 217, 198, 238, 178, 165, 214, 168, 252, 107,
267            123, 214, 33, 17, 114, 35, 23, 172, 145, 5, 39, 16, 33, 92, 163, 132, 240, 167, 128,
268            226, 165, 80, 9, 153, 252, 139, 0, 139, 0, 54, 188, 253, 141, 2, 78, 97, 53, 214, 173,
269            155, 84, 98, 51, 70, 110, 91, 181, 229, 231, 27, 225, 185, 143, 63,
270        ];
271
272        let seed = Seed(1u128);
273
274        let mut rng = G::new(seed);
275        let bytes = rng.take(N_BYTES).collect::<Vec<_>>();
276        assert_eq!(bytes, EXPECTED_BYTE);
277    }
278
279    pub fn test_vectors_xof_seed<G: RandomGenerator>() {
280        // Number of random bytes to generate,
281        // this should be 2 batch worth of aes calls (where a batch is 8 aes)
282        const N_BYTES: usize = 16 * 2 * 8;
283
284        const EXPECTED_BYTE: [u8; N_BYTES] = [
285            181, 134, 231, 117, 200, 60, 174, 158, 95, 80, 64, 236, 147, 204, 196, 251, 198, 110,
286            155, 74, 69, 162, 251, 224, 46, 46, 83, 209, 224, 89, 108, 68, 240, 37, 16, 109, 194,
287            92, 3, 164, 21, 167, 224, 205, 31, 90, 178, 59, 150, 142, 238, 113, 144, 181, 118, 160,
288            72, 187, 38, 29, 61, 189, 229, 66, 22, 4, 38, 210, 63, 232, 182, 115, 49, 96, 6, 120,
289            226, 40, 51, 144, 59, 136, 224, 252, 195, 50, 250, 134, 45, 149, 220, 32, 27, 35, 225,
290            190, 73, 161, 182, 250, 149, 153, 131, 220, 143, 181, 152, 187, 25, 62, 197, 24, 10,
291            142, 57, 172, 15, 17, 244, 242, 232, 51, 50, 244, 85, 58, 69, 28, 113, 151, 143, 138,
292            166, 198, 16, 210, 46, 234, 138, 32, 124, 98, 167, 141, 251, 60, 13, 158, 106, 29, 86,
293            63, 73, 42, 138, 174, 195, 192, 72, 122, 74, 54, 134, 107, 144, 241, 12, 33, 70, 27,
294            116, 154, 123, 1, 252, 141, 73, 79, 30, 162, 43, 57, 8, 99, 62, 222, 117, 232, 147, 81,
295            189, 54, 17, 233, 33, 41, 132, 155, 246, 185, 189, 17, 77, 32, 107, 134, 61, 174, 64,
296            174, 80, 229, 239, 243, 143, 152, 249, 254, 125, 42, 0, 170, 253, 34, 57, 100, 82, 244,
297            9, 101, 126, 138, 218, 215, 55, 58, 177, 154, 5, 28, 113, 89, 123, 129, 254, 212, 191,
298            162, 44, 120, 67, 241, 157, 31, 162, 113,
299        ];
300
301        let seed = 1u128;
302        let xof_seed = XofSeed::new_u128(seed, *b"abcdefgh");
303
304        let mut rng = G::new(xof_seed);
305        let bytes = rng.take(N_BYTES).collect::<Vec<_>>();
306        assert_eq!(bytes, EXPECTED_BYTE);
307    }
308
309    pub fn test_vectors_xof_seed_bytes<G: RandomGenerator>() {
310        // Number of random bytes to generate,
311        // this should be 2 batch worth of aes calls (where a batch is 8 aes)
312        const N_BYTES: usize = 16 * 2 * 8;
313
314        const EXPECTED_BYTE: [u8; N_BYTES] = [
315            122, 21, 82, 236, 82, 18, 196, 63, 129, 54, 134, 70, 114, 199, 200, 11, 5, 52, 170,
316            218, 49, 127, 45, 5, 252, 214, 82, 127, 196, 241, 83, 161, 79, 139, 183, 33, 122, 126,
317            177, 23, 36, 161, 122, 7, 112, 237, 154, 195, 90, 202, 218, 64, 90, 86, 190, 139, 169,
318            192, 105, 248, 220, 126, 133, 60, 124, 81, 72, 183, 238, 253, 138, 141, 144, 167, 168,
319            94, 19, 172, 92, 235, 113, 185, 31, 150, 143, 165, 220, 115, 83, 180, 1, 10, 130, 140,
320            32, 74, 132, 76, 22, 120, 126, 68, 154, 95, 61, 202, 79, 126, 38, 217, 181, 243, 6,
321            218, 75, 232, 235, 194, 255, 254, 184, 18, 122, 51, 222, 61, 167, 175, 97, 188, 186,
322            217, 105, 72, 205, 130, 3, 204, 157, 252, 27, 20, 212, 136, 70, 65, 215, 164, 130, 242,
323            107, 214, 150, 211, 59, 92, 13, 148, 219, 96, 181, 5, 38, 170, 48, 218, 111, 131, 246,
324            102, 169, 17, 182, 253, 41, 209, 185, 79, 245, 30, 142, 192, 127, 78, 178, 68, 223, 89,
325            210, 27, 84, 164, 163, 216, 188, 190, 128, 154, 224, 160, 53, 249, 10, 250, 95, 160,
326            94, 28, 41, 34, 254, 232, 137, 185, 82, 82, 192, 74, 197, 19, 46, 180, 169, 182, 216,
327            221, 127, 196, 185, 156, 82, 32, 133, 97, 140, 183, 67, 37, 110, 31, 210, 197, 27, 81,
328            197, 132, 136, 98, 78, 218, 252, 247, 239, 205, 21, 166,
329        ];
330
331        let seed = vec![
332            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23,
333            24, 25, 26, 27, 28, 29, 30, 31,
334        ];
335        let xof_seed = XofSeed::new(seed, *b"abcdefgh");
336
337        let mut rng = G::new(xof_seed);
338        let bytes = rng.take(N_BYTES).collect::<Vec<_>>();
339        assert_eq!(bytes, EXPECTED_BYTE);
340    }
341}