add sev_enabled subnet invariant

r-birkner · r-birkner · commit 9570db88246a · 2026-02-25T07:06:09.000Z
diff --git a/rs/registry/canister/src/common/test_helpers.rs b/rs/registry/canister/src/common/test_helpers.rs
@@ -112,6 +112,28 @@ pub fn prepare_registry_with_nodes_and_node_operator_id(
     start_mutation_id: u8,
     nodes: u64,
     node_operator_id: PrincipalId,
+) -> (RegistryAtomicMutateRequest, BTreeMap<NodeId, PublicKey>) {
+    prepare_registry_raw(start_mutation_id, nodes, node_operator_id, false)
+}
+
+/// Same as above, just with the possibility to have a chip_id.
+pub fn prepare_registry_with_nodes_and_chip_id(
+    start_mutation_id: u8,
+    nodes: u64,
+) -> (RegistryAtomicMutateRequest, BTreeMap<NodeId, PublicKey>) {
+    prepare_registry_raw(
+        start_mutation_id,
+        nodes,
+        PrincipalId::new_user_test_id(999),
+        true,
+    )
+}
+
+pub fn prepare_registry_raw(
+    start_mutation_id: u8,
+    nodes: u64,
+    node_operator_id: PrincipalId,
+    with_chip_id: bool,
 ) -> (RegistryAtomicMutateRequest, BTreeMap<NodeId, PublicKey>) {
     // Prepare a transaction to add the nodes to the registry
     let mut mutations = Vec::<RegistryMutation>::default();
@@ -135,6 +157,11 @@ pub fn prepare_registry_with_nodes_and_node_operator_id(
                 // Preset this field to Some(), in order to allow seamless creation of ApiBoundaryNodeRecord if needed.
                 domain: Some(format!("node{effective_id}.example.com")),
                 node_reward_type: Some(NodeRewardType::Type1 as i32),
+                chip_id: if with_chip_id {
+                    Some(format!("chip-id-{effective_id}").into_bytes())
+                } else {
+                    None
+                },
                 ..Default::default()
             };
             mutations.append(&mut make_add_node_registry_mutations(
diff --git a/rs/registry/canister/src/invariants/subnet.rs b/rs/registry/canister/src/invariants/subnet.rs
@@ -4,7 +4,8 @@ use std::{
 };
 
 use crate::invariants::common::{
-    InvariantCheckError, RegistrySnapshot, get_subnet_ids_from_snapshot,
+    InvariantCheckError, RegistrySnapshot, get_node_record_from_snapshot,
+    get_subnet_ids_from_snapshot,
 };
 
 use ic_base_types::{NodeId, PrincipalId};
@@ -36,6 +37,7 @@ pub(crate) fn check_subnet_invariants(
                 panic!("Subnet {subnet_id:} is in subnet list but no record exists")
             });
 
+        // Each SSH key access list does not contain more than 50 keys
         if subnet_record.ssh_readonly_access.len() > MAX_NUM_SSH_KEYS
             || subnet_record.ssh_backup_access.len() > MAX_NUM_SSH_KEYS
         {
@@ -52,42 +54,43 @@ pub(crate) fn check_subnet_invariants(
             });
         }
 
-        let num_nodes = subnet_record.membership.len();
-        let mut subnet_members: HashSet<NodeId> = subnet_record
+        let subnet_members: HashSet<NodeId> = subnet_record
             .membership
             .iter()
             .map(|v| NodeId::from(PrincipalId::try_from(v).unwrap()))
             .collect();
 
         // Subnet membership must contain registered nodes only
-        subnet_members.retain(|&k| {
-            let node_key = make_node_record_key(k);
-            let node_exists = snapshot.contains_key(node_key.as_bytes());
-            if !node_exists {
-                panic!("Node {k} does not exist in Subnet {subnet_id}");
+        for &node_id in &subnet_members {
+            let node_key = make_node_record_key(node_id);
+            if !snapshot.contains_key(node_key.as_bytes()) {
+                panic!("Node {node_id} does not exist in Subnet {subnet_id}");
             }
-            node_exists
-        });
+        }
 
         // Each node appears at most once in a subnet membership
+        let num_nodes = subnet_record.membership.len();
         if num_nodes > subnet_members.len() {
             panic!("Repeated nodes in subnet {subnet_id:}");
         }
+
         // Each subnet contains at least one node
         if subnet_members.is_empty() {
             panic!("No node in subnet {subnet_id:}");
         }
+
+        // Each node appears at most once in at most one subnet membership
         let intersection = accumulated_nodes_in_subnets
             .intersection(&subnet_members)
             .collect::<HashSet<_>>();
-        // Each node appears at most once in at most one subnet membership
         if !intersection.is_empty() {
             return Err(InvariantCheckError {
                 msg: format!("Nodes in subnet {subnet_id:} also belong to other subnets"),
                 source: None,
             });
         }
         accumulated_nodes_in_subnets.extend(&subnet_members);
+
         // Count occurrence of system subnets
         if subnet_record.subnet_type == i32::from(SubnetType::System) {
             system_subnet_count += 1;
@@ -105,6 +108,26 @@ pub(crate) fn check_subnet_invariants(
                 source: None,
             });
         }
+
+        // SEV-enabled subnets consist of SEV-enabled nodes only (i.e. nodes with a chip ID in the node record)
+        if let Some(features) = subnet_record.features.as_ref() {
+            if features.sev_enabled == Some(true) {
+                for &node_id in &subnet_members {
+                    // handle missing node record
+                    let node_record = get_node_record_from_snapshot(node_id, snapshot)?
+                        .ok_or_else(|| InvariantCheckError {
+                            msg: format!("Subnet {subnet_id} has node {node_id} in its membership but the node record does not exist"),
+                            source: None,
+                    })?;
+
+                    // handle missing chip_id
+                    node_record.chip_id.as_ref().ok_or_else(|| InvariantCheckError {
+                        msg: format!("Subnet {subnet_id} is SEV-enabled but at least one of its nodes is not: {node_id} does not have a chip ID in its node record"),
+                        source: None,
+                    })?;
+                }
+            }
+        }
     }
 
     // There is at least one system subnet. Note that we disable this invariant for benchmarks, as
