feat: Allow recalling GuestOS versions

rs/boundary_node/ic_boundary/src/test_utils.rs

@@ -122,6 +122,7 @@ pub fn test_subnet_record() -> SubnetRecord {
         ssh_backup_access: vec![],
         chain_key_config: None,
         canister_cycles_cost_schedule: CanisterCyclesCostSchedule::Normal as i32,
+        recalled_replica_version_ids: vec![],
     }
 }
 

rs/crypto/temp_crypto/src/lib.rs

@@ -1117,6 +1117,7 @@ impl EcdsaSubnetConfig {
                     max_parallel_pre_signature_transcripts_in_creation: None,
                 }),
                 canister_cycles_cost_schedule: CanisterCyclesCostSchedule::Normal as i32,
+                recalled_replica_version_ids: vec![],
             },
         }
     }

rs/nns/governance/api/src/types.rs

@@ -2979,7 +2979,7 @@ pub struct ListProposalInfoResponse {
 /// Paging is available if the result set is larger than `MAX_LIST_NEURONS_RESULTS`,
 /// which is currently 500 neurons.  If you are unsure of the number of results in a set,
 /// you can use the `total_pages_available` field in the response to determine how many
-/// additional pages need to be queried.  It will be based on your `page_size` parameter.  
+/// additional pages need to be queried.  It will be based on your `page_size` parameter.
 /// When paging through results, it is good to keep in mind that newly inserted neurons
 /// could be missed if they are inserted between calls to pages, and this could result in missing
 /// a neuron in the combined responses.

rs/nns/integration_tests/src/subnet_handler.rs

@@ -52,6 +52,7 @@ fn test_submit_and_accept_update_subnet_proposal() {
                 ssh_backup_access: vec![],
                 chain_key_config: None,
                 canister_cycles_cost_schedule: CanisterCyclesCostSchedule::Normal as i32,
+                recalled_replica_version_ids: vec![],
             };
 
             let key = make_subnet_record_key(subnet_id);
@@ -170,6 +171,7 @@ fn test_submit_and_accept_update_subnet_proposal() {
                     ssh_backup_access: vec!["pub_key_1".to_string()],
                     chain_key_config: None,
                     canister_cycles_cost_schedule: CanisterCyclesCostSchedule::Normal as i32,
+                    recalled_replica_version_ids: vec![],
                 }
             );
             Ok(())

rs/orchestrator/Cargo.toml

@@ -71,13 +71,16 @@ qrcode = { workspace = true }
 [dev-dependencies]
 async-stream = { workspace = true }
 hyper-util = { workspace = true }
+ic-certification-test-utils = { path = "../certification/test-utils" }
 ic-crypto-temp-crypto = { path = "../crypto/temp_crypto" }
 ic-crypto-test-utils-canister-threshold-sigs = { path = "../crypto/test_utils/canister_threshold_sigs" }
 ic-crypto-test-utils-crypto-returning-ok = { path = "../crypto/test_utils/crypto_returning_ok" }
 ic-crypto-test-utils-ni-dkg = { path = "../crypto/test_utils/ni-dkg" }
 ic-crypto-test-utils-reproducible-rng = { path = "../crypto/test_utils/reproducible_rng" }
 ic-crypto-test-utils-tls = { path = "../crypto/test_utils/tls" }
 ic-crypto-tls-interfaces-mocks = { path = "../crypto/tls_interfaces/mocks" }
+ic-crypto-tree-hash = { path = "../crypto/tree_hash" }
+ic-registry-client = { path = "../registry/client" }
 ic-registry-client-fake = { path = "../registry/fake" }
 ic-registry-subnet-type = { path = "../registry/subnet_type" }
 ic-registry-proto-data-provider = { path = "../registry/proto_data_provider" }

rs/orchestrator/registry_replicator/Cargo.toml

@@ -7,6 +7,7 @@ description.workspace = true
 documentation.workspace = true
 
 [dependencies]
+async-trait = { workspace = true }
 clap = { workspace = true }
 ic-config = { path = "../../config" }
 ic-crypto-utils-threshold-sig-der = { path = "../../crypto/utils/threshold_sig_der" }

rs/orchestrator/registry_replicator/src/internal_state.rs

@@ -577,6 +577,7 @@ mod test {
             ssh_backup_access: vec![],
             chain_key_config: None,
             canister_cycles_cost_schedule: 0,
+            recalled_replica_version_ids: vec![],
         }
     }
 

rs/orchestrator/registry_replicator/src/lib.rs

@@ -63,6 +63,24 @@ use url::Url;
 pub mod args;
 mod internal_state;
 pub mod metrics;
+pub mod mock;
+
+/// Trait for registry replication functionality.
+/// This allows for mocking the registry replicator in tests.
+#[async_trait::async_trait]
+pub trait RegistryReplicatorTrait: Send + Sync {
+    /// Polls the registry once, fetching and applying updates.
+    async fn poll(&self) -> Result<(), String>;
+
+    /// Returns the registry client used by this replicator.
+    fn get_registry_client(&self) -> Arc<dyn RegistryClient>;
+
+    /// Returns the local store used by this replicator.
+    fn get_local_store(&self) -> Arc<dyn LocalStore>;
+
+    /// Stops polling and sets the local registry data to what is contained in the provided local store.
+    async fn stop_polling_and_set_local_registry_data(&self, new_local_store: &dyn LocalStore);
+}
 
 trait PollableRegistryClient: RegistryClient {
     /// Polls the registry once, updating its cache by polling the latest local store changes.
@@ -486,6 +504,25 @@ impl Drop for RegistryReplicator {
     }
 }
 
