feat(governance): LoadCanisterSnapshot proposal type.

rs/nervous_system/clients/src/lib.rs

@@ -3,6 +3,7 @@ pub mod canister_metadata;
 pub mod canister_status;
 pub mod delete_canister;
 pub mod ledger_client;
+pub mod load_canister_snapshot;
 pub mod management_canister_client;
 pub mod stop_canister;
 pub mod take_canister_snapshot;

rs/nervous_system/clients/src/load_canister_snapshot.rs

@@ -0,0 +1,9 @@
+use ic_management_canister_types_private::{IC_00, LoadCanisterSnapshotArgs};
+use ic_nervous_system_runtime::Runtime;
+
+pub async fn load_canister_snapshot<Rt>(args: LoadCanisterSnapshotArgs) -> Result<(), (i32, String)>
+where
+    Rt: Runtime,
+{
+    Rt::call_with_cleanup(IC_00, "load_canister_snapshot", (args,)).await
+}

rs/nervous_system/clients/src/management_canister_client.rs

@@ -3,6 +3,7 @@ use crate::{
     canister_metadata::canister_metadata,
     canister_status::{CanisterStatusResultFromManagementCanister, canister_status},
     delete_canister::delete_canister,
+    load_canister_snapshot::load_canister_snapshot,
     stop_canister::stop_canister,
     take_canister_snapshot::take_canister_snapshot,
     update_settings::{UpdateSettings, update_settings},
@@ -12,7 +13,7 @@ use candid::Encode;
 use ic_base_types::PrincipalId;
 use ic_error_types::RejectCode;
 use ic_management_canister_types_private::{
-    CanisterSnapshotResponse, IC_00, TakeCanisterSnapshotArgs,
+    CanisterSnapshotResponse, IC_00, LoadCanisterSnapshotArgs, TakeCanisterSnapshotArgs,
 };
 use ic_nervous_system_proxied_canister_calls_tracker::ProxiedCanisterCallsTracker;
 use ic_nervous_system_runtime::Runtime;
@@ -62,6 +63,11 @@ pub trait ManagementCanisterClient {
         &self,
         args: TakeCanisterSnapshotArgs,
     ) -> Result<CanisterSnapshotResponse, (i32, String)>;
+
+    async fn load_canister_snapshot(
+        &self,
+        args: LoadCanisterSnapshotArgs,
+    ) -> Result<(), (i32, String)>;
 }
 
 /// An example implementation of the ManagementCanisterClient trait.
@@ -197,6 +203,24 @@ impl<Rt: Runtime + Sync> ManagementCanisterClient for ManagementCanisterClientIm
 
         take_canister_snapshot::<Rt>(args).await
     }
+
+    async fn load_canister_snapshot(
+        &self,
+        args: LoadCanisterSnapshotArgs,
+    ) -> Result<(), (i32, String)> {
+        let _tracker = self.proxied_canister_calls_tracker.map(|tracker| {
+            let encoded_args = Encode!(&args).unwrap_or_default();
+            ProxiedCanisterCallsTracker::start_tracking(
+                tracker,
+                dfn_core::api::caller(),
+                IC_00,
+                "load_canister_snapshot",
+                &encoded_args,
+            )
+        });
+
+        load_canister_snapshot::<Rt>(args).await
+    }
 }
 
 /// A ManagementCanisterClient that wraps another ManagementCanisterClient.
@@ -322,6 +346,14 @@ where
         let _loan = self.try_borrow_slot()?;
         self.inner.take_canister_snapshot(args).await
     }
+
+    async fn load_canister_snapshot(
+        &self,
+        args: LoadCanisterSnapshotArgs,
+    ) -> Result<(), (i32, String)> {
+        let _loan = self.try_borrow_slot()?;
+        self.inner.load_canister_snapshot(args).await
+    }
 }
 
 /// Increments available_slot_count by used_slot_count when dropped.
@@ -376,6 +408,7 @@ pub enum MockManagementCanisterClientCall {
     StopCanister(CanisterIdRecord),
     DeleteCanister(CanisterIdRecord),
     TakeCanisterSnapshot(TakeCanisterSnapshotArgs),
+    LoadCanisterSnapshot(LoadCanisterSnapshotArgs),
 }
 
 #[derive(Clone, Eq, PartialEq, Debug)]
@@ -387,6 +420,7 @@ pub enum MockManagementCanisterClientReply {
     StopCanister(Result<(), (i32, String)>),
     DeleteCanister(Result<(), (i32, String)>),
     TakeCanisterSnapshot(Result<CanisterSnapshotResponse, (i32, String)>),
+    LoadCanisterSnapshot(Result<(), (i32, String)>),
 }
 
 #[async_trait]