diff --git a/rs/registry/canister/src/invariants/subnet/tests.rs b/rs/registry/canister/src/invariants/subnet/tests.rs
@@ -4,7 +4,9 @@ use crate::invariants::common::RegistrySnapshot;
 use ic_base_types::SubnetId;
 use ic_protobuf::registry::{
     node::v1::NodeRecord,
-    subnet::v1::{CanisterCyclesCostSchedule, SubnetListRecord, SubnetRecord, SubnetType},
+    subnet::v1::{
+        CanisterCyclesCostSchedule, SubnetFeatures, SubnetListRecord, SubnetRecord, SubnetType,
+    },
 };
 use ic_registry_keys::{make_subnet_list_record_key, make_subnet_record_key};
 use ic_test_utilities_types::ids::{node_test_id, subnet_test_id};
@@ -17,6 +19,8 @@ fn only_application_subnets_can_be_free_cycles_cost_schedule() {
         setup_minimal_registry_snapshot_for_check_subnet_invariants(
             system_subnet_id,
             test_subnet_id,
+            1,
+            false,
         );
 
     // Trivial case. (Never forget the trivial case, because this is an edge
@@ -54,9 +58,94 @@ fn only_application_subnets_can_be_free_cycles_cost_schedule() {
     );
 }
 
+#[test]
+fn only_sev_enabled_subnets_consist_of_sev_enabled_nodes() {
+    let system_subnet_id = subnet_test_id(1);
+    let test_subnet_id = subnet_test_id(2);
+    let test_node_id = node_test_id(103);
+
+    // get a snapshot without chip IDs
+    let (mut snapshot, mut test_subnet_record) =
+        setup_minimal_registry_snapshot_for_check_subnet_invariants(
+            system_subnet_id,
+            test_subnet_id,
+            4,
+            false,
+        );
+
+    // a non SEV-enabled subnet only with nodes without chip ID is compliant
+    check_subnet_invariants(&snapshot).unwrap();
+
+    // a non SEV-enabled subnet with some nodes with and some without chip IDs is compliant
+    let mut test_node_record = NodeRecord::default();
+    test_node_record.chip_id = Some("a chip id".into());
+    snapshot.insert(
+        make_node_record_key(test_node_id).into_bytes(),
+        test_node_record.encode_to_vec(),
+    );
+    check_subnet_invariants(&snapshot).unwrap();
+
+    // an SEV-enabled subnet only with nodes without chip ID is NOT compliant
+    test_subnet_record.features = Some(SubnetFeatures {
+        sev_enabled: Some(true),
+        ..Default::default()
+    });
+    snapshot.insert(
+        make_subnet_record_key(test_subnet_id).into_bytes(),
+        test_subnet_record.encode_to_vec(),
+    );
+    let mut test_node_record = NodeRecord::default();
+    test_node_record.chip_id = None;
+    snapshot.insert(
+        make_node_record_key(test_node_id).into_bytes(),
+        test_node_record.encode_to_vec(),
+    );
+    assert_non_compliant_record(
+        &snapshot,
+        "subnet fbysm-3acaa-aaaaa-aaaap-yai is sev-enabled but at least one of its nodes is not",
+    );
+
+    // get a snapshot with chip IDs
+    let (mut snapshot, mut test_subnet_record) =
+        setup_minimal_registry_snapshot_for_check_subnet_invariants(
+            system_subnet_id,
+            test_subnet_id,
+            4,
+            true,
+        );
+
+    // a non SEV-enabled subnet only with nodes without chip ID is compliant
+    check_subnet_invariants(&snapshot).unwrap();
+
+    // an SEV-enabled subnet only with nodes with chip ID is compliant
+    test_subnet_record.features = Some(SubnetFeatures {
+        sev_enabled: Some(true),
+        ..Default::default()
+    });
+    snapshot.insert(
+        make_subnet_record_key(test_subnet_id).into_bytes(),
+        test_subnet_record.encode_to_vec(),
+    );
+    check_subnet_invariants(&snapshot).unwrap();
+
+    // an SEV-enabled subnet with some nodes with and some without chip ID is NOT compliant
+    let mut test_node_record = NodeRecord::default();
+    test_node_record.chip_id = None;
+    snapshot.insert(
+        make_node_record_key(test_node_id).into_bytes(),
+        test_node_record.encode_to_vec(),
+    );
+    assert_non_compliant_record(
+        &snapshot,
+        "subnet fbysm-3acaa-aaaaa-aaaap-yai is sev-enabled but at least one of its nodes is not",
+    );
+}
+
 fn setup_minimal_registry_snapshot_for_check_subnet_invariants(
     system_subnet_id: SubnetId,
     test_subnet_id: SubnetId,
+    num_nodes_in_test_subnet: usize,
+    with_chip_id: bool,
 ) -> (RegistrySnapshot, SubnetRecord) {
     let mut snapshot = RegistrySnapshot::new();
 
@@ -84,11 +173,18 @@ fn setup_minimal_registry_snapshot_for_check_subnet_invariants(
     );
 
     // Add a test subnet in the subnet list.
-    let test_node_id = node_test_id(100);
-    snapshot.insert(
-        make_node_record_key(test_node_id.to_owned()).into_bytes(),
-        NodeRecord::default().encode_to_vec(),
-    );
+    for i in 0..num_nodes_in_test_subnet {
+        let node_id = node_test_id((i + 100) as u64);
+        let mut node_record = NodeRecord::default();
+        if with_chip_id {
+            node_record.chip_id = Some(format!("chip-id-{i}").into_bytes());
+        }
+        snapshot.insert(
+            make_node_record_key(node_id.to_owned()).into_bytes(),
+            node_record.encode_to_vec(),
+        );
+    }
+
     let subnet_list_record = SubnetListRecord {
         subnets: vec![
             system_subnet_id.get().into_vec(),
@@ -99,8 +195,12 @@ fn setup_minimal_registry_snapshot_for_check_subnet_invariants(
         make_subnet_list_record_key().into_bytes(),
         subnet_list_record.encode_to_vec(),
     );
+    let membership = (0..num_nodes_in_test_subnet)
+        .map(|i| node_test_id((i + 100) as u64).get().to_vec())
+        .collect();
+
     let test_subnet_record = SubnetRecord {
-        membership: vec![test_node_id.get().to_vec()],
+        membership: membership,
         ..Default::default()
     };
     snapshot.insert(
@@ -116,5 +216,6 @@ fn assert_non_compliant_record(snapshot: &RegistrySnapshot, error_msg: &str) {
         panic!("Expected Err, but got Ok!");
     };
     let message = err.msg.to_lowercase();
+    println!("Error message: {message}");
     assert!(message.contains(error_msg));
 }
diff --git a/rs/registry/canister/src/mutations/do_update_subnet.rs b/rs/registry/canister/src/mutations/do_update_subnet.rs
@@ -508,7 +508,7 @@ mod tests {
     use super::*;
     use crate::common::test_helpers::{
         add_fake_subnet, get_invariant_compliant_subnet_record, invariant_compliant_registry,
-        prepare_registry_with_nodes,
+        prepare_registry_with_nodes, prepare_registry_with_nodes_and_chip_id,
     };
     use ic_management_canister_types_private::{
         EcdsaCurve, EcdsaKeyId, SchnorrAlgorithm, SchnorrKeyId, VetKdCurve, VetKdKeyId,
@@ -979,7 +979,7 @@ mod tests {
     fn make_registry_for_update_subnet_tests() -> (Registry, SubnetId) {
         let mut registry = invariant_compliant_registry(0);
 
-        let (mutate_request, node_ids_and_dkg_pks) = prepare_registry_with_nodes(1, 2);
+        let (mutate_request, node_ids_and_dkg_pks) = prepare_registry_with_nodes_and_chip_id(1, 2);
         registry.maybe_apply_mutation_internal(mutate_request.mutations);
 
         let mut subnet_list_record = registry.get_subnet_list_record();
