smol_mpc/utils/
prg.rs

1//! Implementation of a PRG based on AES-CTR.
2//!
3//! The idea is that a pseudo-random block of bits is generated by computing
4//!    $$\textsf{Block} = \textsf{AES}(\textsf{Seed}, \textsf{CurrentCounter})$$
5//! where $\textsf{CurrentCounter} = \textsf{Nonce} \Vert \textsf{Counter}$, and the counter
6//! is increased when a new block is generated.
7//!
8//! This implementation is based on the one used in [Secure Computation Library].
9//!
10//! [Secure Computation Library]: https://github.com/anderspkd/secure-computation-library/blob/master/include/scl/util/prg.h
11
12use aes::cipher::{KeyIvInit, StreamCipher};
13use std::vec;
14
15type Aes128Ctr64LE = ctr::Ctr64LE<aes::Aes128>;
16
17/// Defines a pseudo-random number generator.
18pub struct Prg {
19    seed: Vec<u8>,
20    counter: u64,
21}
22
23impl Prg {
24    const PRG_NONCE: u64 = 0x0123456789ABCDEF;
25    const PRG_INITIAL_COUNTER: u64 = 0;
26
27    // All the lengths are in bytes
28    const KEY_LEN: usize = 16;
29    const IV_LEN: usize = 16;
30    const BLOCK_LEN: usize = 16;
31
32    /// Creates a new PRG.
33    ///
34    /// The seed can be provided or generated by default
35    /// filled with zeros. If a seed is provided, then there are two cases. If
36    /// the key is longer than the key length for the encryption scheme, it is
37    /// cropped to fit in the specification. If the key is shorter than the
38    /// expected key, the key is padded with zeros. The key will be divided in
39    /// two halves to generate the encryption key and the initialization vector
40    /// for the CTR mode.
41    pub fn new(seed: Option<Vec<u8>>) -> Prg {
42        let cropped_seed = if let Some(mut value_seed) = seed {
43            if value_seed.len() > Self::KEY_LEN + Self::IV_LEN {
44                value_seed[0..Self::KEY_LEN + Self::IV_LEN].to_vec()
45            } else {
46                let original_seed_length = value_seed.len();
47                let mut appended_seed = Vec::new();
48                appended_seed.append(&mut value_seed);
49                appended_seed.append(&mut vec![
50                    0;
51                    Self::KEY_LEN + Self::IV_LEN - original_seed_length
52                ]);
53                appended_seed
54            }
55        } else {
56            vec![0; Self::KEY_LEN + Self::IV_LEN]
57        };
58
59        let counter = Self::PRG_INITIAL_COUNTER;
60
61        let mut prg = Prg {
62            seed: cropped_seed,
63            counter,
64        };
65        prg.init();
66        prg
67    }
68
69    /// Initializes the PRG.
70    pub fn init(&mut self) {
71        self.counter = Self::PRG_INITIAL_COUNTER;
72    }
73
74    /// Resets the PRG.
75    pub fn reset(&mut self) {
76        self.init()
77    }
78
79    /// Returns the current state of the counter in the PRG.
80    pub fn counter(&self) -> u64 {
81        self.counter
82    }
83
84    /// Generates a stream of random bytes.
85    ///
86    /// The method divides the seed into two halves: the first part will be used
87    /// as the key for the AES encryption and the second part will be used as
88    /// the initialization vector for the encryption.
89    pub fn next(&mut self, n_bytes: usize) -> Vec<u8> {
90        if n_bytes == 0 {
91            return Vec::new();
92        }
93
94        // Compute the number of blocks needed
95        let mut n_blocks = n_bytes / Self::BLOCK_LEN;
96        if n_bytes % Self::BLOCK_LEN != 0 {
97            n_blocks += 1;
98        }
99
100        let key = &self.seed[0..Self::KEY_LEN];
101        let iv = &self.seed[Self::KEY_LEN..];
102
103        let mut cipher = Aes128Ctr64LE::new(key.into(), iv.into());
104
105        let mut out = Vec::new();
106        for _ in 0..n_blocks {
107            let mut buffer = [Self::PRG_NONCE.to_ne_bytes(), self.counter.to_ne_bytes()].concat();
108            cipher.apply_keystream(&mut buffer);
109            out.append(&mut buffer);
110
111            self.counter += 1;
112        }
113
114        out[..n_bytes].to_vec()
115    }
116}