@@ -555,6 +589,32 @@ impl ManagementCanisterClient for MockManagementCanisterClient {
             ),
         }
     }
+
+    async fn load_canister_snapshot(
+        &self,
+        args: LoadCanisterSnapshotArgs,
+    ) -> Result<(), (i32, String)> {
+        self.calls
+            .lock()
+            .unwrap()
+            .push_back(MockManagementCanisterClientCall::LoadCanisterSnapshot(args));
+
+        let reply = self
+            .replies
+            .lock()
+            .unwrap()
+            .pop_front()
+            .expect("Expected a MockManagementCanisterClientCall to be on the queue.");
+
+        match reply {
+            MockManagementCanisterClientReply::LoadCanisterSnapshot(result) => result,
+            err => panic!(
+                "Expected MockManagementCanisterClientReply::LoadCanisterSnapshot to be at \
+                the front of the queue. Had {:?}",
+                err
+            ),
+        }
+    }
 }
 
 impl Drop for MockManagementCanisterClient {

rs/nervous_system/clients/src/management_canister_client/tests.rs

@@ -5,6 +5,9 @@ use crate::canister_status::{
 };
 use candid::Nat;
 use ic_base_types::{CanisterId, PrincipalId};
+use ic_management_canister_types_private::{
+    CanisterSnapshotResponse, LoadCanisterSnapshotArgs, TakeCanisterSnapshotArgs,
+};
 use rand::{Rng, thread_rng};
 use std::time::Duration;
 
@@ -91,6 +94,13 @@ async fn test_limit_outstanding_calls() {
             ) -> Result<CanisterSnapshotResponse, (i32, String)> {
                 unimplemented!();
             }
+
+            async fn load_canister_snapshot(
+                &self,
+                _args: LoadCanisterSnapshotArgs,
+            ) -> Result<(), (i32, String)> {
+                unimplemented!();
+            }
         }
 
         impl Drop for MockManagementCanisterClient {

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

@@ -4206,6 +4206,8 @@ pub enum NnsFunction {
     /// controlled by the NNS Root canister. This restriction could be relaxed
     /// later. See nns/.../root.did for the payload type.
     TakeCanisterSnapshot = 56,
+    /// Loads a canister snapshot.
+    LoadCanisterSnapshot = 57,
 }
 impl NnsFunction {
     /// String value of the enum field names used in the ProtoBuf definition.
@@ -4290,6 +4292,7 @@ impl NnsFunction {
             NnsFunction::UnpauseCanisterMigrations => "NNS_FUNCTION_UNPAUSE_CANISTER_MIGRATIONS",
             NnsFunction::SetSubnetOperationalLevel => "NNS_FUNCTION_SET_SUBNET_OPERATIONAL_LEVEL",
             NnsFunction::TakeCanisterSnapshot => "NNS_FUNCTION_TAKE_CANISTER_SNAPSHOT",
+            NnsFunction::LoadCanisterSnapshot => "NNS_FUNCTION_LOAD_CANISTER_SNAPSHOT",
         }
     }
     /// Creates an enum from field names used in the ProtoBuf definition.
@@ -4371,6 +4374,7 @@ impl NnsFunction {
             "NNS_FUNCTION_UNPAUSE_CANISTER_MIGRATIONS" => Some(Self::UnpauseCanisterMigrations),
             "NNS_FUNCTION_SET_SUBNET_OPERATIONAL_LEVEL" => Some(Self::SetSubnetOperationalLevel),
             "NNS_FUNCTION_TAKE_CANISTER_SNAPSHOT" => Some(Self::TakeCanisterSnapshot),
+            "NNS_FUNCTION_LOAD_CANISTER_SNAPSHOT" => Some(Self::LoadCanisterSnapshot),
             _ => None,
         }
     }

rs/nns/governance/proto/ic_nns_governance/pb/v1/governance.proto

@@ -498,6 +498,9 @@ enum NnsFunction {
 
   // Take a canister snapshot.
   NNS_FUNCTION_TAKE_CANISTER_SNAPSHOT = 56;
+
+  // Load a canister snapshot.
+  NNS_FUNCTION_LOAD_CANISTER_SNAPSHOT = 57;
 }
 
 // Payload of a proposal that calls a function on another NNS

rs/nns/governance/src/canister_state.rs

@@ -395,7 +395,8 @@ fn get_effective_payload(
         | ValidNnsFunction::PauseCanisterMigrations
         | ValidNnsFunction::UnpauseCanisterMigrations
         | ValidNnsFunction::SetSubnetOperationalLevel
-        | ValidNnsFunction::TakeCanisterSnapshot => Ok(payload.clone()),
+        | ValidNnsFunction::TakeCanisterSnapshot
+        | ValidNnsFunction::LoadCanisterSnapshot => Ok(payload.clone()),
     }
 }
 

