diff --git a/rust/candid/src/parser/random.rs b/rust/candid/src/parser/random.rs index 79ce9019..4b38dc59 100644 --- a/rust/candid/src/parser/random.rs +++ b/rust/candid/src/parser/random.rs @@ -140,7 +140,7 @@ impl<'a> GenState<'a> { TypeInner::Int64 => IDLValue::Int64(arbitrary_num(u, self.config.range)?), TypeInner::Float32 => IDLValue::Float32(u.arbitrary()?), TypeInner::Float64 => IDLValue::Float64(u.arbitrary()?), - TypeInner::Principal => IDLValue::Principal(crate::Principal::anonymous()), + TypeInner::Principal => IDLValue::Principal(crate::Principal::arbitrary(u)?), TypeInner::Text => { IDLValue::Text(arbitrary_text(u, &self.config.text, &self.config.width)?) } @@ -204,6 +204,9 @@ impl<'a> GenState<'a> { }; IDLValue::Variant(VariantValue(Box::new(field), idx as u64)) } + TypeInner::Func(_) => { + IDLValue::Func(crate::Principal::arbitrary(u)?, String::arbitrary(u)?) + } _ => unimplemented!(), }); self.pop_state(old_config, ty, false); diff --git a/rust/candid/src/types/principal.rs b/rust/candid/src/types/principal.rs index f5c8d776..623ade5d 100644 --- a/rust/candid/src/types/principal.rs +++ b/rust/candid/src/types/principal.rs @@ -1,4 +1,6 @@ use super::{CandidType, Serializer, Type, TypeInner}; +#[cfg(feature = "random")] +use arbitrary::{Arbitrary, Result as ArbitraryResult, Unstructured}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha224}; use std::convert::TryFrom; @@ -365,3 +367,28 @@ impl<'de> serde::Deserialize<'de> for Principal { } } } + +#[cfg(feature = "random")] +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) + } +}