diff --git a/Cargo.lock b/Cargo.lock index d5af46befeef..a4955bbe682b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -850,13 +850,8 @@ dependencies = [ "der 0.7.9", "getrandom 0.2.15", "hex", - "ic-interfaces-registry", - "ic-protobuf", - "ic-registry-client-helpers", - "ic-types", "ic-utils 0.9.0", "ic-utils-rustfmt", - "itertools 0.12.1", "prost 0.13.4", "prost-build 0.13.4", "rand 0.8.5", @@ -873,13 +868,7 @@ version = "0.0.0" dependencies = [ "attestation", "der 0.7.9", - "ic-interfaces-registry-mocks", - "ic-protobuf", - "ic-registry-keys", - "ic-types", - "mockall", "p384", - "prost 0.13.4", "rand 0.8.5", "rsa", "sev", @@ -5378,6 +5367,7 @@ dependencies = [ "ic-crypto-utils-threshold-sig-der", "ic-interfaces-registry", "ic-registry-client", + "ic-registry-client-helpers", "ic-registry-nns-data-provider-wrappers", "rcgen 0.13.2", "rustls 0.23.27", @@ -5403,6 +5393,7 @@ dependencies = [ "hyper 1.8.1", "hyper-util", "ic-interfaces-registry", + "ic-registry-client-helpers", "rcgen 0.13.2", "rustls 0.23.27", "sev", @@ -5435,7 +5426,6 @@ version = "0.0.0" dependencies = [ "anyhow", "attestation", - "attestation_testing", "config_types", "futures", "guest_upgrade_client", @@ -5443,6 +5433,9 @@ dependencies = [ "guest_upgrade_shared", "ic-interfaces-registry", "ic-protobuf", + "ic-registry-client-fake", + "ic-registry-proto-data-provider", + "ic-test-utilities-registry", "rustls 0.23.27", "sev_guest", "sev_guest_testing", diff --git a/rs/ic_os/guest_upgrade/client/BUILD.bazel b/rs/ic_os/guest_upgrade/client/BUILD.bazel index 64ff8f00e008..d82a2d84033b 100644 --- a/rs/ic_os/guest_upgrade/client/BUILD.bazel +++ b/rs/ic_os/guest_upgrade/client/BUILD.bazel @@ -23,6 +23,7 @@ DEPENDENCIES = [ "//rs/ic_os/sev/guest", "//rs/interfaces/registry", "//rs/registry/client", + "//rs/registry/helpers", "//rs/registry/nns_data_provider_wrappers", ] diff --git a/rs/ic_os/guest_upgrade/client/Cargo.toml b/rs/ic_os/guest_upgrade/client/Cargo.toml index a12fc1dbc0bf..dc8704a97122 100644 --- a/rs/ic_os/guest_upgrade/client/Cargo.toml +++ b/rs/ic_os/guest_upgrade/client/Cargo.toml @@ -18,6 +18,7 @@ guest_disk = { path = "../../os_tools/guest_disk" } guest_upgrade_shared = { path = "../shared" } ic-interfaces-registry = { path = "../../../interfaces/registry" } ic-registry-client = { path = "../../../registry/client" } +ic-registry-client-helpers = { path = "../../../registry/helpers" } ic-registry-nns-data-provider-wrappers = { path = "../../../registry/nns_data_provider_wrappers" } ic-crypto-utils-threshold-sig-der = { path = "../../../crypto/utils/threshold_sig_der" } sev_guest = { path = "../../sev/guest" } diff --git a/rs/ic_os/guest_upgrade/client/src/lib.rs b/rs/ic_os/guest_upgrade/client/src/lib.rs index 4741f0d41736..8e0c0024de40 100644 --- a/rs/ic_os/guest_upgrade/client/src/lib.rs +++ b/rs/ic_os/guest_upgrade/client/src/lib.rs @@ -3,7 +3,6 @@ use anyhow::{Context, Error, Result, anyhow, bail}; use attestation::attestation_package::{ AttestationPackageVerifier, ParsedSevAttestationPackage, SevRootCertificateVerification, }; -use attestation::registry::get_blessed_guest_launch_measurements_from_registry; use config_types::GuestOSConfig; use der::asn1::OctetStringRef; use guest_upgrade_shared::STORE_DEVICE; @@ -16,6 +15,7 @@ use hyper_util::rt::TokioIo; use ic_crypto_utils_threshold_sig_der::parse_threshold_sig_key_from_pem_file; use ic_interfaces_registry::RegistryClient; use ic_registry_client::client::RegistryClientImpl; +use ic_registry_client_helpers::blessed_replica_version::BlessedReplicaVersionRegistry; use ic_registry_nns_data_provider_wrappers::CertifiedNnsDataProvider; use rcgen::CertifiedKey; use rustls::ClientConfig; @@ -165,9 +165,11 @@ impl DiskEncryptionKeyExchangeClientAgent { .sev_attestation_package .context("Server attestation report is missing")?; - let blessed_measurements = - get_blessed_guest_launch_measurements_from_registry(&*self.nns_registry_client) - .map_err(|e| anyhow!("Failed to get blessed measurements from registry: {e}"))?; + let registry_version = self.nns_registry_client.get_latest_version(); + let blessed_measurements = self + .nns_registry_client + .get_blessed_guest_launch_measurements(registry_version) + .map_err(|e| anyhow!("Failed to get blessed measurements from registry: {e}"))?; // Verify the server's attestation report. This is to ensure that the key comes from a // trusted source. Without this check, an attacker could start with a malicious GuestOS, diff --git a/rs/ic_os/guest_upgrade/server/Cargo.toml b/rs/ic_os/guest_upgrade/server/Cargo.toml index fbde0128de7a..faaa240fb889 100644 --- a/rs/ic_os/guest_upgrade/server/Cargo.toml +++ b/rs/ic_os/guest_upgrade/server/Cargo.toml @@ -10,6 +10,7 @@ config_types = { path = "../../config/types" } der = { workspace = true } guest_upgrade_shared = { path = "../shared" } ic-interfaces-registry = { path = "../../../interfaces/registry" } +ic-registry-client-helpers = { path = "../../../registry/helpers" } sev_guest = { path = "../../sev/guest" } rcgen = { workspace = true } sev = { workspace = true } diff --git a/rs/ic_os/guest_upgrade/server/src/lib.rs b/rs/ic_os/guest_upgrade/server/src/lib.rs index e332d5179647..fbe97ac9b0a7 100644 --- a/rs/ic_os/guest_upgrade/server/src/lib.rs +++ b/rs/ic_os/guest_upgrade/server/src/lib.rs @@ -1,9 +1,9 @@ use crate::service::DiskEncryptionKeyExchangeServiceImpl; use attestation::attestation_package::SevRootCertificateVerification; -use attestation::registry::get_blessed_guest_launch_measurements_from_registry; use config_types::TrustedExecutionEnvironmentConfig; use guest_upgrade_shared::DEFAULT_SERVER_PORT; use ic_interfaces_registry::RegistryClient; +use ic_registry_client_helpers::blessed_replica_version::BlessedReplicaVersionRegistry; use server::DiskEncryptionKeyExchangeServer; use sev_guest::firmware::SevGuestFirmware; use std::sync::Arc; @@ -96,14 +96,15 @@ impl DiskEncryptionKeyExchangeServerAgent { )) })?; - let blessed_measurements = get_blessed_guest_launch_measurements_from_registry( - &*self.registry_client, - ) - .map_err(|err| { - DiskEncryptionKeyExchangeError::ServerStartError(format!( - "Failed to get blessed measurements: {err}" - )) - })?; + let registry_version = self.registry_client.get_latest_version(); + let blessed_measurements = self + .registry_client + .get_blessed_guest_launch_measurements(registry_version) + .map_err(|err| { + DiskEncryptionKeyExchangeError::ServerStartError(format!( + "Failed to get blessed measurements: {err}" + )) + })?; let upgrade_service = Arc::new(DiskEncryptionKeyExchangeServiceImpl::new( self.sev_firmware_factory.clone(), self.sev_root_certificate_verification, diff --git a/rs/ic_os/guest_upgrade/tests/BUILD.bazel b/rs/ic_os/guest_upgrade/tests/BUILD.bazel index 68defcf0486d..d50d81bc266a 100644 --- a/rs/ic_os/guest_upgrade/tests/BUILD.bazel +++ b/rs/ic_os/guest_upgrade/tests/BUILD.bazel @@ -12,12 +12,14 @@ rust_test( "//rs/ic_os/guest_upgrade/server", "//rs/ic_os/guest_upgrade/shared", "//rs/ic_os/sev/attestation", - "//rs/ic_os/sev/attestation/testing", "//rs/ic_os/sev/guest", "//rs/ic_os/sev/guest/testing", "//rs/ic_os/vsock/vsock_lib", "//rs/interfaces/registry", "//rs/protobuf", + "//rs/registry/fake", + "//rs/registry/proto_data_provider", + "//rs/test_utilities/registry", "@crate_index//:anyhow", "@crate_index//:futures", "@crate_index//:rustls", diff --git a/rs/ic_os/guest_upgrade/tests/Cargo.toml b/rs/ic_os/guest_upgrade/tests/Cargo.toml index 477867068c50..a178f7d985e3 100644 --- a/rs/ic_os/guest_upgrade/tests/Cargo.toml +++ b/rs/ic_os/guest_upgrade/tests/Cargo.toml @@ -7,17 +7,19 @@ edition = "2021" [dependencies] anyhow = { workspace = true } attestation = { path = "../../sev/attestation" } -attestation_testing = { path = "../../sev/attestation/testing" } +config_types = { path = "../../config/types" } futures = { workspace = true } guest_upgrade_client = { path = "../client" } guest_upgrade_server = { path = "../server" } guest_upgrade_shared = { path = "../shared" } -tokio = { workspace = true } -vsock_lib = { path = "../../vsock/vsock_lib" } -config_types = { path = "../../config/types" } +ic-interfaces-registry = { path = "../../../interfaces/registry" } ic-protobuf = { path = "../../../protobuf" } +ic-registry-client-fake = { path = "../../../registry/fake" } +ic-registry-proto-data-provider = { path = "../../../registry/proto_data_provider" } +ic-test-utilities-registry = { path = "../../../test_utilities/registry" } +rustls = { workspace = true } sev_guest = { path = "../../sev/guest" } sev_guest_testing = { path = "../../sev/guest/testing" } -ic-interfaces-registry = { path = "../../../interfaces/registry" } tempfile = { workspace = true } -rustls = { workspace = true } +tokio = { workspace = true } +vsock_lib = { path = "../../vsock/vsock_lib" } diff --git a/rs/ic_os/guest_upgrade/tests/src/lib.rs b/rs/ic_os/guest_upgrade/tests/src/lib.rs index 7d721841b413..0c050c591985 100644 --- a/rs/ic_os/guest_upgrade/tests/src/lib.rs +++ b/rs/ic_os/guest_upgrade/tests/src/lib.rs @@ -2,7 +2,6 @@ use anyhow::bail; use attestation::attestation_package::SevRootCertificateVerification; -use attestation_testing::registry::setup_mock_registry_client_with_blessed_versions; use config_types::{ GuestOSConfig, GuestOSUpgradeConfig, GuestVMType, ICOSSettings, TrustedExecutionEnvironmentConfig, @@ -15,6 +14,9 @@ use guest_upgrade_shared::{DEFAULT_SERVER_PORT, STORE_DEVICE}; use ic_protobuf::registry::replica_version::v1::{ GuestLaunchMeasurement, GuestLaunchMeasurements, ReplicaVersionRecord, }; +use ic_registry_client_fake::FakeRegistryClient; +use ic_registry_proto_data_provider::ProtoRegistryDataProvider; +use ic_test_utilities_registry::{add_blessed_replica_versions, add_replica_version_record}; use sev_guest::key_deriver::{Key, derive_key_from_sev_measurement}; use sev_guest_testing::{FakeAttestationReportSigner, MockSevGuestFirmwareBuilder}; use std::future::Future; @@ -98,29 +100,34 @@ impl DiskEncryptionKeyExchangeTestFixture { fn new(config: TestConfig) -> Self { let _ = rustls::crypto::ring::default_provider().install_default(); - let registry_client = Arc::new(setup_mock_registry_client_with_blessed_versions( - 1.into(), - &[REPLICA_VERSION], - &[( - REPLICA_VERSION, - ReplicaVersionRecord { - release_package_sha256_hex: "abc".to_string(), - guest_launch_measurements: Some(GuestLaunchMeasurements { - guest_launch_measurements: vec![ - GuestLaunchMeasurement { - measurement: DEFAULT_CLIENT_MEASUREMENT.into(), - metadata: None, - }, - GuestLaunchMeasurement { - measurement: DEFAULT_SERVER_MEASUREMENT.into(), - metadata: None, - }, - ], - }), - release_package_urls: vec![], - }, - )], - )); + let registry_data_provider = Arc::new(ProtoRegistryDataProvider::new()); + + add_blessed_replica_versions(®istry_data_provider, 1, &[REPLICA_VERSION]); + + add_replica_version_record( + ®istry_data_provider, + 1, + REPLICA_VERSION, + ReplicaVersionRecord { + release_package_sha256_hex: "abc".to_string(), + guest_launch_measurements: Some(GuestLaunchMeasurements { + guest_launch_measurements: vec![ + GuestLaunchMeasurement { + measurement: DEFAULT_CLIENT_MEASUREMENT.into(), + metadata: None, + }, + GuestLaunchMeasurement { + measurement: DEFAULT_SERVER_MEASUREMENT.into(), + metadata: None, + }, + ], + }), + release_package_urls: vec![], + }, + ); + + let registry_client = Arc::new(FakeRegistryClient::new(registry_data_provider)); + registry_client.update_to_latest_version(); let fake_attestation_report_signer = FakeAttestationReportSigner::default(); diff --git a/rs/ic_os/sev/attestation/BUILD.bazel b/rs/ic_os/sev/attestation/BUILD.bazel index 1ac17f528c90..d3e564dc3a14 100644 --- a/rs/ic_os/sev/attestation/BUILD.bazel +++ b/rs/ic_os/sev/attestation/BUILD.bazel @@ -13,15 +13,10 @@ rust_library( srcs = glob(["src/**/*.rs"]), crate_name = "attestation", deps = [ - "//rs/interfaces/registry", - "//rs/protobuf", - "//rs/registry/helpers", - "//rs/types/types", "//rs/utils", "@crate_index//:candid", "@crate_index//:der", "@crate_index//:hex", - "@crate_index//:itertools", "@crate_index//:prost", "@crate_index//:rand", "@crate_index//:serde", diff --git a/rs/ic_os/sev/attestation/Cargo.toml b/rs/ic_os/sev/attestation/Cargo.toml index 0db8c8c4501c..1ae3388714f0 100644 --- a/rs/ic_os/sev/attestation/Cargo.toml +++ b/rs/ic_os/sev/attestation/Cargo.toml @@ -6,7 +6,7 @@ edition.workspace = true candid = { workspace = true } der = { workspace = true, features = ["alloc", "derive", "std"] } hex = { workspace = true } -itertools = { workspace = true } +ic-utils = { path = "../../../utils" } prost = { workspace = true } rand = { workspace = true } serde = { workspace = true, features = ["derive"] } @@ -14,12 +14,6 @@ sev = { workspace = true } sha2 = { workspace = true } thiserror = { workspace = true } -ic-interfaces-registry = { path = "../../../interfaces/registry" } -ic-protobuf = { path = "../../../protobuf" } -ic-registry-client-helpers = { path = "../../../registry/helpers" } -ic-types = { path = "../../../types/types" } -ic-utils = { path = "../../../utils" } - [target.'cfg(target_arch = "wasm32")'.dependencies] getrandom = { version = "0.2", features = ["custom"] } diff --git a/rs/ic_os/sev/attestation/src/lib.rs b/rs/ic_os/sev/attestation/src/lib.rs index e6011522908c..c9a7e9d90d91 100644 --- a/rs/ic_os/sev/attestation/src/lib.rs +++ b/rs/ic_os/sev/attestation/src/lib.rs @@ -4,7 +4,6 @@ use std::fmt::{Debug, Display, Formatter}; pub mod attestation_package; pub mod custom_data; mod proto_gen; -pub mod registry; #[cfg(test)] mod e2e_tests; diff --git a/rs/ic_os/sev/attestation/src/registry.rs b/rs/ic_os/sev/attestation/src/registry.rs deleted file mode 100644 index 1d2706a8e9c2..000000000000 --- a/rs/ic_os/sev/attestation/src/registry.rs +++ /dev/null @@ -1,104 +0,0 @@ -use ic_interfaces_registry::RegistryClient; -use ic_registry_client_helpers::blessed_replica_version::BlessedReplicaVersionRegistry; -use ic_registry_client_helpers::subnet::SubnetRegistry; -use ic_types::ReplicaVersion; -use itertools::Itertools; -use std::str::FromStr; - -pub fn get_blessed_guest_launch_measurements_from_registry( - nns_registry_client: &dyn RegistryClient, -) -> Result>, String> { - let latest_registry_version = nns_registry_client.get_latest_version(); - let blessed_replica_versions = nns_registry_client - .get_blessed_replica_versions(latest_registry_version) - .map_err(|err| format!("Failed to get blessed replica versions: {err}"))? - .ok_or_else(|| "Blessed replica versions not found in registry".to_string())?; - - let measurements = blessed_replica_versions - .blessed_version_ids - .iter() - .filter_map(|version_id| ReplicaVersion::from_str(version_id).ok()) - .filter_map(|replica_version| { - nns_registry_client - .get_replica_version_record_from_version_id( - &replica_version, - latest_registry_version, - ) - .ok() - .flatten() - }) - .flat_map(|record| { - record - .guest_launch_measurements - .unwrap_or_default() - .guest_launch_measurements - }) - .map(|measurement| measurement.measurement) - .collect_vec(); - - Ok(measurements) -} - -#[cfg(test)] -mod tests { - use super::*; - use attestation_testing::registry::setup_mock_registry_client_with_blessed_versions; - use ic_protobuf::registry::replica_version::v1::{ - GuestLaunchMeasurement, GuestLaunchMeasurements, ReplicaVersionRecord, - }; - use ic_types::RegistryVersion; - - fn create_replica_record( - package_hash: &str, - measurements: &[impl AsRef<[u8]>], - ) -> ReplicaVersionRecord { - ReplicaVersionRecord { - release_package_sha256_hex: package_hash.to_string(), - release_package_urls: vec![], - guest_launch_measurements: Some(GuestLaunchMeasurements { - guest_launch_measurements: measurements - .iter() - .map(|m| GuestLaunchMeasurement { - measurement: m.as_ref().to_vec(), - metadata: None, - }) - .collect(), - }), - } - } - - #[test] - fn test_get_blessed_guest_launch_measurements_from_registry() { - let registry_version = RegistryVersion::from(42); - let blessed_versions = ["version1", "version2"]; - - let measurement1 = [1, 2, 3, 4, 5]; - let measurement2 = [6, 7, 8, 9, 10]; - let measurement3 = [11, 12, 13, 14, 15]; - let measurement4 = [16, 17, 18, 19, 20]; // From unblessed version - - let replica_versions_and_records = vec![ - ( - "version1", - create_replica_record("12345", &[measurement1, measurement2]), - ), - ("version2", create_replica_record("abcde", &[measurement3])), - ("version3", create_replica_record("99999", &[measurement4])), - ]; - - let mock_registry_client = setup_mock_registry_client_with_blessed_versions( - registry_version, - &blessed_versions, - &replica_versions_and_records, - ); - - assert_eq!( - get_blessed_guest_launch_measurements_from_registry(&mock_registry_client), - Ok(vec![ - measurement1.to_vec(), - measurement2.to_vec(), - measurement3.to_vec() - ]) - ); - } -} diff --git a/rs/ic_os/sev/attestation/testing/BUILD.bazel b/rs/ic_os/sev/attestation/testing/BUILD.bazel index e327b08b2b31..e80120a8b9f8 100644 --- a/rs/ic_os/sev/attestation/testing/BUILD.bazel +++ b/rs/ic_os/sev/attestation/testing/BUILD.bazel @@ -9,14 +9,8 @@ rust_library( crate_name = "attestation_testing", deps = [ "//rs/ic_os/sev/attestation", - "//rs/interfaces/registry/mocks", - "//rs/protobuf", - "//rs/registry/keys", - "//rs/types/types", "@crate_index//:der", - "@crate_index//:mockall", "@crate_index//:p384", - "@crate_index//:prost", "@crate_index//:rand", "@crate_index//:rsa", "@crate_index//:sev", diff --git a/rs/ic_os/sev/attestation/testing/Cargo.toml b/rs/ic_os/sev/attestation/testing/Cargo.toml index 7ad5ec45eea3..bf567fa4cccd 100644 --- a/rs/ic_os/sev/attestation/testing/Cargo.toml +++ b/rs/ic_os/sev/attestation/testing/Cargo.toml @@ -11,10 +11,3 @@ rsa = { workspace = true } sev = { workspace = true } sha2 = { workspace = true } x509-cert = { workspace = true } - -ic-interfaces-registry-mocks = { path = "../../../../interfaces/registry/mocks" } -ic-protobuf = { path = "../../../../protobuf" } -ic-registry-keys = { path = "../../../../registry/keys" } -ic-types = { path = "../../../../types/types" } -mockall = { workspace = true } -prost = { workspace = true } diff --git a/rs/ic_os/sev/attestation/testing/src/lib.rs b/rs/ic_os/sev/attestation/testing/src/lib.rs index 437cfc5f05f2..f07c875afc5c 100644 --- a/rs/ic_os/sev/attestation/testing/src/lib.rs +++ b/rs/ic_os/sev/attestation/testing/src/lib.rs @@ -1,3 +1,2 @@ pub mod attestation_package; pub mod attestation_report; -pub mod registry; diff --git a/rs/ic_os/sev/attestation/testing/src/registry.rs b/rs/ic_os/sev/attestation/testing/src/registry.rs deleted file mode 100644 index a62de0058150..000000000000 --- a/rs/ic_os/sev/attestation/testing/src/registry.rs +++ /dev/null @@ -1,45 +0,0 @@ -use ic_interfaces_registry_mocks::MockRegistryClient; -use ic_protobuf::registry::replica_version::v1::{BlessedReplicaVersions, ReplicaVersionRecord}; -use ic_registry_keys::{make_blessed_replica_versions_key, make_replica_version_key}; -use ic_types::RegistryVersion; -use mockall::predicate::eq; -use prost::Message; - -/// Create a mock registry client with the specified blessed replica versions and their -/// corresponding version records. -pub fn setup_mock_registry_client_with_blessed_versions( - registry_version: RegistryVersion, - blessed_version_ids: &[&str], - replica_versions_and_records: &[(&str, ReplicaVersionRecord)], -) -> MockRegistryClient { - let mut mock_client = MockRegistryClient::new(); - - mock_client - .expect_get_latest_version() - .return_const(registry_version); - - let blessed_versions = BlessedReplicaVersions { - blessed_version_ids: blessed_version_ids.iter().map(|x| x.to_string()).collect(), - }; - - mock_client - .expect_get_value() - .with( - eq(make_blessed_replica_versions_key()), - eq(registry_version), - ) - .returning(move |_, _| Ok(Some(blessed_versions.encode_to_vec()))); - - for (version_key, record) in replica_versions_and_records { - let encoded_record = record.encode_to_vec(); - mock_client - .expect_get_value() - .with( - eq(make_replica_version_key(version_key)), - eq(registry_version), - ) - .returning(move |_, _| Ok(Some(encoded_record.clone()))); - } - - mock_client -} diff --git a/rs/registry/helpers/src/blessed_replica_version.rs b/rs/registry/helpers/src/blessed_replica_version.rs index dfecd9667756..fdf920b95e5b 100644 --- a/rs/registry/helpers/src/blessed_replica_version.rs +++ b/rs/registry/helpers/src/blessed_replica_version.rs @@ -1,14 +1,27 @@ use crate::deserialize_registry_value; +use crate::subnet::SubnetRegistry; use ic_base_types::RegistryVersion; use ic_interfaces_registry::{RegistryClient, RegistryClientResult}; use ic_protobuf::registry::replica_version::v1::BlessedReplicaVersions; use ic_registry_keys::make_blessed_replica_versions_key; +use ic_types::ReplicaVersion; +use std::str::FromStr; pub trait BlessedReplicaVersionRegistry { fn get_blessed_replica_versions( &self, version: RegistryVersion, ) -> RegistryClientResult; + + /// Returns all guest launch measurements from all blessed replica versions. + /// + /// This method fetches the blessed replica versions from the registry at the given version, + /// then retrieves the replica version records for each blessed version, and finally collects + /// all guest launch measurements from those records. + fn get_blessed_guest_launch_measurements( + &self, + version: RegistryVersion, + ) -> Result>, String>; } impl BlessedReplicaVersionRegistry for T { @@ -18,4 +31,126 @@ impl BlessedReplicaVersionRegistry for T { ) -> RegistryClientResult { deserialize_registry_value(self.get_value(&make_blessed_replica_versions_key(), version)) } + + fn get_blessed_guest_launch_measurements( + &self, + version: RegistryVersion, + ) -> Result>, String> { + let blessed_replica_versions = self + .get_blessed_replica_versions(version) + .map_err(|err| format!("Failed to get blessed replica versions: {err}"))? + .ok_or_else(|| "Blessed replica versions not found in registry".to_string())?; + + let measurements = blessed_replica_versions + .blessed_version_ids + .iter() + .filter_map(|version_id| ReplicaVersion::from_str(version_id).ok()) + .filter_map(|replica_version| { + self.get_replica_version_record_from_version_id(&replica_version, version) + .ok() + .flatten() + }) + .flat_map(|record| { + record + .guest_launch_measurements + .unwrap_or_default() + .guest_launch_measurements + }) + .map(|measurement| measurement.measurement) + .collect(); + + Ok(measurements) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ic_protobuf::registry::replica_version::v1::{ + GuestLaunchMeasurement, GuestLaunchMeasurements, ReplicaVersionRecord, + }; + use ic_registry_client_fake::FakeRegistryClient; + use ic_registry_keys::make_replica_version_key; + use ic_registry_proto_data_provider::ProtoRegistryDataProvider; + use std::sync::Arc; + + fn create_replica_record( + package_hash: &str, + measurements: &[impl AsRef<[u8]>], + ) -> ReplicaVersionRecord { + ReplicaVersionRecord { + release_package_sha256_hex: package_hash.to_string(), + release_package_urls: vec![], + guest_launch_measurements: Some(GuestLaunchMeasurements { + guest_launch_measurements: measurements + .iter() + .map(|m| GuestLaunchMeasurement { + measurement: m.as_ref().to_vec(), + metadata: None, + }) + .collect(), + }), + } + } + + #[test] + fn test_get_blessed_guest_launch_measurements() { + let registry_version = RegistryVersion::from(42); + let blessed_versions = ["version1", "version2"]; + + let measurement1 = [1, 2, 3, 4, 5]; + let measurement2 = [6, 7, 8, 9, 10]; + let measurement3 = [11, 12, 13, 14, 15]; + let measurement4 = [16, 17, 18, 19, 20]; // From unblessed version + + let replica_versions_and_records = vec![ + ( + "version1", + create_replica_record("12345", &[measurement1, measurement2]), + ), + ("version2", create_replica_record("abcde", &[measurement3])), + ("version3", create_replica_record("99999", &[measurement4])), + ]; + + // Set up registry data provider + let data_provider = ProtoRegistryDataProvider::new(); + + // Add blessed replica versions + let blessed_versions_proto = BlessedReplicaVersions { + blessed_version_ids: blessed_versions.iter().map(|x| x.to_string()).collect(), + }; + data_provider + .add( + &make_blessed_replica_versions_key(), + registry_version, + Some(blessed_versions_proto), + ) + .expect("Failed to add blessed replica versions"); + + // Add replica version records + for (version_id, record) in &replica_versions_and_records { + data_provider + .add( + &make_replica_version_key(version_id), + registry_version, + Some(record.clone()), + ) + .expect("Failed to add replica version record"); + } + + // Create registry client and update to latest version + let registry_client = FakeRegistryClient::new(Arc::new(data_provider)); + registry_client.update_to_latest_version(); + + let result = registry_client.get_blessed_guest_launch_measurements(registry_version); + + assert_eq!( + result, + Ok(vec![ + measurement1.to_vec(), + measurement2.to_vec(), + measurement3.to_vec() + ]) + ); + } } diff --git a/rs/test_utilities/registry/src/lib.rs b/rs/test_utilities/registry/src/lib.rs index 91b59a3c31c4..491b6b59a69d 100644 --- a/rs/test_utilities/registry/src/lib.rs +++ b/rs/test_utilities/registry/src/lib.rs @@ -3,6 +3,7 @@ use ic_limits::INITIAL_NOTARY_DELAY; use ic_management_canister_types_private::VetKdKeyId; use ic_protobuf::registry::crypto::v1::AlgorithmId; use ic_protobuf::registry::crypto::v1::PublicKey as PublicKeyProto; +use ic_protobuf::registry::replica_version::v1::{BlessedReplicaVersions, ReplicaVersionRecord}; use ic_protobuf::registry::subnet::v1::ChainKeyInitialization; use ic_protobuf::registry::subnet::v1::chain_key_initialization::Initialization; use ic_protobuf::registry::subnet::v1::{ @@ -12,7 +13,8 @@ use ic_protobuf::registry::subnet::v1::{ use ic_protobuf::types::v1::master_public_key_id::KeyId; use ic_registry_client_fake::FakeRegistryClient; use ic_registry_keys::{ - make_catch_up_package_contents_key, make_crypto_threshold_signing_pubkey_key, + make_blessed_replica_versions_key, make_catch_up_package_contents_key, + make_crypto_threshold_signing_pubkey_key, make_replica_version_key, make_subnet_list_record_key, make_subnet_record_key, }; use ic_registry_local_store::{LocalStoreImpl, compact_delta_to_changelog}; @@ -80,6 +82,42 @@ pub fn setup_registry_non_final( (registry_data_provider, registry) } +/// Add blessed replica versions to the registry. +pub fn add_blessed_replica_versions( + registry_data_provider: &Arc, + version: u64, + blessed_version_ids: &[&str], +) { + let registry_version = RegistryVersion::from(version); + let blessed_versions = BlessedReplicaVersions { + blessed_version_ids: blessed_version_ids.iter().map(|x| x.to_string()).collect(), + }; + registry_data_provider + .add( + &make_blessed_replica_versions_key(), + registry_version, + Some(blessed_versions), + ) + .expect("Failed to add blessed replica versions."); +} + +/// Add a replica version record to the registry. +pub fn add_replica_version_record( + registry_data_provider: &Arc, + version: u64, + replica_version_id: &str, + record: ReplicaVersionRecord, +) { + let registry_version = RegistryVersion::from(version); + registry_data_provider + .add( + &make_replica_version_key(replica_version_id), + registry_version, + Some(record), + ) + .expect("Failed to add replica version record."); +} + pub fn insert_initial_dkg_transcript( version: u64, subnet_id: SubnetId,