rs/nns/governance/src/gen/ic_nns_governance.pb.v1.rs

@@ -4822,6 +4822,8 @@ pub enum NnsFunction {
     SetSubnetOperationalLevel = 55,
     /// Take a canister snapshot.
     TakeCanisterSnapshot = 56,
+    /// Load a canister snapshot.
+    LoadCanisterSnapshot = 57,
 }
 impl NnsFunction {
     /// String value of the enum field names used in the ProtoBuf definition.
@@ -4896,6 +4898,7 @@ impl NnsFunction {
             Self::UnpauseCanisterMigrations => "NNS_FUNCTION_UNPAUSE_CANISTER_MIGRATIONS",
             Self::SetSubnetOperationalLevel => "NNS_FUNCTION_SET_SUBNET_OPERATIONAL_LEVEL",
             Self::TakeCanisterSnapshot => "NNS_FUNCTION_TAKE_CANISTER_SNAPSHOT",
+            Self::LoadCanisterSnapshot => "NNS_FUNCTION_LOAD_CANISTER_SNAPSHOT",
         }
     }
     /// Creates an enum from field names used in the ProtoBuf definition.
@@ -4977,6 +4980,7 @@ impl NnsFunction {
             "NNS_FUNCTION_UNPAUSE_CANISTER_MIGRATIONS" => Some(Self::UnpauseCanisterMigrations),
             "NNS_FUNCTION_SET_SUBNET_OPERATIONAL_LEVEL" => Some(Self::SetSubnetOperationalLevel),
             "NNS_FUNCTION_TAKE_CANISTER_SNAPSHOT" => Some(Self::TakeCanisterSnapshot),
+            "NNS_FUNCTION_LOAD_CANISTER_SNAPSHOT" => Some(Self::LoadCanisterSnapshot),
             _ => None,
         }
     }

rs/nns/governance/src/governance.rs

@@ -4815,6 +4815,13 @@ impl Governance {
                     ));
                 }
             }
+            ValidNnsFunction::LoadCanisterSnapshot => {
+                if !are_canister_snapshot_proposals_enabled() {
+                    return Err(invalid_proposal_error(
+                        "LoadCanisterSnapshot proposals are not yet enabled.".to_string(),
+                    ));
+                }
+            }
             _ => {}
         };
 

rs/nns/governance/src/pb/conversions/mod.rs

@@ -3988,6 +3988,7 @@ impl From<pb::NnsFunction> for pb_api::NnsFunction {
                 pb_api::NnsFunction::SetSubnetOperationalLevel
             }
             pb::NnsFunction::TakeCanisterSnapshot => pb_api::NnsFunction::TakeCanisterSnapshot,
+            pb::NnsFunction::LoadCanisterSnapshot => pb_api::NnsFunction::LoadCanisterSnapshot,
         }
     }
 }
@@ -4098,6 +4099,7 @@ impl From<pb_api::NnsFunction> for pb::NnsFunction {
                 pb::NnsFunction::SetSubnetOperationalLevel
             }
             pb_api::NnsFunction::TakeCanisterSnapshot => pb::NnsFunction::TakeCanisterSnapshot,
+            pb_api::NnsFunction::LoadCanisterSnapshot => pb::NnsFunction::LoadCanisterSnapshot,
         }
     }
 }

rs/nns/governance/src/proposals/execute_nns_function.rs

@@ -154,6 +154,7 @@ pub enum ValidNnsFunction {
     UnpauseCanisterMigrations,
     SetSubnetOperationalLevel,
     TakeCanisterSnapshot,
+    LoadCanisterSnapshot,
 }
 
 impl ValidNnsFunction {
@@ -283,6 +284,7 @@ impl ValidNnsFunction {
                 (REGISTRY_CANISTER_ID, "set_subnet_operational_level")
             }
             ValidNnsFunction::TakeCanisterSnapshot => (ROOT_CANISTER_ID, "take_canister_snapshot"),
+            ValidNnsFunction::LoadCanisterSnapshot => (ROOT_CANISTER_ID, "load_canister_snapshot"),
         }
     }
 
@@ -338,13 +340,13 @@ impl ValidNnsFunction {
             | ValidNnsFunction::HardResetNnsRootToVersion
             | ValidNnsFunction::BitcoinSetConfig
             | ValidNnsFunction::PauseCanisterMigrations
-            | ValidNnsFunction::UnpauseCanisterMigrations => Topic::ProtocolCanisterManagement,
+            | ValidNnsFunction::UnpauseCanisterMigrations
+            | ValidNnsFunction::TakeCanisterSnapshot
+            | ValidNnsFunction::LoadCanisterSnapshot => Topic::ProtocolCanisterManagement,
 
             ValidNnsFunction::AddSnsWasm | ValidNnsFunction::InsertSnsWasmUpgradePathEntries => {
                 Topic::ServiceNervousSystemManagement
             }
-
-            ValidNnsFunction::TakeCanisterSnapshot => Topic::ProtocolCanisterManagement,
         }
     }
 
