fix: use correct key id

dashpay · Oct 2, 2024 · 0879ca0 · 0879ca0
1 parent e03289d
commit 0879ca0
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 11 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
@@ -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(
@@ -584,10 +585,10 @@ impl IdentityPublicKey {
             ));
         }
 
-        if key_count > 16 {
+        if key_count > MAX_RANDOM_KEYS as u32 {
             return Err(ProtocolError::PublicKeyGenerationError(format!(
-                "too many keys requested: {}, max is 16",
-                key_count
+                "too many keys requested: {}, max is {}",
+                key_count, MAX_RANDOM_KEYS
             )));
         }
 
@@ -616,18 +617,23 @@ 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
-        for i in 6..key_count {
+        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,
             )?;
             main_keys.push(privkey);

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
@@ -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,9 +71,8 @@ impl IdentityPublicKeyV0 {
         platform_version: &PlatformVersion,
     ) -> Result<(Self, Vec<u8>), ProtocolError> {
         // we have 16 different permutations possible
-        const MAX_KEYS: KeyCount = 16;
 
-        let mut binding = [false; MAX_KEYS as usize].to_vec();
+        let mut binding = [false; MAX_RANDOM_KEYS].to_vec();
         let (key_count, key_matrix) = used_key_matrix.unwrap_or((0, &mut binding));
 
         // input validation
@@ -83,14 +83,14 @@ impl IdentityPublicKeyV0 {
         }
 
         // we need space for at least one additional key
-        if key_count > MAX_KEYS - 1 {
+        if key_count > MAX_RANDOM_KEYS as u32 - 1 {
             return Err(ProtocolError::PublicKeyGenerationError(
                 "too many keys already created".to_string(),
             ));
         }
 
         // max_key_number is the number of keys that can be created
-        let max_key_number = MAX_KEYS - key_count;
+        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