1
//! Code for deterministic and/or reproducible use of PRNGs in tests.
2
//!
3
//! Often in testing we want to test a random scenario, but we want to be sure
4
//! of our ability to reproduce the scenario if the test fails.
5
//!
6
//! To achieve this,  just have your test use [`testing_rng()`] in place of
7
//! [`rand::rng()`].  Then the test will (by default) choose a new random
8
//! seed for every run, and print that seed to standard output.  If the test
9
//! fails, the seed will be displayed as part of the failure message, and you
10
//! will be able to use it to recreate the same PRNG seed as the one that caused
11
//! the failure.
12
//!
13
//! If you're running your tests in a situation where deterministic behavior is
14
//! key, you can also enable this via the environment.
15
//!
16
//! The run-time behavior is controlled using the `ARTI_TEST_PRNG` variable; you
17
//! can set it to any of the following:
18
//!   * `random` for a randomly seeded PRNG. (This is the default).
19
//!   * `deterministic` for an arbitrary seed that is the same on every run of
20
//!     the program. (You can use this in cases where even a tiny chance of
21
//!     stochastic behavior in your tests is unacceptable.)
22
//!   * A hexadecimal string, to specify a given seed to reuse from a previous
23
//!     test run.
24
//!
25
//! # WARNING
26
//!
27
//! This is for testing only!  Never ever use it in non-testing code.  Doing so
28
//! may compromise your security.
29
//!
30
//! You may wish to use clippy's `disallowed-methods` lint to ensure you aren't
31
//! using it outside of your tests.
32
//!
33
//! # Examples
34
//!
35
//! Here's a simple example of a test that verifies that integer sorting works
36
//! correctly by shuffling a short sequence and then re-sorting it.
37
//!
38
//! ```
39
//! use tor_basic_utils::test_rng::testing_rng;
40
//! use rand::{seq::SliceRandom};
41
//! let mut rng = testing_rng();
42
//!
43
//! let mut v = vec![-10, -3, 0, 1, 2, 3];
44
//! v.shuffle(&mut rng);
45
//! v.sort();
46
//! assert_eq!(&v, &[-10, -3, 0, 1, 2, 3])
47
//! ```
48
//!
49
//! Here's a trickier example of how you might write a test to override the
50
//! default behavior.  (For example, you might want to do this if the test is
51
//! unreliable and you don't have time to hunt down the issues.)
52
//!
53
//! ```
54
//! use tor_basic_utils::test_rng::Config;
55
//! let mut rng = Config::from_env()
56
//!     .unwrap_or(Config::Deterministic)
57
//!     .into_rng();
58
//! ```
59

            
60
// We allow printing to stdout and stderr in this module, since it's intended to
61
// be used by tests, where this is the preferred means of communication with the user.
62
#![allow(clippy::print_stdout, clippy::print_stderr)]
63

            
64
use rand::{RngCore, SeedableRng};
65
// We'll use the same PRNG as the (current) standard.  We specify it here rather
66
// than using StdRng, since we want determinism in the future.
67
pub use rand_chacha::ChaCha12Rng as TestingRng;
68

            
69
/// The seed type for the RNG we're returning.
70
type Seed = <TestingRng as SeedableRng>::Seed;
71

            
72
/// Default seed for deterministic RNG usage.
73
///
74
/// This is the seed we use when we're told to use a deterministic RNG with no
75
/// specific seed.
76
const DEFAULT_SEED: Seed = *b"4   // chosen by fair dice roll.";
77

            
78
/// The environment variable that we inspect.
79
const PRNG_VAR: &str = "ARTI_TEST_PRNG";
80

            
81
/// Return a new, possibly deterministic, RNG for use in tests.
82
///
83
/// This function is **only** for testing: using it elsewhere may make your code
84
/// insecure!
85
///
86
/// The type of this RNG will depend on the value of `ARTI_TEST_PRNG`:
87
///   * If ARTI_TEST_PRNG is `random` or unset, we'll use a real seeded PRNG.
88
///   * If ARTI_TEST_PRNG is `deterministic`, we'll use a standard canned PRNG
89
///     seed.
90
///   * If ARTI_TEST_PRNG is a hexadecimal string, we'll use that as the PRNG
91
///     seed.
92
///
93
/// We'll print the value of this RNG seed to stdout, so that if the test fails,
94
/// you'll know what seed to use in reproducing it.
95
///
96
/// # Panics
97
///
98
/// Panics if the environment variable is set to an invalid value.
99
///
100
/// (If your code must not panic, then it is not test code, and you should not
101
/// be using this function.)
102
47122
pub fn testing_rng() -> TestingRng {
103
47122
    // Somewhat controversially, we prefer a Random prng by default.  Our
104
47122
    // rationale is that, if this weren't the default, nobody would ever set it,
105
47122
    // and we'd never find out about busted tests or code.
106
47122
    Config::from_env().unwrap_or(Config::Random).into_rng()
107
47122
}
108

            
109
/// Type describing a testing_rng configuration.
110
///
111
/// This is a separate type so that you can pick different defaults, or inspect
112
/// the configuration before using it.
113
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
114
#[non_exhaustive]
115
pub enum Config {
116
    /// Use a PRNG with a randomly chosen seed.
117
    Random,
118
    /// Use a PRNG with a (default) pre-selected seed.
119
    Deterministic,
120
    /// Use a specific seed value for the PRNG.
121
    Seeded(Seed),
122
}
123

            
124
impl Config {
125
    /// Return the testing PRNG from the environment, if one is configured.
126
    ///
127
    /// # Panics
128
    ///
129
    /// Panics if the environment variable is set to an invalid value.
130
    ///
131
    /// (If your code must not panic, then it is not test code, and you should not
132
    /// be using this function.)
133
47282
    pub fn from_env() -> Option<Self> {
134
47282
        match Self::from_env_result(std::env::var(PRNG_VAR)) {
135
47282
            Ok(c) => c,
136
            Err(e) => {
137
                panic!(
138
                    "Bad value for {}: {}\n\
139
                    We recognize `random`, `deterministic`, or a hexadecimal seed.",
140
                    PRNG_VAR, e
141
                );
142
            }
143
        }
144
47282
    }
145

            
146
    /// Read the configuration from the result of `std::env::var()`.
147
    ///
148
    /// Return None if there was no option.
149
47296
    fn from_env_result(var: Result<String, std::env::VarError>) -> Result<Option<Self>, Error> {
150
10
        match var {
151
10
            Ok(s) if s.is_empty() => Ok(None),
152
8
            Ok(s) => Ok(Some(Config::from_str(&s)?)),
153
47284
            Err(std::env::VarError::NotPresent) => Ok(None),
154
2
            Err(std::env::VarError::NotUnicode(_)) => Err(Error::InvalidUnicode),
155
        }
156
47296
    }
157

            
158
    /// Read the configuration from a provided string.
159
    ///
160
    /// The string format is as described in [`testing_rng`].
161
    ///
162
    /// Return None if this string can't be interpreted as a [`Config`]
163
24
    fn from_str(s: &str) -> Result<Self, Error> {
164
24
        Ok(if s == "random" {
165
4
            Self::Random
166
20
        } else if s == "deterministic" {
167
4
            Self::Deterministic
168
16
        } else if let Some(seed) = decode_seed_bytes(s) {
169
10
            Self::Seeded(seed)
170
        } else {
171
6
            return Err(Error::UnrecognizedValue(s.to_string()));
172
        })
173
24
    }
174

            
175
    /// Consume this `Config` and return a `Seed`.
176
50974
    fn into_seed(self) -> Seed {
177
50974
        match self {
178
3684
            Config::Deterministic => DEFAULT_SEED,
179
162
            Config::Seeded(seed) => seed,
180
            Config::Random => {
181
47128
                let mut seed = Seed::default();
182
47128
                rand::rng().fill_bytes(&mut seed[..]);
183
47128
                seed
184
            }
185
        }
186
50974
    }
187

            
188
    /// Consume this `Config` and return a `TestingRng`.
189
50966
    pub fn into_rng(self) -> TestingRng {
190
50966
        let seed = self.into_seed();
191
50966
        println!("  Using RNG seed {}={}", PRNG_VAR, format_seed_bytes(&seed));
192
50966
        TestingRng::from_seed(seed)
193
50966
    }
194
}
195

            
196
/// Format `seed` in the format expected by [`decode_seed_bytes`].
197
///
198
/// This is a separate function to make it clearer what the tests are testing.
199
50968
fn format_seed_bytes(seed: &Seed) -> String {
200
50968
    hex::encode(seed)
201
50968
}
202

            
203
/// Try to see whether a literal seed can be decoded from a given string.  If
204
/// so, return it.
205
///
206
/// We currently use a hex encoding, truncating or zero-extending the provided
207
/// seed as needed.
208
18
fn decode_seed_bytes(s: &str) -> Option<Seed> {
209
18
    if s.is_empty() {
210
        // Do not accept the empty string.
211
2
        return None;
212
16
    }
213
16
    let bytes = hex::decode(s).ok()?;
214
12
    let mut seed = Seed::default();
215
12
    let n = std::cmp::min(seed.len(), bytes.len());
216
12
    seed[..n].copy_from_slice(&bytes[..n]);
217
12
    Some(seed)
218
18
}
219

            
220
/// An error from trying to decode a [`Config`] from a string.
221
#[derive(Clone, Debug, thiserror::Error, Eq, PartialEq)]
222
enum Error {
223
    /// We got a value that wasn't unicode.
224
    #[error("Value was not UTF-8")]
225
    InvalidUnicode,
226
    /// We got a value that we otherwise couldn't decode.
227
    #[error("Could not interpret {0:?} as a PRNG seed.")]
228
    UnrecognizedValue(String),
229
}
230

            
231
#[cfg(test)]
232
mod test {
233
    // @@ begin test lint list maintained by maint/add_warning @@
234
    #![allow(clippy::bool_assert_comparison)]
235
    #![allow(clippy::clone_on_copy)]
236
    #![allow(clippy::dbg_macro)]
237
    #![allow(clippy::mixed_attributes_style)]
238
    #![allow(clippy::print_stderr)]
239
    #![allow(clippy::print_stdout)]
240
    #![allow(clippy::single_char_pattern)]
241
    #![allow(clippy::unwrap_used)]
242
    #![allow(clippy::unchecked_duration_subtraction)]
243
    #![allow(clippy::useless_vec)]
244
    #![allow(clippy::needless_pass_by_value)]
245
    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
246
    use std::env::VarError;
247

            
248
    use super::*;
249

            
250
    #[test]
251
    fn from_str() {
252
        assert_eq!(Ok(Config::Deterministic), Config::from_str("deterministic"));
253
        assert_eq!(Ok(Config::Random), Config::from_str("random"));
254
        assert_eq!(Ok(Config::Seeded([0x00; 32])), Config::from_str("00"));
255
        {
256
            let s = "aaaaaaaa";
257
            let seed = [
258
                0xaa, 0xaa, 0xaa, 0xaa, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
259
                0, 0, 0, 0, 0, 0, 0, 0,
260
            ];
261
            assert_eq!(Ok(Config::Seeded(seed)), Config::from_str(s));
262
        }
263
        {
264
            let seed = *b"hello world. this is a longer st";
265
            let mut s = hex::encode(seed);
266
            assert_eq!(Ok(Config::Seeded(seed)), Config::from_str(&s));
267
            // we can make it longer, and it just gets truncated.
268
            s.push_str("aabbccddeeff");
269
            assert_eq!(Ok(Config::Seeded(seed)), Config::from_str(&s));
270
        }
271

            
272
        assert_eq!(
273
            Err(Error::UnrecognizedValue("".to_string())),
274
            Config::from_str("")
275
        );
276

            
277
        assert_eq!(
278
            Err(Error::UnrecognizedValue("return 4".to_string())),
279
            Config::from_str("return 4")
280
        );
281
    }
282

            
283
    #[test]
284
    fn from_env() {
285
        assert_eq!(
286
            Ok(Some(Config::Deterministic)),
287
            Config::from_env_result(Ok("deterministic".to_string()))
288
        );
289
        assert_eq!(
290
            Ok(Some(Config::Random)),
291
            Config::from_env_result(Ok("random".to_string()))
292
        );
293
        assert_eq!(
294
            Ok(Some(Config::Seeded([0xcd; 32]))),
295
            Config::from_env_result(Ok("cd".repeat(32)))
296
        );
297
        assert_eq!(Ok(None), Config::from_env_result(Ok("".to_string())));
298
        assert_eq!(Ok(None), Config::from_env_result(Err(VarError::NotPresent)));
299
        assert_eq!(
300
            Err(Error::InvalidUnicode),
301
            Config::from_env_result(Err(VarError::NotUnicode("3".into())))
302
        );
303
        assert_eq!(
304
            Err(Error::UnrecognizedValue("123".to_string())),
305
            Config::from_env_result(Ok("123".to_string()))
306
        );
307
    }
308

            
309
    #[test]
310
    fn make_seed() {
311
        assert_eq!(Config::Deterministic.into_seed(), DEFAULT_SEED);
312
        assert_eq!(Config::Seeded([0x24; 32]).into_seed(), [0x24; 32]);
313

            
314
        let s1 = Config::Random.into_seed();
315
        let s2 = Config::Random.into_seed();
316
        assert_ne!(s1, s2);
317
    }
318

            
319
    #[test]
320
    fn code_decode() {
321
        assert_eq!(
322
            decode_seed_bytes(&format_seed_bytes(&DEFAULT_SEED)).unwrap(),
323
            DEFAULT_SEED
324
        );
325
    }
326

            
327
    #[test]
328
    fn determinism() {
329
        let mut d_rng = Config::Deterministic.into_rng();
330
        let values: Vec<_> = std::iter::repeat_with(|| d_rng.next_u32())
331
            .take(8)
332
            .collect();
333

            
334
        // This should be the same every time.
335
        let deterministic_values = vec![
336
            4222362647, 2976626662, 1407369338, 1087750672, 196711223, 996083910, 836259566,
337
            2589890951,
338
        ];
339
        assert_eq!(values, deterministic_values);
340

            
341
        // But if we use a random RNG, we'll get different values
342
        // (with P=1-2^-256)
343
        let mut r_rng = Config::Random.into_rng();
344
        let values: Vec<_> = std::iter::repeat_with(|| r_rng.next_u32())
345
            .take(8)
346
            .collect();
347
        assert_ne!(values, deterministic_values);
348
    }
349
}