fix(dpp): panic when generating more than 12 identity keys

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<bool>;
+pub const MAX_RANDOM_KEYS: usize = 16;
 
 impl IdentityPublicKey {
     pub fn random_key(id: KeyID, seed: Option<u64>, platform_version: &PlatformVersion) -> Self {
@@ -542,7 +543,7 @@ impl IdentityPublicKey {
         rng: &mut StdRng,
         platform_version: &PlatformVersion,
     ) -> Result<Vec<Self>, 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)
     }
 

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,29 +71,53 @@ impl IdentityPublicKeyV0 {
         platform_version: &PlatformVersion,
     ) -> Result<(Self, Vec<u8>), 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 {
                     needed_pos = Some(pos as u8);
                     *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();

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;