fix(drive): add masternode identity keys through a dedicated operation

packages/rs-drive-abci/src/execution/platform_events/core_based_updates/update_masternode_identities/update_masternode_identities/v1/mod.rs

@@ -19,7 +19,7 @@ use drive::drive::identity::key::fetch::{
 use drive::util::batch::DriveOperation;
 use drive::util::batch::DriveOperation::IdentityOperation;
 use drive::util::batch::IdentityOperationType::{
-    AddNewIdentity, AddNewKeysToIdentity, DisableIdentityKeys, ReEnableIdentityKeys,
+    AddNewIdentity, AddNewKeysToMasternodeIdentity, DisableIdentityKeys, ReEnableIdentityKeys,
 };
 
 use crate::platform_types::platform_state::PlatformStateV0Methods;
@@ -287,9 +287,8 @@ where
             }
         }
         if let Some(key) = new_key {
-            drive_operations.push(IdentityOperation(AddNewKeysToIdentity {
+            drive_operations.push(IdentityOperation(AddNewKeysToMasternodeIdentity {
                 identity_id: owner_identifier,
-                unique_keys_to_add: vec![],
                 non_unique_keys_to_add: vec![key],
             }));
         }

packages/rs-drive-abci/src/execution/platform_events/core_based_updates/update_masternode_identities/update_operator_identity/v0/mod.rs

@@ -24,7 +24,7 @@ use drive::grovedb::Transaction;
 use drive::util::batch::DriveOperation;
 use drive::util::batch::DriveOperation::IdentityOperation;
 use drive::util::batch::IdentityOperationType::{
-    AddNewIdentity, AddNewKeysToIdentity, DisableIdentityKeys, ReEnableIdentityKeys,
+    AddNewIdentity, AddNewKeysToMasternodeIdentity, DisableIdentityKeys, ReEnableIdentityKeys,
 };
 
 impl<C> Platform<C>
@@ -212,9 +212,8 @@ where
             }
 
             if !non_unique_keys_to_add.is_empty() {
-                drive_operations.push(IdentityOperation(AddNewKeysToIdentity {
+                drive_operations.push(IdentityOperation(AddNewKeysToMasternodeIdentity {
                     identity_id: operator_identifier.to_buffer(),
-                    unique_keys_to_add: vec![],
                     non_unique_keys_to_add,
                 }));
             }
@@ -374,9 +373,8 @@ where
 
                 if !non_unique_keys_to_add.is_empty() {
                     // We should add keys that didn't exist before
-                    drive_operations.push(IdentityOperation(AddNewKeysToIdentity {
+                    drive_operations.push(IdentityOperation(AddNewKeysToMasternodeIdentity {
                         identity_id: operator_identifier.to_buffer(),
-                        unique_keys_to_add: vec![],
                         non_unique_keys_to_add,
                     }));
                 }
@@ -1172,7 +1170,7 @@ mod tests {
 
         let added_platform_node_key = drive_operations.iter().any(|operation| {
             if let drive::util::batch::DriveOperation::IdentityOperation(
-                drive::util::batch::IdentityOperationType::AddNewKeysToIdentity {
+                drive::util::batch::IdentityOperationType::AddNewKeysToMasternodeIdentity {
                     identity_id,
                     non_unique_keys_to_add,
                     ..

packages/rs-drive-abci/src/execution/platform_events/core_based_updates/update_masternode_identities/update_owner_withdrawal_address/v0/mod.rs

@@ -16,7 +16,7 @@ use drive::grovedb::Transaction;
 use drive::util::batch::DriveOperation;
 use drive::util::batch::DriveOperation::IdentityOperation;
 use drive::util::batch::IdentityOperationType::{
-    AddNewKeysToIdentity, DisableIdentityKeys, ReEnableIdentityKeys,
+    AddNewKeysToMasternodeIdentity, DisableIdentityKeys, ReEnableIdentityKeys,
 };
 impl<C> Platform<C>
 where
@@ -114,9 +114,8 @@ where
                 "add new withdrawal key to owner identity"
             );
 
-            drive_operations.push(IdentityOperation(AddNewKeysToIdentity {
+            drive_operations.push(IdentityOperation(AddNewKeysToMasternodeIdentity {
                 identity_id: owner_identifier,
-                unique_keys_to_add: vec![],
                 non_unique_keys_to_add: vec![new_owner_withdrawal_key],
             }));
         }

packages/rs-drive-abci/src/execution/platform_events/core_based_updates/update_masternode_identities/update_owner_withdrawal_address/v1/mod.rs

@@ -16,7 +16,7 @@ use drive::grovedb::Transaction;
 use drive::util::batch::DriveOperation;
 use drive::util::batch::DriveOperation::IdentityOperation;
 use drive::util::batch::IdentityOperationType::{
-    AddNewKeysToIdentity, DisableIdentityKeys, ReEnableIdentityKeys,
+    AddNewKeysToMasternodeIdentity, DisableIdentityKeys, ReEnableIdentityKeys,
 };
 impl<C> Platform<C>
 where
@@ -124,9 +124,8 @@ where
                 "add new withdrawal key to owner identity"
             );
 
-            drive_operations.push(IdentityOperation(AddNewKeysToIdentity {
+            drive_operations.push(IdentityOperation(AddNewKeysToMasternodeIdentity {
                 identity_id: owner_identifier,
-                unique_keys_to_add: vec![],
                 non_unique_keys_to_add: vec![new_owner_withdrawal_key],
             }));
         }

packages/rs-drive/src/drive/identity/update/methods/add_new_keys_to_identity/mod.rs

@@ -66,4 +66,58 @@ impl Drive {
             })),
         }
     }
+
+    /// The operations for adding new non-unique keys to a masternode identity. This function is
+    /// version controlled by the version of `add_new_keys_to_identity_operations`, whose
+    /// operations it emits when estimating costs and for every key whose hash is not a unique
+    /// key hash.
+    ///
+    /// A key whose hash is already a unique key hash is added without a non-unique key hash
+    /// reference, matching how keys are stored when a masternode identity is created.
+    ///
+    /// # Arguments
+    ///
+    /// * `identity_id` - The ID of the masternode identity to which keys are to be added.
+    /// * `non_unique_keys_to_add` - The keys to be added, all of non-unique key types.
+    /// * `epoch` - The current epoch.
+    /// * `estimated_costs_only_with_layer_info` - The estimated costs with layer information.
+    /// * `transaction` - The current transaction.
+    /// * `platform_version` - The platform version.
+    ///
+    /// # Returns
+    ///
+    /// * `Result<Vec<LowLevelDriveOperation>, Error>` - The resulting low level drive operations if successful, or an error.
+    pub fn add_new_keys_to_masternode_identity_operations(
+        &self,
+        identity_id: [u8; 32],
+        non_unique_keys_to_add: Vec<IdentityPublicKey>,
+        epoch: &Epoch,
+        estimated_costs_only_with_layer_info: &mut Option<
+            HashMap<KeyInfoPath, EstimatedLayerInformation>,
+        >,
+        transaction: TransactionArg,
+        platform_version: &PlatformVersion,
+    ) -> Result<Vec<LowLevelDriveOperation>, Error> {
+        match platform_version
+            .drive
+            .methods
+            .identity
+            .update
+            .add_new_keys_to_identity
+        {
+            0 => self.add_new_keys_to_masternode_identity_operations_v0(
+                identity_id,
+                non_unique_keys_to_add,
+                epoch,
+                estimated_costs_only_with_layer_info,
+                transaction,
+                platform_version,
+            ),
+            version => Err(Error::Drive(DriveError::UnknownVersionMismatch {
+                method: "add_new_keys_to_masternode_identity_operations".to_string(),
+                known_versions: vec![0],
+                received: version,
+            })),
+        }
+    }
 }

packages/rs-drive/src/drive/identity/update/methods/add_new_keys_to_identity/v0/mod.rs

@@ -1,6 +1,7 @@
 use crate::drive::Drive;
 use crate::error::Error;
 use crate::fees::op::LowLevelDriveOperation;
+use dpp::identity::identity_public_key::methods::hash::IdentityPublicKeyHashMethodsV0;
 use dpp::identity::IdentityPublicKey;
 
 use dpp::block::epoch::Epoch;
@@ -64,4 +65,53 @@ impl Drive {
         }
         Ok(drive_operations)
     }
+
+    /// The operations for adding new non-unique keys to a masternode identity
+    #[inline(always)]
+    pub(super) fn add_new_keys_to_masternode_identity_operations_v0(
+        &self,
+        identity_id: [u8; 32],
+        non_unique_keys_to_add: Vec<IdentityPublicKey>,
+        epoch: &Epoch,
+        estimated_costs_only_with_layer_info: &mut Option<
+            HashMap<KeyInfoPath, EstimatedLayerInformation>,
+        >,
+        transaction: TransactionArg,
+        platform_version: &PlatformVersion,
+    ) -> Result<Vec<LowLevelDriveOperation>, Error> {
+        let mut drive_operations: Vec<LowLevelDriveOperation> = vec![];
+        if let Some(estimated_costs_only_with_layer_info) = estimated_costs_only_with_layer_info {
+            Self::add_estimation_costs_for_keys_for_identity_id(
+                identity_id,
+                estimated_costs_only_with_layer_info,
+                &platform_version.drive,
+            )?;
+        }
+
+        for key in non_unique_keys_to_add {
+            // When estimating costs the reference is assumed, as by `AddNewKeysToIdentity`.
+            // Otherwise a key whose hash is already a unique key hash is added without a
+            // non-unique key hash reference. The probe records no cost operation, so every
+            // other key gets exactly the operations of `AddNewKeysToIdentity`.
+            let with_reference_to_non_unique_key = estimated_costs_only_with_layer_info.is_some()
+                || !self.has_unique_public_key_hash(
+                    key.public_key_hash()?,
+                    transaction,
+                    &platform_version.drive,
+                )?;
+
+            self.insert_new_non_unique_key_operations(
+                identity_id,
+                key,
+                with_reference_to_non_unique_key,
+                true,
+                epoch,
+                estimated_costs_only_with_layer_info,
+                transaction,
+                &mut drive_operations,
+                platform_version,
+            )?;
+        }
+        Ok(drive_operations)
+    }
 }

packages/rs-drive/src/drive/identity/update/mod.rs

@@ -397,6 +397,101 @@ mod tests {
         }
     }
 
+    mod add_new_keys_to_masternode_identity {
+        use super::*;
+        use crate::util::batch::drive_op_batch::DriveLowLevelOperationConverter;
+        use crate::util::batch::IdentityOperationType;
+        use dpp::block::block_info::BlockInfo;
+        use dpp::identity::identity_public_key::v0::IdentityPublicKeyV0;
+        use dpp::identity::{IdentityV0, KeyID, KeyType, Purpose, SecurityLevel};
+        use dpp::platform_value::BinaryData;
+        use dpp::version::PlatformVersion;
+        use std::collections::{BTreeMap, HashMap};
+
+        const MASTERNODE_IDENTITY_ID: [u8; 32] = [1; 32];
+
+        fn masternode_key(id: KeyID, key_type: KeyType, purpose: Purpose) -> IdentityPublicKey {
+            IdentityPublicKeyV0 {
+                id,
+                key_type,
+                purpose,
+                security_level: SecurityLevel::CRITICAL,
+                read_only: true,
+                data: BinaryData::new(vec![(id + 2) as u8; 20]),
+                disabled_at: None,
+                contract_bounds: None,
+            }
+            .into()
+        }
+
+        #[test]
+        fn should_emit_the_operations_of_add_new_keys_to_identity_at_every_protocol_version() {
+            for protocol_version in 1..=PlatformVersion::latest().protocol_version {
+                let platform_version =
+                    PlatformVersion::get(protocol_version).expect("expected a platform version");
+                let drive = setup_drive_with_initial_state_structure(Some(platform_version));
+                let masternode_identity: Identity = IdentityV0 {
+                    id: MASTERNODE_IDENTITY_ID.into(),
+                    public_keys: BTreeMap::from([(
+                        0,
+                        masternode_key(0, KeyType::ECDSA_HASH160, Purpose::TRANSFER),
+                    )]),
+                    balance: 0,
+                    revision: 0,
+                }
+                .into();
+                drive
+                    .add_new_identity(
+                        masternode_identity,
+                        true,
+                        &BlockInfo::default(),
+                        true,
+                        None,
+                        platform_version,
+                    )
+                    .expect("expected to insert the masternode identity");
+
+                let keys = vec![
+                    masternode_key(1, KeyType::ECDSA_HASH160, Purpose::TRANSFER),
+                    masternode_key(2, KeyType::EDDSA_25519_HASH160, Purpose::SYSTEM),
+                ];
+                let operations = |operation: IdentityOperationType, estimate: bool| {
+                    let mut estimated_costs = estimate.then(HashMap::new);
+                    let operations = operation
+                        .into_low_level_drive_operations(
+                            &drive,
+                            &mut estimated_costs,
+                            &BlockInfo::default(),
+                            None,
+                            platform_version,
+                        )
+                        .expect("expected operations");
+                    (operations, estimated_costs)
+                };
+                for estimate in [false, true] {
+                    assert_eq!(
+                        operations(
+                            IdentityOperationType::AddNewKeysToMasternodeIdentity {
+                                identity_id: MASTERNODE_IDENTITY_ID,
+                                non_unique_keys_to_add: keys.clone(),
+                            },
+                            estimate
+                        ),
+                        operations(
+                            IdentityOperationType::AddNewKeysToIdentity {
+                                identity_id: MASTERNODE_IDENTITY_ID,
+                                unique_keys_to_add: vec![],
+                                non_unique_keys_to_add: keys.clone(),
+                            },
+                            estimate
+                        ),
+                        "protocol version {protocol_version}, estimate {estimate}"
+                    );
+                }
+            }
+        }
+    }
+
     mod disable_identity_keys {
         use super::*;
         use chrono::Utc;

packages/rs-drive/src/util/batch/drive_op_batch/identity.rs

@@ -55,6 +55,15 @@ pub enum IdentityOperationType {
         /// The non unique keys to be added
         non_unique_keys_to_add: Vec<IdentityPublicKey>,
     },
+    /// Adds an array of non unique keys to a masternode identity
+    /// A key whose hash is already a unique key hash is added without a non unique key hash
+    /// reference, matching how keys are stored when a masternode identity is created
+    AddNewKeysToMasternodeIdentity {
+        /// The identity id of the masternode identity
+        identity_id: [u8; 32],
+        /// The non unique keys to be added
+        non_unique_keys_to_add: Vec<IdentityPublicKey>,
+    },
     /// Disable Identity Keys
     DisableIdentityKeys {
         /// The identity id of the identity
@@ -180,6 +189,17 @@ impl DriveLowLevelOperationConverter for IdentityOperationType {
                 transaction,
                 platform_version,
             ),
+            IdentityOperationType::AddNewKeysToMasternodeIdentity {
+                identity_id,
+                non_unique_keys_to_add,
+            } => drive.add_new_keys_to_masternode_identity_operations(
+                identity_id,
+                non_unique_keys_to_add,
+                &block_info.epoch,
+                estimated_costs_only_with_layer_info,
+                transaction,
+                platform_version,
+            ),
             IdentityOperationType::DisableIdentityKeys {
                 identity_id,
                 keys_ids,