refactor: Move registry functionality out of the attestation crate

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",
 ]
 

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" }

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,

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 }

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,

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",

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" }

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(&registry_data_provider, 1, &[REPLICA_VERSION]);
+
+        add_replica_version_record(
+            &registry_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();
 

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",

rs/ic_os/sev/attestation/Cargo.toml

@@ -6,20 +6,14 @@ 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"] }
 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"] }
 

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;