From faff80fcc64747cd56d71e9cc9ab380f2341e17b Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 14 Dec 2023 03:20:47 -0500 Subject: [PATCH] deterministic RNGs for tests --- Cargo.lock | 1 + crates/_types/tokens/Cargo.toml | 4 +- crates/_types/tokens/src/deterministic_rng.rs | 43 +++++++++++++++++++ crates/_types/tokens/src/lib.rs | 1 + crates/_types/tokens/src/macros.rs | 35 +++++++++++++++ .../tokens/tests/media_file_token_tests.rs | 24 +++++++++++ .../weights/list_featured_weights.rs | 1 - 7 files changed, 107 insertions(+), 2 deletions(-) create mode 100644 crates/_types/tokens/src/deterministic_rng.rs diff --git a/Cargo.lock b/Cargo.lock index 19da97d557..af9894f874 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6442,6 +6442,7 @@ name = "tokens" version = "0.0.1" dependencies = [ "my-workspace-hack", + "once_cell", "rand 0.8.5", "serde", "serde_derive", diff --git a/crates/_types/tokens/Cargo.toml b/crates/_types/tokens/Cargo.toml index 581f8c65de..94152dee1b 100644 --- a/crates/_types/tokens/Cargo.toml +++ b/crates/_types/tokens/Cargo.toml @@ -25,7 +25,9 @@ serde = "1.0.125" serde_derive = "1.0.125" # Random -rand = "0.8.5" +rand = { version = "0.8.5", features = ["std_rng"] } + +once_cell = { version = "1.18.0", features = [] } sqlx = { version = "0.7.2", features = [ "mysql", "runtime-tokio", "tls-rustls", "chrono" ] } sqlx-core = { version = "0.7.2" } diff --git a/crates/_types/tokens/src/deterministic_rng.rs b/crates/_types/tokens/src/deterministic_rng.rs new file mode 100644 index 0000000000..90ad47f5ab --- /dev/null +++ b/crates/_types/tokens/src/deterministic_rng.rs @@ -0,0 +1,43 @@ +use std::sync::{Arc, LockResult, RwLock, RwLockWriteGuard}; + +use once_cell::sync::Lazy; +use rand::rngs::StdRng; +use rand::SeedableRng; + +// FIXME(bt,2023-12-14): Is this threadsafe? Looks like it should be fine. +/// Only use this in tests (and dev seed tools), never in production! +#[derive(Clone)] +pub struct DeterministicRng { + rng : Arc> +} + +impl DeterministicRng { + pub fn get_instance() -> LockResult> { + static DETERMINISTIC_RNG : Lazy>> = Lazy::new(|| { + Arc::new(RwLock::new(DeterministicRng::new())) + }); + + DETERMINISTIC_RNG.write() + } + + fn new() -> Self { + Self { + rng: Arc::new(RwLock::new(StdRng::seed_from_u64(0))) + } + } + + /// Reset the deterministic RNG's seed. + pub fn reset_rng(&self) { + match self.rng.write() { + Err(err) => panic!(format!("test panic due to rng failure {err}")), + Ok(mut lock) => { + *lock = StdRng::seed_from_u64(0); + } + } + } + + /// Grab the deterministic RNG + pub fn get_rng(&self) -> LockResult>{ + self.rng.write() + } +} diff --git a/crates/_types/tokens/src/lib.rs b/crates/_types/tokens/src/lib.rs index df6e81b4f2..b220bdb751 100644 --- a/crates/_types/tokens/src/lib.rs +++ b/crates/_types/tokens/src/lib.rs @@ -49,6 +49,7 @@ pub(crate) static MINIMUM_CHARACTER_ENTROPY : usize = 8; #[macro_use] mod macros; +pub (crate) mod deterministic_rng; pub (crate) mod prefixes; pub mod tokens; diff --git a/crates/_types/tokens/src/macros.rs b/crates/_types/tokens/src/macros.rs index 5964b6109a..4df28db41a 100644 --- a/crates/_types/tokens/src/macros.rs +++ b/crates/_types/tokens/src/macros.rs @@ -60,6 +60,41 @@ macro_rules! impl_crockford_generator { $t(token) } + /// Constructor for a new token to be used **DETERMINISTICALLY** in tests and seeding. + /// This will not use a real RNG, so do not use in production code! + /// Once called, it will output the same random numbers in sequence. + pub fn reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously() { + let deterministic = crate::deterministic_rng::DeterministicRng::get_instance().expect("tests should not fail due to mutex"); + deterministic.reset_rng(); + } + + /// Constructor for a new token to be used **DETERMINISTICALLY** in tests and seeding. + /// This will not use a real RNG, so do not use in production code! + /// Once called, it will output the same random numbers in sequence. + pub fn generate_for_testing_and_dev_seeding_never_use_in_production_seriously() -> Self { + use rand::Rng; + + // FIXME(bt,2023-12-14): This is bad error handling. (Maybe panic if lock is poisoned?) + let deterministic = crate::deterministic_rng::DeterministicRng::get_instance().expect("tests should not fail due to mutex"); + let mut rng = deterministic.get_rng().expect("tests should not fail due to mutex"); + + let charset = Self::token_character_set(); + let entropy_length = Self::entropic_character_len(); + + let entropy_part: String = (0..entropy_length) + .map(|_| { + let idx = rng.gen_range(0..charset.len()); + charset[idx] as char + }) + .collect(); + + let token_prefix = Self::token_prefix(); + + let token = format!("{}{}", token_prefix, entropy_part); + + $t(token) + } + #[inline] pub fn entropy_suffix(&self) -> &str { use crate::prefixes::PrefixGenerator; diff --git a/crates/_types/tokens/tests/media_file_token_tests.rs b/crates/_types/tokens/tests/media_file_token_tests.rs index dc09692485..9a8704e1b2 100644 --- a/crates/_types/tokens/tests/media_file_token_tests.rs +++ b/crates/_types/tokens/tests/media_file_token_tests.rs @@ -21,6 +21,30 @@ mod interface { assert!(token.to_string().starts_with("m_")); } + #[test] + fn generate_for_testing_and_dev_seeding_never_use_in_production_seriously_1() { + MediaFileToken::reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(); + assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_q8sz47gmfw2zx02snrbz88ns9m16ab"); + assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_ma1xetxrwbh39vg639a9zrq8b9wk6h"); + assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_4tswec8z27wnm01njypx4vmfhgj41e"); + } + + #[test] + fn generate_for_testing_and_dev_seeding_never_use_in_production_seriously_2() { + MediaFileToken::reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(); + assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_q8sz47gmfw2zx02snrbz88ns9m16ab"); + assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_ma1xetxrwbh39vg639a9zrq8b9wk6h"); + assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_4tswec8z27wnm01njypx4vmfhgj41e"); + } + + #[test] + fn generate_for_testing_and_dev_seeding_never_use_in_production_seriously_3() { + MediaFileToken::reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(); + assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_q8sz47gmfw2zx02snrbz88ns9m16ab"); + assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_ma1xetxrwbh39vg639a9zrq8b9wk6h"); + assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_4tswec8z27wnm01njypx4vmfhgj41e"); + } + #[test] fn new() { let token = MediaFileToken::new("m_foo".to_string()); diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_featured_weights.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_featured_weights.rs index 33f921b722..87d75eed0f 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_featured_weights.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_featured_weights.rs @@ -122,7 +122,6 @@ pub async fn list_featured_weights_handler( }; } - // generate parse a response let response = ListFeaturedWeightsSuccessResponse { success: true, weights: weights.into_iter()