diff --git a/rs/registry/canister/src/invariants/api_boundary_node.rs b/rs/registry/canister/src/invariants/api_boundary_node.rs index 7df059af4301..9aaf95e3f554 100644 --- a/rs/registry/canister/src/invariants/api_boundary_node.rs +++ b/rs/registry/canister/src/invariants/api_boundary_node.rs @@ -6,7 +6,7 @@ use std::{ use ic_base_types::NodeId; use super::common::{ - InvariantCheckError, RegistrySnapshot, get_api_boundary_node_ids_from_snapshot, + InvariantCheckError, RegistrySnapshot, get_api_boundary_node_records_from_snapshot, get_node_record_from_snapshot, }; @@ -28,8 +28,7 @@ pub(crate) fn check_api_boundary_node_invariants( // - An attempt to read the related NodeRecord for an API BN would fail and cause ReadRegistryError::Transient() // - Transient registry errors are retried in `message_route.rs` code. However, in this case it's not helpful, the error is persistent in nature // - As a result, the subnet is stalled - let api_boundary_node_ids = get_api_boundary_node_ids_from_snapshot(snapshot)?; - for api_bn_id in api_boundary_node_ids { + for api_bn_id in get_api_boundary_node_records_from_snapshot(snapshot).into_keys() { let node_record = get_node_record_from_snapshot(api_bn_id, snapshot)?; let Some(node_record) = node_record else { return Err(InvariantCheckError { diff --git a/rs/registry/canister/src/invariants/common.rs b/rs/registry/canister/src/invariants/common.rs index b9b4eaa625a8..0bfb2d67ef96 100644 --- a/rs/registry/canister/src/invariants/common.rs +++ b/rs/registry/canister/src/invariants/common.rs @@ -3,6 +3,7 @@ use std::{ convert::TryFrom, error, fmt::{Display, Formatter, Result as FmtResult}, + ops::Bound::{Included, Unbounded}, }; use ic_base_types::{NodeId, PrincipalId, SubnetId}; @@ -12,8 +13,8 @@ use ic_protobuf::registry::{ replica_version::v1::ReplicaVersionRecord, subnet::v1::SubnetListRecord, }; use ic_registry_keys::{ - 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, + API_BOUNDARY_NODE_RECORD_KEY_PREFIX, CHAIN_KEY_ENABLED_SUBNET_LIST_KEY_PREFIX, + HOSTOS_VERSION_KEY_PREFIX, NODE_RECORD_KEY_PREFIX, REPLICA_VERSION_KEY_PREFIX, make_node_record_key, make_subnet_list_record_key, }; use prost::Message; @@ -46,124 +47,97 @@ impl error::Error for InvariantCheckError { } } -/// Returns all node records in the snapshot. -pub(crate) fn get_all_node_records(snapshot: &RegistrySnapshot) -> BTreeMap { +fn scan_snapshot_for_prefix( + snapshot: &RegistrySnapshot, + prefix: &str, +) -> Result, InvariantCheckError> { let mut nodes = BTreeMap::new(); - for (k, v) in snapshot { - if let Some(id) = get_node_record_node_id(str::from_utf8(k).unwrap()) { - let record = NodeRecord::decode(v.as_slice()).unwrap(); - nodes.insert(NodeId::from(id), record); - } + + // Note: effectively .range(PREFIX..) + for (k, v) in snapshot.range::<[u8], _>((Included(prefix.as_bytes()), Unbounded)) { + let Some(k) = k.strip_prefix(prefix.as_bytes()) else { + break; + }; + + let key = str::from_utf8(k) + .map_err(|err| InvariantCheckError { + msg: format!("Failed to decode keys from the RegistrySnapshot: {err}"), + source: None, + })? + .to_string(); + let value = T::decode(v.as_slice()).map_err(|err| InvariantCheckError { + msg: format!("Deserialize registry value failed with {err}"), + source: None, + })?; + + nodes.insert(key, value); } - nodes + Ok(nodes) +} + +/// Returns all node records in the snapshot. +pub(crate) fn get_all_node_records(snapshot: &RegistrySnapshot) -> BTreeMap { + scan_snapshot_for_prefix::(snapshot, NODE_RECORD_KEY_PREFIX) + .unwrap() + .into_iter() + .map(|(k, v)| (NodeId::from(k.parse::().unwrap()), v)) + .collect() } /// Returns all replica version records in the snapshot. pub(crate) fn get_all_replica_version_records( snapshot: &RegistrySnapshot, ) -> BTreeMap { - let mut replica_versions = BTreeMap::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.insert(key.to_owned(), record); - } - } - - replica_versions -} - -pub(crate) fn get_value_from_snapshot( - snapshot: &RegistrySnapshot, - key: String, -) -> Option { - snapshot - .get(key.as_bytes()) - .map(|v| T::decode(v.as_slice()).unwrap()) + scan_snapshot_for_prefix::(snapshot, REPLICA_VERSION_KEY_PREFIX).unwrap() } // Retrieve all records that serve as lists of subnets that can sign with chain keys -#[allow(dead_code)] pub(crate) fn get_all_chain_key_signing_subnet_list_records( snapshot: &RegistrySnapshot, ) -> BTreeMap { - let mut result = BTreeMap::::new(); - for key in snapshot.keys() { - let signing_subnet_list_key = String::from_utf8(key.clone()).unwrap(); - if signing_subnet_list_key.starts_with(CHAIN_KEY_ENABLED_SUBNET_LIST_KEY_PREFIX) { - let chain_key_signing_subnet_list_record = match snapshot.get(key) { - Some(chain_key_signing_subnet_list_record) => ChainKeyEnabledSubnetList::decode( - chain_key_signing_subnet_list_record.as_slice(), - ) - .unwrap(), - None => panic!("Cannot fetch ChainKeySigningSubnetList record for an existing key"), - }; - result.insert( - signing_subnet_list_key, - chain_key_signing_subnet_list_record, - ); - } - } - result + scan_snapshot_for_prefix::( + snapshot, + CHAIN_KEY_ENABLED_SUBNET_LIST_KEY_PREFIX, + ) + .unwrap() + .into_iter() + // Preserve the prefix for crypto code, which depends on it downstream. + .map(|(k, v)| (format!("{CHAIN_KEY_ENABLED_SUBNET_LIST_KEY_PREFIX}{k}"), v)) + .collect() } // Retrieve all HostOS version records pub(crate) fn get_all_hostos_version_records( snapshot: &RegistrySnapshot, -) -> BTreeSet { - let mut result = BTreeSet::new(); - 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.insert(hostos_version_record); - } - } - - result +) -> BTreeMap { + scan_snapshot_for_prefix::(snapshot, HOSTOS_VERSION_KEY_PREFIX).unwrap() } /// Returns all api boundary node records from the snapshot. pub(crate) fn get_api_boundary_node_records_from_snapshot( snapshot: &RegistrySnapshot, ) -> BTreeMap { - let mut result = BTreeMap::::new(); - for (key, value) in snapshot.iter() { - let key_str = - String::from_utf8(key.clone()).expect("failed to convert UTF-8 byte vector to string"); - if let Some(principal_id) = get_api_boundary_node_record_node_id(&key_str) { - // This is indeed an api boundary node record - let api_boundary_node_record = ApiBoundaryNodeRecord::decode(value.as_slice()).unwrap(); - let node_id = NodeId::from(principal_id); - result.insert(node_id, api_boundary_node_record); - } - } - result + scan_snapshot_for_prefix::(snapshot, API_BOUNDARY_NODE_RECORD_KEY_PREFIX) + .unwrap() + .into_iter() + .map(|(k, v)| (NodeId::from(k.parse::().unwrap()), v)) + .collect() } -/// Returns an all api boundary node ids record from the registry snapshot. -pub(crate) fn get_api_boundary_node_ids_from_snapshot( +pub(crate) fn get_value_from_snapshot( snapshot: &RegistrySnapshot, -) -> Result, InvariantCheckError> { + key: String, +) -> Result, InvariantCheckError> { snapshot - .keys() - .cloned() - .map(|key| { - String::from_utf8(key).map_err(|_| InvariantCheckError { - msg: "Failed to decode keys from the RegistrySnapshot".to_string(), + .get(key.as_bytes()) + .map(|v| { + T::decode(v.as_slice()).map_err(|err| InvariantCheckError { + msg: format!("Deserialize registry value failed with {err}"), source: None, }) }) - .collect::, InvariantCheckError>>() - .map(|keys| { - keys.into_iter() - .filter_map(|key_str| get_api_boundary_node_record_node_id(&key_str)) - .map(NodeId::from) - .collect() - }) + .transpose() } /// Returns node record from the snapshot corresponding to a key. @@ -171,39 +145,21 @@ pub(crate) fn get_node_record_from_snapshot( key: NodeId, snapshot: &RegistrySnapshot, ) -> Result, InvariantCheckError> { - let key = make_node_record_key(key); - let value = snapshot.get(key.as_bytes()); - value - .map(|bytes| { - NodeRecord::decode(bytes.as_slice()).map_err(|err| InvariantCheckError { - msg: format!("Deserialize registry value failed with {err}"), - source: None, - }) - }) - .transpose() + get_value_from_snapshot::(snapshot, make_node_record_key(key)) } pub(crate) fn get_subnet_ids_from_snapshot(snapshot: &RegistrySnapshot) -> BTreeSet { get_value_from_snapshot::(snapshot, make_subnet_list_record_key()) + .unwrap() .map(|r| { r.subnets .iter() - .map(|s| SubnetId::from(PrincipalId::try_from(s.clone().as_slice()).unwrap())) + .map(|s| SubnetId::from(PrincipalId::try_from(s).unwrap())) .collect() }) .unwrap_or_default() } -pub(crate) fn assert_sha256(s: &str) { - if s.bytes().any(|x| !x.is_ascii_hexdigit()) { - panic!("Hash contains at least one invalid character: `{s}`"); - } - - if s.len() != 64 { - panic!("Hash is an invalid length: `{s}`"); - } -} - pub(crate) fn assert_valid_urls_and_hash(urls: &[String], hash: &str, allow_file_url: bool) { // Either both, the URL and the hash are set, or both are not set. if (urls.is_empty() as i32 ^ hash.is_empty() as i32) > 0 { @@ -213,7 +169,13 @@ pub(crate) fn assert_valid_urls_and_hash(urls: &[String], hash: &str, allow_file return; } - assert_sha256(hash); + if hash.bytes().any(|x| !x.is_ascii_hexdigit()) { + panic!("Hash contains at least one invalid character: `{hash}`"); + } + + if hash.len() != 64 { + panic!("Hash is an invalid length: `{hash}`"); + } urls.iter().for_each(|url| // File URLs are used in test deployments. We only disallow non-ASCII. @@ -254,7 +216,7 @@ mod tests { } #[test] - #[should_panic(expected = "DecodeError")] + #[should_panic(expected = "failed to decode")] fn test_get_api_boundary_node_records_from_snapshot_with_wrongly_encoded_record() { let mut snapshot = RegistrySnapshot::new(); let node_id: NodeId = PrincipalId::new_node_test_id(0).into(); @@ -267,7 +229,7 @@ mod tests { } #[test] - #[should_panic(expected = "DecodeError")] + #[should_panic(expected = "failed to decode")] fn test_get_value_from_snapshot_panics() { let mut snapshot = RegistrySnapshot::new(); let node_id: NodeId = PrincipalId::new_node_test_id(0).into(); @@ -277,6 +239,6 @@ mod tests { vec![0], // incorrect value, not an encoded ApiBoundaryNodeRecord ); // this call should panic when decoding the ApiBoundaryNodeRecord - get_value_from_snapshot::(&snapshot, key); + get_value_from_snapshot::(&snapshot, key).unwrap(); } } diff --git a/rs/registry/canister/src/invariants/crypto.rs b/rs/registry/canister/src/invariants/crypto.rs index 423218a43d6e..0c42bd2c091d 100644 --- a/rs/registry/canister/src/invariants/crypto.rs +++ b/rs/registry/canister/src/invariants/crypto.rs @@ -432,9 +432,9 @@ fn check_high_threshold_public_key_matches_the_one_in_cup( let high_threshold_public_key_bytes: Option = get_value_from_snapshot( snapshot, make_crypto_threshold_signing_pubkey_key(subnet_id), - ); + )?; let cup_contents_bytes: Option = - get_value_from_snapshot(snapshot, make_catch_up_package_contents_key(subnet_id)); + get_value_from_snapshot(snapshot, make_catch_up_package_contents_key(subnet_id))?; if let (Some(high_threshold_public_key_proto), Some(cup_contents)) = (high_threshold_public_key_bytes, cup_contents_bytes) { diff --git a/rs/registry/canister/src/invariants/firewall.rs b/rs/registry/canister/src/invariants/firewall.rs index d0d544984b75..2c4fedf0a425 100644 --- a/rs/registry/canister/src/invariants/firewall.rs +++ b/rs/registry/canister/src/invariants/firewall.rs @@ -282,6 +282,8 @@ fn get_firewall_rules(snapshot: &RegistrySnapshot, record_key: String) -> Option if snapshot.contains_key(record_key.as_bytes()) { Some( get_value_from_snapshot(snapshot, record_key.clone()) + .ok() + .flatten() .unwrap_or_else(|| panic!("Could not find firewall rules: {record_key}")), ) } else { diff --git a/rs/registry/canister/src/invariants/hostos_version.rs b/rs/registry/canister/src/invariants/hostos_version.rs index ce488d87eda0..d5ca6f40da04 100644 --- a/rs/registry/canister/src/invariants/hostos_version.rs +++ b/rs/registry/canister/src/invariants/hostos_version.rs @@ -1,10 +1,10 @@ +use std::collections::BTreeSet; + use crate::invariants::common::{ InvariantCheckError, RegistrySnapshot, assert_valid_urls_and_hash, - get_all_hostos_version_records, get_all_node_records, get_value_from_snapshot, + get_all_hostos_version_records, get_all_node_records, }; -use ic_protobuf::registry::hostos_version::v1::HostosVersionRecord; -use ic_registry_keys::make_hostos_version_key; use ic_types::hostos_version::HostosVersion; /// A predicate on the HostOS version records contained in a registry @@ -20,34 +20,36 @@ use ic_types::hostos_version::HostosVersion; pub(crate) fn check_hostos_version_invariants( snapshot: &RegistrySnapshot, ) -> Result<(), InvariantCheckError> { - let mut all_versions = Vec::new(); - // Collect all referenced HostOS versions - let node_versions = get_all_hostos_versions_of_nodes(snapshot); - all_versions.extend(node_versions); + let versions_in_use = get_all_hostos_versions_of_nodes(snapshot); - // Get the current list of registered HostOS versions - let registered_versions = get_all_hostos_version_records(snapshot); + // Re-collect since we can't compare `BTreeSet` with `BTreeSet<&String>` with `is_superset`. + let versions_in_use: BTreeSet<_> = versions_in_use.iter().collect(); - all_versions.extend(registered_versions.into_iter().map(|v| v.hostos_version_id)); - all_versions.dedup(); - - for version in all_versions { + // Get the current list of registered HostOS versions + let elected_versions = get_all_hostos_version_records(snapshot); + let elected_set: BTreeSet<_> = elected_versions.keys().collect(); + assert!( + elected_set.is_superset(&versions_in_use), + "Using a version that isn't elected. Elected versions: {elected_set:?}, in use: {versions_in_use:?}." + ); + assert!( + elected_set.iter().all(|v| !v.trim().is_empty()), + "Elected an empty version ID." + ); + + for (key, record) in elected_versions { // Enforce that the version ID is well-formed, so that consumers reading // it back out of the Registry can turn it into a HostosVersion. - if let Err(err) = HostosVersion::try_from(version.as_str()) { + if let Err(err) = HostosVersion::try_from(key.as_str()) { panic!("Registered an invalid HostOS version ID: {err}"); } - // Check that every referenced version exists, i.e. we can only set a - // Node's version to one that has already been added to the registry. - let r = get_hostos_version_record(snapshot, version); - // Check whether release package URLs (update image) and corresponding hash are well-formed. // As file-based URLs are only used in test-deployments, we disallow file:/// URLs. assert_valid_urls_and_hash( - &r.release_package_urls, - &r.release_package_sha256_hex, + &record.release_package_urls, + &record.release_package_sha256_hex, false, ); } @@ -55,14 +57,9 @@ pub(crate) fn check_hostos_version_invariants( Ok(()) } -fn get_hostos_version_record(snapshot: &RegistrySnapshot, version: String) -> HostosVersionRecord { - get_value_from_snapshot(snapshot, make_hostos_version_key(version.clone())) - .unwrap_or_else(|| panic!("Could not find HostOS version: {version}")) -} - /// 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 { +fn get_all_hostos_versions_of_nodes(snapshot: &RegistrySnapshot) -> BTreeSet { get_all_node_records(snapshot) .into_values() .filter_map(|node_record| node_record.hostos_version_id) @@ -71,9 +68,9 @@ fn get_all_hostos_versions_of_nodes(snapshot: &RegistrySnapshot) -> Vec #[cfg(test)] mod tests { - use super::*; - use crate::common::test_helpers::invariant_compliant_registry; + use ic_protobuf::registry::hostos_version::v1::HostosVersionRecord; + use ic_registry_keys::make_hostos_version_key; use ic_registry_transport::{insert, pb::v1::RegistryMutation}; use prost::Message; diff --git a/rs/registry/canister/src/invariants/replica_version.rs b/rs/registry/canister/src/invariants/replica_version.rs index cd31b0c95fcd..02a4b25f405b 100644 --- a/rs/registry/canister/src/invariants/replica_version.rs +++ b/rs/registry/canister/src/invariants/replica_version.rs @@ -55,9 +55,11 @@ pub(crate) fn check_replica_version_invariants( versions_in_use.append(&mut get_all_standard_engine_replica_versions(snapshot)); versions_in_use.append(&mut get_all_api_boundary_node_versions(snapshot)); - let elected_set: BTreeSet<_> = get_all_replica_version_records(snapshot) - .into_keys() - .collect(); + // Re-collect since we can't compare `BTreeSet` with `BTreeSet<&String>` with `is_superset`. + let versions_in_use: BTreeSet<_> = versions_in_use.iter().collect(); + + let elected_versions = get_all_replica_version_records(snapshot); + let elected_set: BTreeSet<_> = elected_versions.keys().collect(); assert!( elected_set.is_superset(&versions_in_use), "Using a version that isn't elected. Elected versions: {elected_set:?}, in use: {versions_in_use:?}." @@ -67,44 +69,39 @@ pub(crate) fn check_replica_version_invariants( "Elected an empty version ID." ); - for version in elected_set { + for (key, record) in elected_versions { // Enforce that the version ID is well-formed, so that consumers reading // it back out of the Registry can turn it into a ReplicaVersion. - if let Err(err) = ReplicaVersion::try_from(version.as_str()) { + if let Err(err) = ReplicaVersion::try_from(key.as_str()) { panic!("Elected an invalid version ID: {err}"); } - let r = get_replica_version_record(snapshot, &version); - // Check whether release package URLs (update image) and corresponding hash are well-formed. // As file-based URLs are only used in test-deployments, we disallow file:/// URLs. assert_valid_urls_and_hash( - &r.release_package_urls, - &r.release_package_sha256_hex, + &record.release_package_urls, + &record.release_package_sha256_hex, false, // allow_file_url ); // Check that all measured versions are valid - if let Some(Err(defects)) = r.guest_launch_measurements.map(|v| v.validate()) { + if let Some(Err(defects)) = record.guest_launch_measurements.map(|v| v.validate()) { panic!("guest_launch_measurements are not valid. Defects: {defects:?}"); } // Enforce that the stored version always matches the key - if let Some(replica_version_id) = r.replica_version_id { - assert_eq!(replica_version_id, version); + if let Some(replica_version_id) = record.replica_version_id { + assert_eq!(replica_version_id, key); } } Ok(()) } -fn get_replica_version_record(snapshot: &RegistrySnapshot, version: &str) -> ReplicaVersionRecord { - get_value_from_snapshot(snapshot, make_replica_version_key(version)) - .unwrap_or_else(|| panic!("Could not find replica version: {version}")) -} - fn get_subnet_record(snapshot: &RegistrySnapshot, subnet_id: SubnetId) -> SubnetRecord { get_value_from_snapshot(snapshot, make_subnet_record_key(subnet_id)) + .ok() + .flatten() .unwrap_or_else(|| panic!("Could not get subnet record for subnet: {subnet_id}")) } @@ -164,6 +161,8 @@ pub(crate) fn has_launch_measurements( snapshot, make_replica_version_key(replica_version_id), ) + .ok() + .flatten() .and_then(|replica_version_record| replica_version_record.guest_launch_measurements) .is_some() } diff --git a/rs/registry/canister/src/invariants/standard_engine_replica_version.rs b/rs/registry/canister/src/invariants/standard_engine_replica_version.rs index 8be4bb6c85c2..983373766c35 100644 --- a/rs/registry/canister/src/invariants/standard_engine_replica_version.rs +++ b/rs/registry/canister/src/invariants/standard_engine_replica_version.rs @@ -22,7 +22,8 @@ pub(crate) fn check_standard_engine_replica_version_invariants( let Some(record) = get_value_from_snapshot::( snapshot, make_standard_engine_replica_version_record_key(), - ) else { + )? + else { // If there is no record yet, then we are trivially valid. return Ok(()); }; diff --git a/rs/registry/canister/src/invariants/subnet.rs b/rs/registry/canister/src/invariants/subnet.rs index 99b51aa80ccf..87ecf76170b1 100644 --- a/rs/registry/canister/src/invariants/subnet.rs +++ b/rs/registry/canister/src/invariants/subnet.rs @@ -201,7 +201,8 @@ fn check_default_initial_dkg_subnet_invariant( let Some(subnet_id_proto) = get_value_from_snapshot::( snapshot, make_default_initial_dkg_subnet_id_key(), - ) else { + )? + else { return Ok(()); }; diff --git a/rs/registry/canister/src/invariants/unassigned_nodes_config.rs b/rs/registry/canister/src/invariants/unassigned_nodes_config.rs index 8bfc9bc73474..b1b13f9f3b02 100644 --- a/rs/registry/canister/src/invariants/unassigned_nodes_config.rs +++ b/rs/registry/canister/src/invariants/unassigned_nodes_config.rs @@ -20,7 +20,7 @@ pub(crate) fn check_unassigned_nodes_config_invariants( if let Some(config) = get_value_from_snapshot::( snapshot, make_unassigned_nodes_config_record_key(), - ) && config.ssh_readonly_access.len() > MAX_NUM_SSH_KEYS + )? && config.ssh_readonly_access.len() > MAX_NUM_SSH_KEYS { return Err(InvariantCheckError { msg: format!(