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

packages/rs-dpp/src/identity/identity_public_key/random.rs

@@ -583,6 +583,14 @@ impl IdentityPublicKey {
                 "at least 2 keys must be created".to_string(),
             ));
         }
+
+        if key_count > 16 {
+            return Err(ProtocolError::PublicKeyGenerationError(format!(
+                "too many keys requested: {}, max is 16",
+                key_count
+            )));
+        }
+
         //create a master and a high level key
         let mut main_keys = if key_count == 2 {
             vec![
@@ -615,15 +623,15 @@ impl IdentityPublicKey {
         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(
+        for i in 6..key_count {
+            let privkey = Self::random_authentication_key_with_private_key_with_rng(
                 i,
                 rng,
                 Some((i, &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

@@ -70,29 +70,54 @@ impl IdentityPublicKeyV0 {
         platform_version: &PlatformVersion,
     ) -> Result<(Self, Vec<u8>), ProtocolError> {
         // we have 16 different permutations possible
-        let mut binding = [false; 16].to_vec();
+        const MAX_KEYS: KeyCount = 16;
+
+        let mut binding = [false; MAX_KEYS as usize].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_KEYS - 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_KEYS - 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;