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

dashpay · Oct 1, 2024 · 5c598ab · 5c598ab
1 parent 023ff85
commit 5c598ab
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/packages/rs-dpp/src/identity/identity_public_key/random.rs b/packages/rs-dpp/src/identity/identity_public_key/random.rs
@@ -615,15 +615,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 3..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)
     }
 

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
@@ -69,15 +69,16 @@ impl IdentityPublicKeyV0 {
         used_key_matrix: Option<(KeyCount, &mut UsedKeyMatrix)>,
         platform_version: &PlatformVersion,
     ) -> Result<(Self, Vec<u8>), ProtocolError> {
+        const MAX_KEYS: u8 = 16;
         // we have 16 different permutations possible
-        let mut binding = [false; 16].to_vec();
+        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 {
+        if key_count > MAX_KEYS as u32 {
             return Err(ProtocolError::PublicKeyGenerationError(
                 "too many keys already created".to_string(),
             ));
         }
-        let key_number = rng.gen_range(0..(12 - key_count as u8));
+        let key_number = rng.gen_range(0..(MAX_KEYS - key_count as u8));
         // now we need to find the first bool that isn't set to true
         let mut needed_pos = None;
         let mut counter = 0;