+#[async_trait::async_trait]
+impl RegistryReplicatorTrait for RegistryReplicator {
+    async fn poll(&self) -> Result<(), String> {
+        self.poll().await
+    }
+
+    fn get_registry_client(&self) -> Arc<dyn RegistryClient> {
+        self.get_registry_client()
+    }
+
+    fn get_local_store(&self) -> Arc<dyn LocalStore> {
+        self.get_local_store()
+    }
+
+    async fn stop_polling_and_set_local_registry_data(&self, new_local_store: &dyn LocalStore) {
+        self.stop_polling_and_set_local_registry_data(new_local_store).await
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

rs/orchestrator/registry_replicator/src/mock.rs

@@ -0,0 +1,96 @@
+//! Mock implementation of RegistryReplicatorTrait for testing.
+
+use crate::RegistryReplicatorTrait;
+use ic_interfaces_registry::{RegistryClient, RegistryDataProvider};
+use ic_registry_client::client::RegistryClientImpl;
+use ic_registry_local_store::{KeyMutation, LocalStore, LocalStoreImpl, LocalStoreWriter};
+use ic_types::RegistryVersion;
+use std::sync::Arc;
+
+/// Mock registry replicator that syncs data from a "remote" registry data provider
+/// to a local store and registry client.
+pub struct MockRegistryReplicator {
+    /// The "remote" registry data provider that simulates the NNS registry
+    remote_data_provider: Arc<dyn RegistryDataProvider>,
+    /// The local store that gets updated during replication
+    local_store: Arc<LocalStoreImpl>,
+    /// The registry client that reads from the local store
+    registry_client: Arc<RegistryClientImpl>,
+}
+
+impl MockRegistryReplicator {
+    /// Creates a new mock registry replicator.
+    ///
+    /// # Arguments
+    /// * `remote_data_provider` - The "remote" registry that will be replicated from
+    /// * `local_store` - The local store that will be updated during replication
+    /// * `registry_client` - The registry client that reads from the local store
+    pub fn new(
+        remote_data_provider: Arc<dyn RegistryDataProvider>,
+        local_store: Arc<LocalStoreImpl>,
+        registry_client: Arc<RegistryClientImpl>,
+    ) -> Self {
+        Self {
+            remote_data_provider,
+            local_store,
+            registry_client,
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl RegistryReplicatorTrait for MockRegistryReplicator {
+    /// Simulates polling the remote registry by fetching updates from the remote
+    /// data provider and applying them to the local store.
+    async fn poll(&self) -> Result<(), String> {
+        // Get the current latest version in the local store
+        let local_latest_version = self.registry_client.get_latest_version();
+
+        // Fetch all updates from the remote data provider since the local version
+        let updates = self
+            .remote_data_provider
+            .get_updates_since(local_latest_version)
+            .map_err(|e| format!("Failed to get updates from remote: {:?}", e))?;
+
+        // Group updates by version
+        let mut version_mutations: std::collections::BTreeMap<RegistryVersion, Vec<KeyMutation>> =
+            std::collections::BTreeMap::new();
+
+        for record in updates {
+            version_mutations
+                .entry(record.version)
+                .or_insert_with(Vec::new)
+                .push(KeyMutation {
+                    key: record.key,
+                    value: record.value,
+                });
+        }
+
+        // Apply updates to the local store
+        for (version, mutations) in version_mutations {
+            self.local_store
+                .store(version, mutations)
+                .map_err(|e| format!("Failed to store updates: {:?}", e))?;
+        }
+
+        // Update the registry client to see the new data
+        self.registry_client
+            .poll_once()
+            .map_err(|e| format!("Failed to update registry client: {:?}", e))?;
+
+        Ok(())
+    }
+
+    fn get_registry_client(&self) -> Arc<dyn RegistryClient> {
+        self.registry_client.clone()
+    }
+
+    fn get_local_store(&self) -> Arc<dyn LocalStore> {
+        self.local_store.clone()
+    }
+
+    async fn stop_polling_and_set_local_registry_data(&self, _new_local_store: &dyn LocalStore) {
+        // Mock implementation does nothing - not needed for testing
+    }
+}
+

rs/orchestrator/src/catch_up_package_provider.rs

@@ -501,7 +501,7 @@ fn get_cup_proto_height(cup: &pb::CatchUpPackage) -> Option<Height> {
 }
 
 #[cfg(test)]
-mod tests {
+pub(crate) mod tests {
     use super::*;
     use crate::{
         catch_up_package_provider::CatchUpPackageProvider, registry_helper::RegistryHelper,
@@ -690,7 +690,7 @@ mod tests {
         ))
     }
 
-    fn mock_tls_config() -> MockTlsConfig {
+    pub(crate) fn mock_tls_config() -> MockTlsConfig {
         #[derive(Debug)]
         struct NoVerify;
         impl ServerCertVerifier for NoVerify {