feat: Stop using blessed versions list for invariants in the registry

rs/nns/integration_tests/src/upgrades_handler.rs

@@ -166,7 +166,7 @@ fn test_submit_and_accept_update_elected_replica_versions_proposal() {
             ),
             (
                 bless_version_payload(""),
-                Some("Blessed an empty version ID"),
+                Some("Elected an empty version ID"),
             ),
             (
                 update_versions_payload(

rs/registry/canister/src/invariants/assignment.rs

@@ -4,8 +4,8 @@ use ic_base_types::{NodeId, PrincipalId};
 
 use super::{
     common::{
-        InvariantCheckError, RegistrySnapshot, get_api_boundary_node_records_from_snapshot,
-        get_node_records_from_snapshot,
+        InvariantCheckError, RegistrySnapshot, get_all_node_records,
+        get_api_boundary_node_records_from_snapshot,
     },
     subnet::get_subnet_records_map,
 };
@@ -30,7 +30,7 @@ pub(crate) fn check_node_assignment_invariants(
         .into_keys()
         .collect();
 
-    let errors: Vec<String> = get_node_records_from_snapshot(snapshot).keys().map(|node_id| {
+    let errors: Vec<String> = get_all_node_records(snapshot).iter().map(|(node_id, _)| {
             let (is_replica, is_api_boundary_node) = (
                 replicas.contains(node_id),
                 api_boundary_nodes.contains(node_id),

rs/registry/canister/src/invariants/common.rs

@@ -8,12 +8,13 @@ use std::{
 use ic_base_types::{NodeId, PrincipalId, SubnetId};
 use ic_protobuf::registry::{
     api_boundary_node::v1::ApiBoundaryNodeRecord, crypto::v1::ChainKeyEnabledSubnetList,
-    hostos_version::v1::HostosVersionRecord, node::v1::NodeRecord, subnet::v1::SubnetListRecord,
+    hostos_version::v1::HostosVersionRecord, node::v1::NodeRecord,
+    replica_version::v1::ReplicaVersionRecord, subnet::v1::SubnetListRecord,
 };
 use ic_registry_keys::{
-    CHAIN_KEY_ENABLED_SUBNET_LIST_KEY_PREFIX, HOSTOS_VERSION_KEY_PREFIX, NODE_RECORD_KEY_PREFIX,
-    get_api_boundary_node_record_node_id, get_node_record_node_id, make_node_record_key,
-    make_subnet_list_record_key,
+    CHAIN_KEY_ENABLED_SUBNET_LIST_KEY_PREFIX, HOSTOS_VERSION_KEY_PREFIX,
+    REPLICA_VERSION_KEY_PREFIX, get_api_boundary_node_record_node_id, get_node_record_node_id,
+    make_node_record_key, make_subnet_list_record_key,
 };
 use prost::Message;
 use url::Url;
@@ -46,17 +47,36 @@ impl error::Error for InvariantCheckError {
 }
 
 /// Returns all node records in the snapshot.
-pub(crate) fn get_all_node_records(snapshot: &RegistrySnapshot) -> Vec<NodeRecord> {
-    let mut nodes: Vec<NodeRecord> = Vec::new();
+pub(crate) fn get_all_node_records(snapshot: &RegistrySnapshot) -> Vec<(NodeId, NodeRecord)> {
+    let mut nodes = Vec::new();
     for (k, v) in snapshot {
-        if k.starts_with(NODE_RECORD_KEY_PREFIX.as_bytes()) {
+        if let Some(id) = get_node_record_node_id(str::from_utf8(k).unwrap()) {
             let record = NodeRecord::decode(v.as_slice()).unwrap();
-            nodes.push(record);
+            nodes.push((NodeId::from(id), record));
         }
     }
+
     nodes
 }
 
+/// Returns all replica version records in the snapshot.
+pub(crate) fn get_all_replica_version_records(
+    snapshot: &RegistrySnapshot,
+) -> Vec<(String, ReplicaVersionRecord)> {
+    let mut replica_versions = Vec::new();
+    for (k, v) in snapshot {
+        if let Some(key) = str::from_utf8(k)
+            .unwrap()
+            .strip_prefix(REPLICA_VERSION_KEY_PREFIX)
+        {
+            let record = ReplicaVersionRecord::decode(v.as_slice()).unwrap();
+            replica_versions.push((key.to_owned(), record));
+        }
+    }
+
+    replica_versions
+}
+
 pub(crate) fn get_value_from_snapshot<T: Message + Default>(
     snapshot: &RegistrySnapshot,
     key: String,
@@ -96,41 +116,13 @@ pub(crate) fn get_all_hostos_version_records(
     snapshot: &RegistrySnapshot,
 ) -> Vec<HostosVersionRecord> {
     let mut result = Vec::new();
-    for key in snapshot.keys() {
-        let hostos_version_key = String::from_utf8(key.clone()).unwrap();
-        if hostos_version_key.starts_with(HOSTOS_VERSION_KEY_PREFIX) {
-            let hostos_version_record = match snapshot.get(key) {
-                Some(hostos_version_record_bytes) => {
-                    HostosVersionRecord::decode(hostos_version_record_bytes.as_slice()).unwrap()
-                }
-                None => panic!("Cannot fetch HostosVersionRecord for an existing key"),
-            };
+    for (k, v) in snapshot {
+        if k.starts_with(HOSTOS_VERSION_KEY_PREFIX.as_bytes()) {
+            let hostos_version_record = HostosVersionRecord::decode(v.as_slice()).unwrap();
             result.push(hostos_version_record);
         }
     }
-    result
-}
 
-/// Returns all node records from the snapshot.
-pub(crate) fn get_node_records_from_snapshot(
-    snapshot: &RegistrySnapshot,
-) -> BTreeMap<NodeId, NodeRecord> {
-    let mut result = BTreeMap::<NodeId, NodeRecord>::new();
-    for key in snapshot.keys() {
-        if let Some(principal_id) =
-            get_node_record_node_id(String::from_utf8(key.clone()).unwrap().as_str())
-        {
-            // This is indeed a node record
-            let node_record = match snapshot.get(key) {
-                Some(node_record_bytes) => {
-                    NodeRecord::decode(node_record_bytes.as_slice()).unwrap()
-                }
-                None => panic!("Cannot fetch node record for an existing key"),
-            };
-            let node_id = NodeId::from(principal_id);
-            result.insert(node_id, node_record);
-        }
-    }
     result
 }
 

rs/registry/canister/src/invariants/crypto.rs

@@ -1,8 +1,8 @@
 use crate::common::LOG_PREFIX;
 use crate::invariants::{
     common::{
-        InvariantCheckError, RegistrySnapshot, get_node_records_from_snapshot,
-        get_subnet_ids_from_snapshot, get_value_from_snapshot,
+        InvariantCheckError, RegistrySnapshot, get_all_node_records, get_subnet_ids_from_snapshot,
+        get_value_from_snapshot,
     },
     subnet::get_subnet_records_map,
 };
@@ -75,13 +75,13 @@ fn check_node_crypto_keys_exist_and_are_unique(
     snapshot: &RegistrySnapshot,
 ) -> Result<(), InvariantCheckError> {
     println!("{LOG_PREFIX}node_crypto_keys_invariants_check_start");
-    let nodes = get_node_records_from_snapshot(snapshot);
+    let nodes = get_all_node_records(snapshot);
     let (pks, certs) = get_all_nodes_public_keys_and_certs(snapshot)?;
 
     let mut ok_node_count = 0;
     let mut bad_node_count = 0;
     let mut maybe_error: Option<Result<(), InvariantCheckError>> = None;
-    for node_id in nodes.keys() {
+    for node_id in nodes.iter().map(|(k, _)| k) {
         // Check that all the nodes' keys and certs are present, and node_id is consistent
         match node_has_all_keys_and_cert_and_valid_node_id(node_id, &pks, &certs) {
             Ok(()) => ok_node_count += 1,

rs/registry/canister/src/invariants/endpoint.rs

@@ -1,6 +1,4 @@
-use crate::invariants::common::{
-    InvariantCheckError, RegistrySnapshot, get_node_records_from_snapshot,
-};
+use crate::invariants::common::{InvariantCheckError, RegistrySnapshot, get_all_node_records};
 
 use std::{
     convert::TryFrom,
@@ -31,7 +29,7 @@ pub(crate) fn check_endpoint_invariants(
     strict: bool,
 ) -> Result<(), InvariantCheckError> {
     let mut valid_endpoints = BTreeSet::<(IpAddr, u16)>::new();
-    let node_records = get_node_records_from_snapshot(snapshot);
+    let node_records = get_all_node_records(snapshot);
     let common_error_prefix = format!(
         "Invariant violation detected among {} node records",
         node_records.len()

rs/registry/canister/src/invariants/firewall.rs

@@ -1,6 +1,6 @@
 use crate::invariants::common::{
-    InvariantCheckError, RegistrySnapshot, get_node_records_from_snapshot,
-    get_subnet_ids_from_snapshot, get_value_from_snapshot,
+    InvariantCheckError, RegistrySnapshot, get_all_node_records, get_subnet_ids_from_snapshot,
+    get_value_from_snapshot,
 };
 
 use std::{
@@ -36,7 +36,7 @@ pub(crate) fn check_firewall_invariants(
 ) -> Result<(), InvariantCheckError> {
     validate_firewall_rule_principals(snapshot)?;
 
-    for node_id in get_node_records_from_snapshot(snapshot).keys() {
+    for node_id in get_all_node_records(snapshot).iter().map(|(k, _)| k) {
         let node_ruleset = get_node_firewall_rules(snapshot, node_id);
         validate_firewall_ruleset(node_ruleset)?;
     }
@@ -63,8 +63,9 @@ pub(crate) fn check_firewall_invariants(
 fn validate_firewall_rule_principals(
     snapshot: &RegistrySnapshot,
 ) -> Result<(), InvariantCheckError> {
-    let mut principal_ids: BTreeSet<PrincipalId> = get_node_records_from_snapshot(snapshot)
-        .keys()
+    let mut principal_ids: BTreeSet<PrincipalId> = get_all_node_records(snapshot)
+        .iter()
+        .map(|(k, _)| k)
         .map(|s| s.get())
         .collect();
 

rs/registry/canister/src/invariants/hostos_version.rs

@@ -1,6 +1,6 @@
 use crate::invariants::common::{
     InvariantCheckError, RegistrySnapshot, assert_valid_urls_and_hash,
-    get_all_hostos_version_records, get_node_records_from_snapshot, get_value_from_snapshot,
+    get_all_hostos_version_records, get_all_node_records, get_value_from_snapshot,
 };
 
 use ic_protobuf::registry::hostos_version::v1::HostosVersionRecord;
@@ -55,8 +55,8 @@ fn get_hostos_version_record(snapshot: &RegistrySnapshot, version: String) -> Ho
 /// Returns the list of HostOS versions where each version is referred to
 /// by at least one node.
 fn get_all_hostos_versions_of_nodes(snapshot: &RegistrySnapshot) -> Vec<String> {
-    get_node_records_from_snapshot(snapshot)
-        .values()
-        .filter_map(|node_record| node_record.hostos_version_id.clone())
+    get_all_node_records(snapshot)
+        .iter()
+        .filter_map(|(_, node_record)| node_record.hostos_version_id.clone())
         .collect()
 }

rs/registry/canister/src/invariants/node_operator.rs

@@ -14,7 +14,7 @@ pub(crate) fn check_node_operator_invariants(
     strict: bool,
 ) -> Result<(), InvariantCheckError> {
     if strict {
-        for node_record in get_all_node_records(snapshot) {
+        for (_, node_record) in get_all_node_records(snapshot) {
             let node_operator_id = PrincipalId::try_from(node_record.node_operator_id).unwrap();
             let key = make_node_operator_record_key(node_operator_id);
             match snapshot.get(key.as_bytes()) {

rs/registry/canister/src/invariants/node_record.rs

@@ -1,28 +1,24 @@
-use std::collections::{BTreeMap, HashMap};
+use std::collections::HashMap;
 
 use ic_base_types::NodeId;
 use ic_nns_common::registry::MAX_NUM_SSH_KEYS;
 use ic_protobuf::registry::node::v1::NodeRecord;
 
-use crate::invariants::common::{
-    InvariantCheckError, RegistrySnapshot, get_node_records_from_snapshot,
-};
+use crate::invariants::common::{InvariantCheckError, RegistrySnapshot, get_all_node_records};
 
 pub(crate) fn check_node_record_invariants(
     snapshot: &RegistrySnapshot,
 ) -> Result<(), InvariantCheckError> {
-    let node_records = get_node_records_from_snapshot(snapshot);
+    let node_records = get_all_node_records(snapshot);
 
     check_ssh_key_limits(&node_records)?;
     check_chip_ids_are_unique(&node_records)?;
 
     Ok(())
 }
 
-fn check_ssh_key_limits(
-    node_records: &BTreeMap<NodeId, NodeRecord>,
-) -> Result<(), InvariantCheckError> {
-    for node_record in node_records.values() {
+fn check_ssh_key_limits(node_records: &[(NodeId, NodeRecord)]) -> Result<(), InvariantCheckError> {
+    for node_record in node_records.iter().map(|(_, v)| v) {
         // Enforce that the ssh_node_state_write_access field does not have too many elements.
         if node_record.ssh_node_state_write_access.len() > MAX_NUM_SSH_KEYS {
             return Err(InvariantCheckError {
@@ -47,7 +43,7 @@ fn check_ssh_key_limits(
 }
 
 fn check_chip_ids_are_unique(
-    node_records: &BTreeMap<NodeId, NodeRecord>,
+    node_records: &[(NodeId, NodeRecord)],
 ) -> Result<(), InvariantCheckError> {
     let mut chip_id_to_node: HashMap<Vec<u8>, NodeId> = HashMap::new();