diff --git a/packages/rs-dpp/src/identity/identity_public_key/random.rs b/packages/rs-dpp/src/identity/identity_public_key/random.rs index 0dea7a5dee..b43155e056 100644 --- a/packages/rs-dpp/src/identity/identity_public_key/random.rs +++ b/packages/rs-dpp/src/identity/identity_public_key/random.rs @@ -9,6 +9,7 @@ use rand::rngs::StdRng; use rand::SeedableRng; pub type UsedKeyMatrix = Vec; +pub const MAX_RANDOM_KEYS: usize = 16; impl IdentityPublicKey { pub fn random_key(id: KeyID, seed: Option, platform_version: &PlatformVersion) -> Self { @@ -542,7 +543,7 @@ impl IdentityPublicKey { rng: &mut StdRng, platform_version: &PlatformVersion, ) -> Result, ProtocolError> { - let mut used_key_matrix = [false; 16].to_vec(); + let mut used_key_matrix = [false; MAX_RANDOM_KEYS].to_vec(); (0..key_count) .map(|i| { Self::random_authentication_key_with_rng( @@ -583,6 +584,14 @@ impl IdentityPublicKey { "at least 2 keys must be created".to_string(), )); } + + if key_count > MAX_RANDOM_KEYS as u32 { + return Err(ProtocolError::PublicKeyGenerationError(format!( + "too many keys requested: {}, max is {}", + key_count, MAX_RANDOM_KEYS + ))); + } + //create a master and a high level key let mut main_keys = if key_count == 2 { vec![ @@ -608,22 +617,27 @@ impl IdentityPublicKey { )?, ] }; - let mut used_key_matrix = [false; 16].to_vec(); + let mut used_key_matrix = [false; MAX_RANDOM_KEYS].to_vec(); used_key_matrix[0] = true; used_key_matrix[1] = true; used_key_matrix[2] = true; used_key_matrix[4] = true; //also a master key used_key_matrix[8] = true; //also a master key used_key_matrix[12] = true; //also a master key - main_keys.extend((3..key_count).map(|i| { - Self::random_authentication_key_with_private_key_with_rng( + let used_key_matrix_count = used_key_matrix.iter().filter(|x| **x).count() as u32; + let main_keys_count = main_keys.len() as u32; + for i in main_keys_count..key_count { + let privkey = Self::random_authentication_key_with_private_key_with_rng( i, rng, - Some((i, &mut used_key_matrix)), + Some(( + used_key_matrix_count + (i - main_keys_count), + &mut used_key_matrix, + )), platform_version, - ) - .unwrap() - })); + )?; + main_keys.push(privkey); + } Ok(main_keys) } diff --git a/packages/rs-dpp/src/identity/identity_public_key/v0/random.rs b/packages/rs-dpp/src/identity/identity_public_key/v0/random.rs index 13afdd673f..8bb6d7751f 100644 --- a/packages/rs-dpp/src/identity/identity_public_key/v0/random.rs +++ b/packages/rs-dpp/src/identity/identity_public_key/v0/random.rs @@ -1,4 +1,5 @@ use crate::identity::contract_bounds::ContractBounds; +use crate::identity::identity_public_key::random::MAX_RANDOM_KEYS; use crate::identity::identity_public_key::v0::IdentityPublicKeyV0; use crate::identity::KeyType::{ECDSA_HASH160, ECDSA_SECP256K1}; use crate::identity::Purpose::{AUTHENTICATION, VOTING}; @@ -70,17 +71,32 @@ impl IdentityPublicKeyV0 { platform_version: &PlatformVersion, ) -> Result<(Self, Vec), ProtocolError> { // we have 16 different permutations possible - let mut binding = [false; 16].to_vec(); + + let mut binding = [false; MAX_RANDOM_KEYS].to_vec(); let (key_count, key_matrix) = used_key_matrix.unwrap_or((0, &mut binding)); - if key_count > 16 { + + // input validation + if key_count != key_matrix.iter().filter(|x| **x).count() as u32 { + return Err(ProtocolError::PublicKeyGenerationError( + "invalid argument: key count in used_key_matrix.0 doesn't match actual number of used keys".to_string(), + )); + } + + // we need space for at least one additional key + if key_count > MAX_RANDOM_KEYS as u32 - 1 { return Err(ProtocolError::PublicKeyGenerationError( "too many keys already created".to_string(), )); } - let key_number = rng.gen_range(0..(12 - key_count as u8)); - // now we need to find the first bool that isn't set to true + + // max_key_number is the number of keys that can be created + let max_key_number = MAX_RANDOM_KEYS as u32 - key_count; + let key_number = rng.gen_range(0..max_key_number); + // now we need to find n'th not used key of this key_matrix (where `n = key_number`), + // that is: the first bool that isn't set to true let mut needed_pos = None; let mut counter = 0; + let mut used = 0; key_matrix.iter_mut().enumerate().for_each(|(pos, is_set)| { if !*is_set { if counter == key_number { @@ -88,11 +104,20 @@ impl IdentityPublicKeyV0 { *is_set = true; } counter += 1; + } else { + used += 1; // for debugging purposes } }); - let needed_pos = needed_pos.ok_or(ProtocolError::PublicKeyGenerationError( - "too many keys already created".to_string(), - ))?; + // should never happen, as we have input validation above + assert_eq!( + used, key_count, + "incorrect number of used keys in key_matrix {:?}", + key_matrix, + ); + let needed_pos = needed_pos.ok_or(ProtocolError::PublicKeyGenerationError(format!( + "too many keys already created: , key_count: {}, random key number: {}, unused key counter: {}, used key counter: {}", + key_count, key_number, counter, used, + )))?; let key_type = needed_pos.div(&4); let security_level = needed_pos.rem(&4); let security_level = SecurityLevel::try_from(security_level).unwrap(); diff --git a/packages/rs-drive-abci/tests/strategy_tests/main.rs b/packages/rs-drive-abci/tests/strategy_tests/main.rs index 98e2b518a2..17a0a35438 100644 --- a/packages/rs-drive-abci/tests/strategy_tests/main.rs +++ b/packages/rs-drive-abci/tests/strategy_tests/main.rs @@ -47,7 +47,6 @@ mod tests { use crate::execution::{continue_chain_for_strategy, run_chain_for_strategy}; use crate::query::QueryStrategy; use crate::strategy::{FailureStrategy, MasternodeListChangesStrategy}; - use assert_matches::assert_matches; use dashcore_rpc::dashcore::hashes::Hash; use dashcore_rpc::dashcore::BlockHash; use dashcore_rpc::json::QuorumType;