feat(NNS): compare node provider ID instead of operator when removing node directly

rs/nns/integration_tests/src/registry_get_changes_since.rs

@@ -37,7 +37,7 @@ fn test_allow_opaque_caller() {
     assert_eq!(error, None);
     // The important thing is that deltas is not empty. The exact number of
     // elements is not so important.
-    assert_eq!(deltas.len(), 13);
+    assert_eq!(deltas.len(), 15);
     assert_eq!(version, 1);
 }
 
@@ -69,6 +69,6 @@ fn test_allow_self_authenticating_caller() {
     assert_eq!(error, None);
     // The important thing is that deltas is not empty. The exact number of
     // elements is not so important.
-    assert_eq!(deltas.len(), 13);
+    assert_eq!(deltas.len(), 15);
     assert_eq!(version, 1);
 }

rs/nns/test_utils/src/registry.rs

@@ -17,6 +17,7 @@ use ic_nervous_system_common_test_keys::{
 };
 use ic_protobuf::registry::{
     crypto::v1::{PublicKey, X509PublicKeyCert},
+    dc::v1::DataCenterRecord,
     node::v1::{ConnectionEndpoint, NodeRecord},
     node_operator::v1::NodeOperatorRecord,
     replica_version::v1::{BlessedReplicaVersions, ReplicaVersionRecord},
@@ -30,8 +31,9 @@ use ic_registry_canister_api::AddNodePayload;
 use ic_registry_keys::{
     make_blessed_replica_versions_key, make_catch_up_package_contents_key, make_crypto_node_key,
     make_crypto_threshold_signing_pubkey_key, make_crypto_tls_cert_key,
-    make_node_operator_record_key, make_node_record_key, make_replica_version_key,
-    make_routing_table_record_key, make_subnet_list_record_key, make_subnet_record_key,
+    make_data_center_record_key, make_node_operator_record_key, make_node_record_key,
+    make_replica_version_key, make_routing_table_record_key, make_subnet_list_record_key,
+    make_subnet_record_key,
 };
 use ic_registry_routing_table::{CanisterIdRange, RoutingTable};
 use ic_registry_subnet_type::SubnetType;
@@ -43,7 +45,7 @@ use ic_registry_transport::{
     },
     serialize_get_value_request, Error,
 };
