diff --git a/Cargo.lock b/Cargo.lock index 6f025f8a..98eaec75 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -855,6 +855,7 @@ dependencies = [ name = "ic_principal" version = "0.1.0" dependencies = [ + "arbitrary", "crc32fast", "data-encoding", "hex", diff --git a/rust/candid/Cargo.toml b/rust/candid/Cargo.toml index dd263d93..e0ff00c2 100644 --- a/rust/candid/Cargo.toml +++ b/rust/candid/Cargo.toml @@ -50,7 +50,7 @@ bignum = ["num-bigint", "num-traits"] value = ["bignum"] printer = ["pretty"] default = ["serde_bytes", "printer", "bignum"] -all = ["default", "value"] +all = ["default", "value", "ic_principal/arbitrary"] [[test]] name = "types" diff --git a/rust/ic_principal/Cargo.toml b/rust/ic_principal/Cargo.toml index 8164b80d..157f5f43 100644 --- a/rust/ic_principal/Cargo.toml +++ b/rust/ic_principal/Cargo.toml @@ -34,6 +34,11 @@ optional = true version = "0.11.5" optional = true +[dependencies.arbitrary] +version = "1.0" +optional = true + [features] # Default features include serde support. default = ['serde', 'serde_bytes'] +arbitrary = ['default', 'dep:arbitrary'] diff --git a/rust/ic_principal/src/lib.rs b/rust/ic_principal/src/lib.rs index b128ccce..a0fa9b92 100644 --- a/rust/ic_principal/src/lib.rs +++ b/rust/ic_principal/src/lib.rs @@ -5,6 +5,9 @@ use std::convert::TryFrom; use std::fmt::Write; use thiserror::Error; +#[cfg(feature = "arbitrary")] +use arbitrary::{Arbitrary, Result as ArbitraryResult, Unstructured}; + /// An error happened while encoding, decoding or serializing a [`Principal`]. #[derive(Error, Clone, Debug, Eq, PartialEq)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] @@ -358,3 +361,28 @@ impl<'de> serde::Deserialize<'de> for Principal { } } } + +#[cfg(feature = "arbitrary")] +impl<'a> Arbitrary<'a> for Principal { + fn arbitrary(u: &mut Unstructured<'a>) -> ArbitraryResult { + let principal = match u8::arbitrary(u)? { + u8::MAX => Principal::management_canister(), + 254u8 => Principal::anonymous(), + _ => { + let length: usize = u.int_in_range(1..=Principal::MAX_LENGTH_IN_BYTES)?; + let mut result: Vec = Vec::with_capacity(length); + for _ in 0..length { + result.push(u8::arbitrary(u)?); + } + // non-anonymous principal cannot have type ANONYMOUS + // adapt by changing the last byte. + let last = result.last_mut().unwrap(); + if *last == 4_u8 { + *last = u8::MAX + } + Principal::try_from(&result[..]).unwrap() + } + }; + Ok(principal) + } +}