@@ -402,6 +404,7 @@ impl ValidNnsFunction {
             ValidNnsFunction::UnpauseCanisterMigrations => "Unpause Canister Migrations",
             ValidNnsFunction::SetSubnetOperationalLevel => "Set Subnet Operational Level",
             ValidNnsFunction::TakeCanisterSnapshot => "Take Canister Snapshot",
+            ValidNnsFunction::LoadCanisterSnapshot => "Load Canister Snapshot",
         }
     }
 
@@ -633,6 +636,13 @@ impl ValidNnsFunction {
                  For an introduction to canister snapshots in general, see \
                  https://docs.internetcomputer.org/building-apps/canister-management/snapshots ."
             }
+            ValidNnsFunction::LoadCanisterSnapshot => {
+                "A proposal to load a canister snapshot into a canister controlled the NNS. \
+                 In other words, to restore the canister to an earlier recorded state, \
+                 which including the code and memory (including stable memory). \
+                 For an introduction to canister snapshots in general, see \
+                 https://docs.internetcomputer.org/building-apps/canister-management/snapshots ."
+            }
         }
     }
 }
@@ -722,6 +732,7 @@ impl TryFrom<NnsFunction> for ValidNnsFunction {
                 Ok(ValidNnsFunction::SetSubnetOperationalLevel)
             }
             NnsFunction::TakeCanisterSnapshot => Ok(ValidNnsFunction::TakeCanisterSnapshot),
+            NnsFunction::LoadCanisterSnapshot => Ok(ValidNnsFunction::LoadCanisterSnapshot),
 
             // Obsolete functions - based on check_obsolete
             NnsFunction::BlessReplicaVersion | NnsFunction::RetireReplicaVersion => {

rs/nns/handlers/root/impl/canister/canister.rs

@@ -30,8 +30,8 @@ use ic_nns_handler_root::{
 };
 use ic_nns_handler_root_interface::{
     ChangeCanisterControllersRequest, ChangeCanisterControllersResponse,
-    TakeCanisterSnapshotRequest, TakeCanisterSnapshotResponse, UpdateCanisterSettingsRequest,
-    UpdateCanisterSettingsResponse,
+    LoadCanisterSnapshotRequest, LoadCanisterSnapshotResponse, TakeCanisterSnapshotRequest,
+    TakeCanisterSnapshotResponse, UpdateCanisterSettingsRequest, UpdateCanisterSettingsResponse,
 };
 use std::cell::RefCell;
 
@@ -242,6 +242,20 @@ async fn take_canister_snapshot(
     .await
 }
 
+/// Loads a snapshot of a canister controlled by NNS Root. Only callable by NNS
+/// Governance.
+#[update]
+async fn load_canister_snapshot(
+    load_canister_snapshot_request: LoadCanisterSnapshotRequest,
+) -> LoadCanisterSnapshotResponse {
+    check_caller_is_governance();
+    canister_management::load_canister_snapshot(
+        load_canister_snapshot_request,
+        &mut new_management_canister_client(),
+    )
+    .await
+}
+
 /// Resources to serve for a given http_request
 /// Serve an HttpRequest made to this canister
 #[query(

rs/nns/handlers/root/impl/canister/root.did

@@ -172,6 +172,26 @@ type TakeCanisterSnapshotError = record {
   description : text;
 };
 
+type LoadCanisterSnapshotRequest = record {
+  canister_id : principal;
+  snapshot_id : blob;
+};
+
+type LoadCanisterSnapshotResponse = variant {
+  Ok : LoadCanisterSnapshotOk;
+  Err : LoadCanisterSnapshotError;
+};
+
+type LoadCanisterSnapshotOk = record {
+  // This intentionally left empty. It is nevertheless here for
+  // a) consistency, and b) in case of future expansion.
+};
+
+type LoadCanisterSnapshotError = record {
+  code : opt int32;
+  description : text;
+};
+
 service : () -> {
   canister_status : (CanisterIdRecord) -> (CanisterStatusResult);
   get_build_metadata : () -> (text) query;
@@ -190,4 +210,7 @@ service : () -> {
   take_canister_snapshot : (TakeCanisterSnapshotArgs) -> (
       TakeCanisterSnapshotResponse,
     );
+  load_canister_snapshot : (LoadCanisterSnapshotRequest) -> (
+      LoadCanisterSnapshotResponse,
+    );
 }