-use ic_test_utilities_types::ids::{subnet_test_id, user_test_id};
+use ic_test_utilities_types::ids::subnet_test_id;
 use ic_types::{
     crypto::{
         threshold_sig::ni_dkg::{NiDkgTag, NiDkgTargetId, NiDkgTranscript},
@@ -249,7 +251,7 @@ pub fn invariant_compliant_mutation_with_subnet_id(
     subnet_pid: SubnetId,
     chain_key_config: Option<ChainKeyConfig>,
 ) -> Vec<RegistryMutation> {
-    let node_operator_pid = user_test_id(TEST_ID);
+    let node_operator_pid = PrincipalId::new_user_test_id(1990);
 
     let (valid_pks, node_id) = new_node_keys_and_node_id();
 
@@ -270,12 +272,14 @@ pub fn invariant_compliant_mutation_with_subnet_id(
             port: 4321,
         };
         NodeRecord {
-            node_operator_id: node_operator_pid.get().to_vec(),
+            node_operator_id: node_operator_pid.to_vec(),
             xnet: Some(xnet_connection_endpoint),
             http: Some(http_connection_endpoint),
             ..Default::default()
         }
     };
+    let dc_record = make_data_center_record("dc1");
+    let node_operator_record = make_node_operator_record(node_operator_pid, "dc1");
     const MOCK_HASH: &str = "d1bc8d3ba4afc7e109612cb73acbdddac052c93025aa1f82942edabb7deb82a1";
     let release_package_url = "http://release_package.tar.zst".to_string();
     let replica_version_id = ReplicaVersion::default().to_string();
@@ -319,6 +323,14 @@ pub fn invariant_compliant_mutation_with_subnet_id(
             blessed_replica_version.encode_to_vec(),
         ),
     ];
+    mutations.push(insert(
+        make_data_center_record_key("dc1").as_bytes(),
+        dc_record.encode_to_vec(),
+    ));
+    mutations.push(insert(
+        make_node_operator_record_key(node_operator_pid).as_bytes(),
+        node_operator_record.encode_to_vec(),
+    ));
     mutations.append(&mut make_add_node_registry_mutations(
         node_id,
         node_record,
@@ -395,11 +407,22 @@ pub fn new_node_crypto_keys_mutations(
     new_current_node_crypto_keys_mutations(node_id, current_npks)
 }
 
+/// Make a `DataCenterRecord` with the provided `dc_id`.
+pub fn make_data_center_record(dc_id: &str) -> DataCenterRecord {
+    DataCenterRecord {
+        id: dc_id.to_string(),
+        region: "Europe,ES,Barcelona".to_string(),
+        owner: "Jorge Gomez".to_string(),
+        ..Default::default()
+    }
+}
+
 /// Make a `NodeOperatorRecord` from the provided `PrincipalId`.
-fn make_node_operator_record(principal_id: PrincipalId) -> NodeOperatorRecord {
+fn make_node_operator_record(principal_id: PrincipalId, dc_id: &str) -> NodeOperatorRecord {
     NodeOperatorRecord {
         node_allowance: 1,
         node_operator_principal_id: principal_id.into(),
+        dc_id: dc_id.to_string(),
         ..Default::default()
     }
 }
@@ -511,11 +534,15 @@ pub fn initial_mutations_for_a_multinode_nns_subnet() -> Vec<RegistryMutation> {
         *TEST_USER6_PRINCIPAL,
         *TEST_USER7_PRINCIPAL,
     ] {
-        node_operator.push(make_node_operator_record(*principal_id));
+        node_operator.push(make_node_operator_record(*principal_id, "dc1"));
     }
 
     let mut add_node_mutations = vec![];
     let mut node_id = vec![];
+    add_node_mutations.push(insert(
+        make_data_center_record_key("dc1").as_bytes(),
+        make_data_center_record("dc1").encode_to_vec(),
+    ));
     for nor in &node_operator {
         let (id, mut mutations) = get_new_node_id_and_mutations(nor, nns_subnet_id);
         node_id.push(id);

rs/registry/canister/src/common/test_helpers.rs

@@ -8,10 +8,14 @@ use ic_nns_test_utils::registry::{
     new_node_keys_and_node_id,
 };
 use ic_protobuf::registry::crypto::v1::PublicKey;
+use ic_protobuf::registry::dc::v1::DataCenterRecord;
 use ic_protobuf::registry::node::v1::IPv4InterfaceConfig;
 use ic_protobuf::registry::node::v1::NodeRecord;
+use ic_protobuf::registry::node_operator::v1::NodeOperatorRecord;
 use ic_protobuf::registry::subnet::v1::SubnetListRecord;
 use ic_protobuf::registry::subnet::v1::SubnetRecord;
+use ic_registry_keys::make_data_center_record_key;
+use ic_registry_keys::make_node_operator_record_key;
 use ic_registry_keys::make_subnet_list_record_key;
 use ic_registry_keys::make_subnet_record_key;
 use ic_registry_transport::pb::v1::{
@@ -21,6 +25,7 @@ use ic_registry_transport::upsert;
 use ic_types::ReplicaVersion;
 use prost::Message;
 use std::collections::BTreeMap;
+use std::f32::consts::E;
 
 pub fn invariant_compliant_registry(mutation_id: u8) -> Registry {
     let mut registry = Registry::new();
@@ -99,6 +104,7 @@ pub fn prepare_registry_with_nodes(
 ) -> (RegistryAtomicMutateRequest, BTreeMap<NodeId, PublicKey>) {
     // Prepare a transaction to add the nodes to the registry
     let mut mutations = Vec::<RegistryMutation>::default();
+    let node_operator_principal_id = PrincipalId::new_user_test_id(1999);
     let node_ids_and_dkg_pks: BTreeMap<NodeId, PublicKey> = (0..nodes)
         .map(|id| {
             let (valid_pks, node_id) = new_node_keys_and_node_id();
@@ -115,7 +121,7 @@ pub fn prepare_registry_with_nodes(
                     ip_addr: format!("128.0.{effective_id}.1"),
                     ..Default::default()
                 }),
-                node_operator_id: PrincipalId::new_user_test_id(999).into_vec(),
+                node_operator_id: node_operator_principal_id.into_vec(),
                 // Preset this field to Some(), in order to allow seamless creation of ApiBoundaryNodeRecord if needed.
                 domain: Some(format!("node{effective_id}.example.com")),
                 ..Default::default()

rs/registry/canister/src/mutations/node_management/common.rs

@@ -55,7 +55,7 @@ pub fn get_subnet_list_record(registry: &Registry) -> SubnetListRecord {
     SubnetListRecord::decode(subnet_list_record_vec.as_slice()).unwrap()
 }
 
-pub fn get_node_operator_id_for_node(
+pub fn get_node_operator_id_for_node_id(
     registry: &Registry,
     node_id: NodeId,
 ) -> Result<PrincipalId, String> {
@@ -80,6 +80,34 @@ pub fn get_node_operator_id_for_node(
         )
 }
 
+pub fn get_node_provider_id_for_operator_id(
+    registry: &Registry,
+    node_operator_id: PrincipalId,
+) -> Result<PrincipalId, String> {
+    let node_operator_key = make_node_operator_record_key(node_operator_id);
+    registry
+        .get(node_operator_key.as_bytes(), registry.latest_version())
+        .map_or(
+            Err(format!(
+                "Node Operator Id {:} not found in the registry.",
+                node_operator_key
+            )),
+            |result| {
+                PrincipalId::try_from(
+                    NodeOperatorRecord::decode(result.value.as_slice())
+                        .unwrap()
+                        .node_provider_principal_id,
+                )
+                .map_err(|_| {
+                    format!(
+                        "Could not decode node_operator_record's node_provider_id for Node Operator Id {}",
+                        node_operator_id
+                    )
+                })
+            },
+        )
+}
+
 pub fn get_node_operator_record(
     registry: &Registry,
     node_operator_id: PrincipalId,

rs/registry/canister/src/mutations/node_management/do_remove_node_directly.rs

@@ -1,7 +1,7 @@
 use crate::mutations::node_management::common::{
-    find_subnet_for_node, get_node_operator_id_for_node, get_node_operator_record,
-    get_subnet_list_record, make_remove_node_registry_mutations,
-    make_update_node_operator_mutation,
+    find_subnet_for_node, get_node_operator_id_for_node_id, get_node_operator_record,
+    get_node_provider_id_for_operator_id, get_subnet_list_record,
+    make_remove_node_registry_mutations, make_update_node_operator_mutation,
 };
 use crate::{common::LOG_PREFIX, registry::Registry};
 use candid::{CandidType, Deserialize};
@@ -26,7 +26,7 @@ impl Registry {
     pub fn do_remove_node(&mut self, payload: RemoveNodeDirectlyPayload, caller_id: PrincipalId) {
         // 1. Find the node operator id for this record
         // and abort if the node record is not found
-        let node_operator_id = get_node_operator_id_for_node(self, payload.node_id)
+        let node_operator_id = get_node_operator_id_for_node_id(self, payload.node_id)
             .map_err(|e| {
                 format!(
                     "{}do_remove_node_directly: Aborting node removal: {}",
@@ -35,11 +35,28 @@ impl Registry {
             })
             .unwrap();
 
-        // 2. Get the caller ID and check that it matches the node's NO
+        // 2. Compare the node provider ID of the caller (Node Operator) and the node's NP
+        let node_provider_caller = get_node_provider_id_for_operator_id(self, caller_id)
+            .map_err(|e| {
+                format!(
+                    "{}do_remove_node_directly: Aborting node removal: {}",
+                    LOG_PREFIX, e
+                )
+            })
+            .unwrap();
+        let node_provider_of_the_node =
+            get_node_provider_id_for_operator_id(self, node_operator_id)
+                .map_err(|e| {
+                    format!(
+                        "{}do_remove_node_directly: Aborting node removal: {}",
+                        LOG_PREFIX, e
+                    )
+                })
+                .unwrap();
         assert_eq!(
-            node_operator_id, caller_id,
-            "The caller {}, does not match this Node's Operator id.",
-            caller_id
+            node_provider_caller, node_provider_of_the_node,
+            "The node provider {} of the caller {}, does not match the provider {} of the node {}.",
+            node_provider_caller, caller_id, node_provider_of_the_node, payload.node_id
         );
 
         // 3. Ensure node is not in a subnet

rs/registry/canister/src/mutations/node_management/do_remove_nodes.rs

@@ -1,5 +1,5 @@
 use crate::mutations::node_management::common::{
-    find_subnet_for_node, get_node_operator_id_for_node, get_node_operator_record,
+    find_subnet_for_node, get_node_operator_id_for_node_id, get_node_operator_record,
     get_subnet_list_record, make_remove_node_registry_mutations,
     make_update_node_operator_mutation,
 };
@@ -36,11 +36,11 @@ impl Registry {
 
                 // 4. Find the node operator id for this record
                 // and abort if the node record is not found
-                let node_operator_id = get_node_operator_id_for_node(self, node_to_remove)
+                let node_operator_id = get_node_operator_id_for_node_id(self, node_to_remove)
                     .map_err(|e| format!("{}do_remove_nodes: Aborting node removal: {}", LOG_PREFIX, e))
                     .unwrap();
 
-                // 5. Ensure node is not in a subnet 
+                // 5. Ensure node is not in a subnet
                 let is_node_in_subnet = find_subnet_for_node(self, node_to_remove, &subnet_list_record);
                 if let Some(subnet_id) = is_node_in_subnet {
                     panic!("{}do_remove_nodes: Cannot remove a node that is a member of a subnet. This node is a member of Subnet: {}",

rs/registry/canister/src/mutations/node_management/do_update_node_domain_directly.rs

@@ -1,4 +1,4 @@
-use crate::mutations::node_management::common::get_node_operator_id_for_node;
+use crate::mutations::node_management::common::get_node_operator_id_for_node_id;
 use crate::{common::LOG_PREFIX, registry::Registry};
 
 use candid::{CandidType, Deserialize};
@@ -43,7 +43,7 @@ impl Registry {
         let mut node_record = self.get_node_or_panic(node_id);
 
         // Ensure caller is an actual node operator of the node
-        let node_operator_id = get_node_operator_id_for_node(self, node_id)
+        let node_operator_id = get_node_operator_id_for_node_id(self, node_id)
             .map_err(|e| format!("Failed to obtain the node operator ID: {}", e))
             .unwrap();
 

rs/registry/canister/src/mutations/node_management/do_update_node_ipv4_config_directly.rs

@@ -1,5 +1,5 @@
 use crate::mutations::node_management::common::{
-    get_node_operator_id_for_node, node_exists_with_ipv4,
+    get_node_operator_id_for_node_id, node_exists_with_ipv4,
 };
 use crate::{common::LOG_PREFIX, mutations::common::node_exists_or_panic, registry::Registry};
 
@@ -84,7 +84,7 @@ impl Registry {
 
     fn check_caller_is_node_operator(&self, caller_id: PrincipalId, node_id: NodeId) {
         // Find the node operator id for this node
-        let node_operator_id = get_node_operator_id_for_node(self, node_id)
+        let node_operator_id = get_node_operator_id_for_node_id(self, node_id)
             .map_err(|e| format!("Failed to obtain the node operator ID: {}", e))
             .unwrap();
 

rs/registry/canister/tests/remove_nodes.rs

@@ -352,7 +352,7 @@ fn remove_nodes_removes_all_keys() {
 
         // Add the node along with keys and certs
         let mut mutations = make_add_node_registry_mutations(node_id, node, valid_pks);
-        // Add node operator records
+        // Add node operator record
         mutations.push(insert(
             make_node_operator_record_key(user_test_id(NO_ID).get()).as_bytes(),
             node_operator.encode_to_vec(),