Added load_canister_snapshot to the Root canister. This commit is analogous to the changes in PR 8159.

Author: daniel-wong-dfinity-org

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,
+    );
 }

rs/nns/handlers/root/impl/src/canister_management.rs

@@ -7,7 +7,7 @@ use ic_cdk::{
 };
 use ic_management_canister_types_private::{
     CanisterInstallMode::Install, CanisterSettingsArgsBuilder, CanisterSnapshotResponse,
-    CreateCanisterArgs, InstallCodeArgs, TakeCanisterSnapshotArgs,
+    CreateCanisterArgs, InstallCodeArgs, LoadCanisterSnapshotArgs, TakeCanisterSnapshotArgs,
 };
 use ic_nervous_system_clients::{
     canister_id_record::CanisterIdRecord,
@@ -24,9 +24,11 @@ use ic_nns_common::{
     types::CallCanisterRequest,
 };
 use ic_nns_handler_root_interface::{
-    ChangeCanisterControllersRequest, ChangeCanisterControllersResponse, TakeCanisterSnapshotError,
-    TakeCanisterSnapshotOk, TakeCanisterSnapshotRequest, TakeCanisterSnapshotResponse,
-    UpdateCanisterSettingsError, UpdateCanisterSettingsRequest, UpdateCanisterSettingsResponse,
+    ChangeCanisterControllersRequest, ChangeCanisterControllersResponse, LoadCanisterSnapshotError,
+    LoadCanisterSnapshotOk, LoadCanisterSnapshotRequest, LoadCanisterSnapshotResponse,
+    TakeCanisterSnapshotError, TakeCanisterSnapshotOk, TakeCanisterSnapshotRequest,
+    TakeCanisterSnapshotResponse, UpdateCanisterSettingsError, UpdateCanisterSettingsRequest,
+    UpdateCanisterSettingsResponse,
 };
 use ic_protobuf::{
     registry::nns::v1::{NnsCanisterRecord, NnsCanisterRecords},
@@ -321,3 +323,50 @@ fn convert_from_canister_snapshot_response_to_take_canister_snapshot_ok(
         total_size,
     }
 }
+
+pub async fn load_canister_snapshot(
+    load_canister_snapshot_request: LoadCanisterSnapshotRequest,
+    management_canister_client: &mut impl ManagementCanisterClient,
+) -> LoadCanisterSnapshotResponse {
+    let LoadCanisterSnapshotRequest {
+        canister_id,
+        snapshot_id,
+    } = load_canister_snapshot_request;
+
+    let snapshot_id = match SnapshotId::try_from(snapshot_id) {
+        Ok(ok) => ok,
+        Err(err) => {
+            return LoadCanisterSnapshotResponse::Err(LoadCanisterSnapshotError {
+                code: None,
+                description: format!("Invalid snapshot ID: {err}"),
+            });
+        }
+    };
+
+    let canister_id = match CanisterId::try_from(canister_id) {
+        Ok(ok) => ok,
+        Err(err) => {
+            return LoadCanisterSnapshotResponse::Err(LoadCanisterSnapshotError {
+                code: None,
+                description: format!("Invalid canister ID: {err}"),
+            });
+        }
+    };
+
+    let load_canister_snapshot_args = LoadCanisterSnapshotArgs::new(
+        canister_id,
+        snapshot_id,
+        management_canister_client.canister_version(),
+    );
+
+    match management_canister_client
+        .load_canister_snapshot(load_canister_snapshot_args)
+        .await
+    {
+        Ok(()) => LoadCanisterSnapshotResponse::Ok(LoadCanisterSnapshotOk {}),
+        Err((code, description)) => LoadCanisterSnapshotResponse::Err(LoadCanisterSnapshotError {
+            code: Some(code),
+            description,
+        }),
+    }
+}

rs/nns/handlers/root/interface/src/lib.rs

@@ -107,3 +107,24 @@ pub struct TakeCanisterSnapshotError {
     pub code: Option<i32>,
     pub description: String,
 }
+
+#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)]
+pub struct LoadCanisterSnapshotRequest {
+    pub canister_id: PrincipalId,
+    pub snapshot_id: Vec<u8>,
+}
+
+#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)]
+pub enum LoadCanisterSnapshotResponse {
+    Ok(LoadCanisterSnapshotOk),
+    Err(LoadCanisterSnapshotError),
+}
+
+#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)]
+pub struct LoadCanisterSnapshotOk {}
+
+#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)]
+pub struct LoadCanisterSnapshotError {
+    pub code: Option<i32>,
+    pub description: String,
+}

rs/nns/handlers/root/unreleased_changelog.md

@@ -9,7 +9,10 @@ on the process that this file is part of, see
 
 ## Added
 
-* Added `take_canister_snapshot` method. It is only callable by the Governance canister though.
+* Added `take_canister_snapshot`, and `load_canister_snapshot` methods. These
+  are only callable by the Governance canister though. What these do is
+  proxy/immediately forward to methods of the same name in the Management
+  (pseudo-)canister.
 
 ## Changed