Skip to content

Commit

Permalink
Generate seeded random test strings (#2265)
Browse files Browse the repository at this point in the history
  • Loading branch information
heaths authored Mar 1, 2025
1 parent 3e41b0b commit a972420
Show file tree
Hide file tree
Showing 4 changed files with 150 additions and 40 deletions.
1 change: 1 addition & 0 deletions .vscode/cspell.json
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
"asyncoperation",
"azsdk",
"azurecli",
"bugbug",
"clippy",
"contoso",
"cplusplus",
Expand Down
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions sdk/core/azure_core_test/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ tokio = { workspace = true, features = [
azure_security_keyvault_secrets = { path = "../../keyvault/azure_security_keyvault_secrets" }
clap.workspace = true
tracing-subscriber = { workspace = true, features = ["env-filter", "fmt"] }
uuid.workspace = true

[target.'cfg(not(target_arch = "wasm32"))'.dev-dependencies]
tokio = { workspace = true, features = ["signal"] }
187 changes: 147 additions & 40 deletions sdk/core/azure_core_test/src/recording.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

//! The [`Recording`] and other types used in recorded tests.
// cspell:ignore csprng seedable
// cspell:ignore csprng seedable tpbwhbkhckmk
use crate::{
proxy::{
client::{
Expand All @@ -26,7 +26,7 @@ use azure_core::{
};
use azure_identity::DefaultAzureCredential;
use rand::{
distributions::{Distribution, Standard},
distributions::{Alphanumeric, DistString, Distribution, Standard},
Rng, SeedableRng,
};
use rand_chacha::ChaCha20Rng;
Expand Down Expand Up @@ -143,10 +143,29 @@ impl Recording {
///
/// # Examples
///
/// Generate a random integer.
///
/// ```
/// # let recording = azure_core_test::Recording::with_seed();
/// let i: i32 = recording.random();
/// # assert_eq!(i, 1054672670);
/// ```
///
/// Generate a symmetric data encryption key (DEK).
///
/// ```no_compile
/// ```
/// # let recording = azure_core_test::Recording::with_seed();
/// let dek: [u8; 32] = recording.random();
/// # assert_eq!(typespec_client_core::base64::encode(dek), "HumPRAN6RqKWf0YhFV2CAFWu/8L/pwh0LRzeam5VlGo=");
/// ```
///
/// Generate a UUID.
///
/// ```
/// use azure_core::Uuid;
/// # let recording = azure_core_test::Recording::with_seed();
/// let uuid: Uuid = Uuid::from_u128(recording.random());
/// # assert_eq!(uuid.to_string(), "fe906b44-5838-cc8f-05e3-c7e93edd071e");
/// ```
///
/// # Panics
Expand All @@ -158,52 +177,68 @@ impl Recording {
where
Standard: Distribution<T>,
{
const NAME: &str = "RandomSeed";

// Use ChaCha20 for a deterministic, portable CSPRNG.
let rng = self.rand.get_or_init(|| match self.test_mode {
TestMode::Live => ChaCha20Rng::from_entropy().into(),
TestMode::Playback => {
let variables = self
.variables
.read()
.map_err(read_lock_error)
.unwrap_or_else(|err| panic!("{err}"));
let seed: String = variables
.get(NAME)
.map(Into::into)
.unwrap_or_else(|| panic!("random seed variable not set"));
let seed = base64::decode(seed)
.unwrap_or_else(|err| panic!("failed to decode random seed: {err}"));
let seed = seed
.first_chunk::<32>()
.unwrap_or_else(|| panic!("insufficient random seed variable"));

ChaCha20Rng::from_seed(*seed).into()
}
TestMode::Record => {
let rng = ChaCha20Rng::from_entropy();
let seed = rng.get_seed();
let seed = base64::encode(seed);
let rng = self.rng();
let Ok(mut rng) = rng.lock() else {
panic!("failed to lock RNG");
};

let mut variables = self
.variables
.write()
.map_err(write_lock_error)
.unwrap_or_else(|err| panic!("{err}"));
variables.insert(NAME.to_string(), Value::from(Some(seed), None));
rng.gen()
}

rng.into()
/// Generate a random string with optional prefix.
///
/// This will always be the OS cryptographically secure pseudo-random number generator (CSPRNG) when running live.
/// When recording, it will initialize from the OS CSPRNG but save the seed value to the recording file.
/// When playing back, the saved seed value is read from the recording to reproduce the same sequence of random data.
///
/// # Examples
///
/// Generate a random string.
///
/// ```
/// # let recording = azure_core_test::Recording::with_seed();
/// let id = recording.random_string::<12>(Some("t")).to_ascii_lowercase();
/// # assert_eq!(id, "tpbwhbkhckmk");
/// ```
///
/// # Panics
///
/// Panics if the recording variables cannot be locked for reading or writing,
/// if the random seed cannot be encoded or decoded properly,
/// if `LEN` is 0,
/// or if the length of `prefix` is greater than or equal to `LEN`.
///
/// ```should_panic
/// # let recording = azure_core_test::Recording::with_seed();
/// let vault_name = recording.random_string::<8>(Some("keyvault"));
/// ```
///
pub fn random_string<const LEN: usize>(&self, prefix: Option<&str>) -> String {
struct NonZero<const N: usize>;
impl<const N: usize> NonZero<N> {
const ASSERT: () = assert!(N > 0, "LEN must be greater than 0");
}
#[allow(clippy::let_unit_value)]
let _ = NonZero::<LEN>::ASSERT;
let len = match prefix {
Some(p) => {
assert!(p.len() < LEN, "prefix length must be less than LEN");
LEN - p.len()
}
});
None => LEN,
};

let rng = self.rng();
let Ok(mut rng) = rng.lock() else {
panic!("failed to lock RNG");
};

rng.gen()
let value = Alphanumeric.sample_string(&mut *rng, len);
match prefix {
Some(prefix) => prefix.to_string() + &value,
None => value,
}
}

/// Removes the list of sanitizers from the recording.
///
/// You can find a list of default sanitizers in [source code](https://github.com/Azure/azure-sdk-tools/blob/main/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/SanitizerDictionary.cs).
Expand Down Expand Up @@ -285,6 +320,8 @@ impl Recording {
}
}

const RANDOM_SEED_NAME: &str = "RandomSeed";

impl Recording {
pub(crate) fn new(
test_mode: TestMode,
Expand All @@ -310,6 +347,28 @@ impl Recording {
}
}

// #[cfg(any(test, doctest))] // BUGBUG: https://github.com/rust-lang/rust/issues/67295
#[doc(hidden)]
pub fn with_seed() -> Self {
let span = tracing::trace_span!("Recording::seeded");
Self {
test_mode: TestMode::Playback,
span: span.entered(),
_proxy: None,
client: None,
policy: OnceCell::new(),
service_directory: String::from("sdk/core"),
recording_file: String::from("none"),
recording_assets_file: None,
id: None,
variables: RwLock::new(HashMap::from([(
RANDOM_SEED_NAME.into(),
"8S9UCR2yV8LU01tq+VNEwGssAXVUbL0Hd488GAYVosM=".into(),
)])),
rand: OnceLock::new(),
}
}

fn env<K>(&self, key: K) -> Option<String>
where
K: AsRef<str>,
Expand All @@ -322,6 +381,45 @@ impl Recording {
.and_then(|v| v.into_string().ok())
}

fn rng(&self) -> &Mutex<ChaCha20Rng> {
// Use ChaCha20 for a deterministic, portable CSPRNG.
self.rand.get_or_init(|| match self.test_mode {
TestMode::Live => ChaCha20Rng::from_entropy().into(),
TestMode::Playback => {
let variables = self
.variables
.read()
.map_err(read_lock_error)
.unwrap_or_else(|err| panic!("{err}"));
let seed: String = variables
.get(RANDOM_SEED_NAME)
.map(Into::into)
.unwrap_or_else(|| panic!("random seed variable not set"));
let seed = base64::decode(seed)
.unwrap_or_else(|err| panic!("failed to decode random seed: {err}"));
let seed = seed
.first_chunk::<32>()
.unwrap_or_else(|| panic!("insufficient random seed variable"));

ChaCha20Rng::from_seed(*seed).into()
}
TestMode::Record => {
let rng = ChaCha20Rng::from_entropy();
let seed = rng.get_seed();
let seed = base64::encode(seed);

let mut variables = self
.variables
.write()
.map_err(write_lock_error)
.unwrap_or_else(|err| panic!("{err}"));
variables.insert(RANDOM_SEED_NAME.to_string(), Value::from(Some(seed), None));

rng.into()
}
})
}

fn set_skip(&self, skip: Option<Skip>) -> azure_core::Result<()> {
let Some(policy) = self.policy.get() else {
return Ok(());
Expand Down Expand Up @@ -502,6 +600,15 @@ impl Value {
}
}

impl From<&str> for Value {
fn from(value: &str) -> Self {
Self {
value: value.into(),
sanitized: None,
}
}
}

impl From<String> for Value {
fn from(value: String) -> Self {
Self {
Expand Down

0 comments on commit a972420

Please sign in to comment.