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}