feat(nns-root): Added take_canister_snapshot to Root canister. (#8159)

Only callable by Governance.

Part of [NNS1-4295].

[NNS1-4295]: https://dfinity.atlassian.net/browse/NNS1-4295

[👈 Previous PR][prev]

[prev]: #8158

Author: daniel-wong-dfinity-org

Cargo.lock

@@ -30,7 +30,8 @@ use ic_nns_handler_root::{
 };
 use ic_nns_handler_root_interface::{
     ChangeCanisterControllersRequest, ChangeCanisterControllersResponse,
-    UpdateCanisterSettingsRequest, UpdateCanisterSettingsResponse,
+    TakeCanisterSnapshotRequest, TakeCanisterSnapshotResponse, UpdateCanisterSettingsRequest,
+    UpdateCanisterSettingsResponse,
 };
 use std::cell::RefCell;
 
@@ -227,6 +228,20 @@ async fn update_canister_settings(
     .await
 }
 
+/// Takes a snapshot of a canister controlled by NNS Root. Only callable by NNS
+/// Governance.
+#[update]
+async fn take_canister_snapshot(
+    take_canister_snapshot_request: TakeCanisterSnapshotRequest,
+) -> TakeCanisterSnapshotResponse {
+    check_caller_is_governance();
+    canister_management::take_canister_snapshot(
+        take_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/canister.rs

@@ -151,6 +151,27 @@ type CallCanisterRequest = record {
   payload : blob;
 };
 
+type TakeCanisterSnapshotArgs = record {
+  canister_id : principal;
+  replace_snapshot : opt blob;
+};
+
+type TakeCanisterSnapshotResponse = variant {
+  Ok : TakeCanisterSnapshotOk;
+  Err : TakeCanisterSnapshotError;
+};
+
+type TakeCanisterSnapshotOk = record {
+  id : blob;
+  taken_at_timestamp : nat64;
+  total_size : nat64;
+};
+
+type TakeCanisterSnapshotError = record {
+  code : opt int32;
+  description : text;
+};
+
 service : () -> {
   canister_status : (CanisterIdRecord) -> (CanisterStatusResult);
   get_build_metadata : () -> (text) query;
@@ -166,4 +187,7 @@ service : () -> {
       UpdateCanisterSettingsResponse,
     );
   call_canister : (CallCanisterRequest) -> ();
+  take_canister_snapshot : (TakeCanisterSnapshotArgs) -> (
+      TakeCanisterSnapshotResponse,
+    );
 }

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

@@ -1,12 +1,13 @@
 #![allow(deprecated)]
 use crate::PROXIED_CANISTER_CALLS_TRACKER;
-use ic_base_types::{CanisterId, PrincipalId};
+use ic_base_types::{CanisterId, PrincipalId, SnapshotId};
 use ic_cdk::{
     api::call::{RejectionCode, call_with_payment},
     call, caller, print,
 };
 use ic_management_canister_types_private::{
-    CanisterInstallMode::Install, CanisterSettingsArgsBuilder, CreateCanisterArgs, InstallCodeArgs,
+    CanisterInstallMode::Install, CanisterSettingsArgsBuilder, CanisterSnapshotResponse,
+    CreateCanisterArgs, InstallCodeArgs, TakeCanisterSnapshotArgs,
 };
 use ic_nervous_system_clients::{
     canister_id_record::CanisterIdRecord,
@@ -23,7 +24,8 @@ use ic_nns_common::{
     types::CallCanisterRequest,
 };
 use ic_nns_handler_root_interface::{
-    ChangeCanisterControllersRequest, ChangeCanisterControllersResponse,
+    ChangeCanisterControllersRequest, ChangeCanisterControllersResponse, TakeCanisterSnapshotError,
+    TakeCanisterSnapshotOk, TakeCanisterSnapshotRequest, TakeCanisterSnapshotResponse,
     UpdateCanisterSettingsError, UpdateCanisterSettingsRequest, UpdateCanisterSettingsResponse,
 };
 use ic_protobuf::{
@@ -251,3 +253,71 @@ pub async fn update_canister_settings(
         }
     }
 }
+
+pub async fn take_canister_snapshot(
+    take_canister_snapshot_request: TakeCanisterSnapshotRequest,
+    management_canister_client: &mut impl ManagementCanisterClient,
+) -> TakeCanisterSnapshotResponse {
+    let TakeCanisterSnapshotRequest {
+        canister_id,
+        replace_snapshot,
+    } = take_canister_snapshot_request;
+
+    let replace_snapshot = match replace_snapshot {
+        None => None,
+        Some(snapshot_id) => {
+            let snapshot_id = match SnapshotId::try_from(&snapshot_id) {
+                Ok(ok) => ok,
+                Err(err) => {
+                    return TakeCanisterSnapshotResponse::Err(TakeCanisterSnapshotError {
+                        code: None,
+                        description: format!("Invalid snapshot ID ({snapshot_id:02X?}): {err}"),
+                    });
+                }
+            };
+
+            Some(snapshot_id)
+        }
+    };
+
+    let take_canister_snapshot_args = TakeCanisterSnapshotArgs {
+        canister_id,
+        replace_snapshot,
+        uninstall_code: None,
+        sender_canister_version: management_canister_client.canister_version(),
+    };
+
+    match management_canister_client
+        .take_canister_snapshot(take_canister_snapshot_args)
+        .await
+    {
+        Ok(result) => {
+            let result =
+                convert_from_canister_snapshot_response_to_take_canister_snapshot_ok(result);
+            TakeCanisterSnapshotResponse::Ok(result)
+        }
+
+        Err((code, description)) => TakeCanisterSnapshotResponse::Err(TakeCanisterSnapshotError {
+            code: Some(code),
+            description,
+        }),
+    }
+}
+
+fn convert_from_canister_snapshot_response_to_take_canister_snapshot_ok(
+    response: CanisterSnapshotResponse,
+) -> TakeCanisterSnapshotOk {
+    let CanisterSnapshotResponse {
+        id,
+        taken_at_timestamp,
+        total_size,
+    } = response;
+
+    let id = id.to_vec();
+
+    TakeCanisterSnapshotOk {
+        id,
+        taken_at_timestamp,
+        total_size,
+    }
+}

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

@@ -7,6 +7,7 @@ DEPENDENCIES = [
     "//rs/nervous_system/clients",
     "//rs/nns/constants",
     "//rs/types/base_types",
+    "//rs/types/management_canister_types",
     "@crate_index//:candid",
     "@crate_index//:ic-cdk",
     "@crate_index//:serde",

rs/nns/handlers/root/interface/BUILD.bazel

@@ -8,6 +8,7 @@ async-trait = { workspace = true }
 candid = { workspace = true }
 ic-base-types = { path = "../../../../types/base_types" }
 ic-cdk = { workspace = true }
+ic-management-canister-types-private = { path = "../../../../types/management_canister_types" }
 ic-nervous-system-clients = { path = "../../../../nervous_system/clients" }
 ic-nns-constants = { path = "../../../../nns/constants" }
 serde = { workspace = true }

rs/nns/handlers/root/interface/Cargo.toml

@@ -82,3 +82,28 @@ pub struct UpdateCanisterSettingsError {
     pub code: Option<i32>,
     pub description: String,
 }
+
+#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)]
+pub struct TakeCanisterSnapshotRequest {
+    pub canister_id: PrincipalId,
+    pub replace_snapshot: Option</* snapshot ID */ Vec<u8>>,
+}
+
+#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)]
+pub enum TakeCanisterSnapshotResponse {
+    Ok(TakeCanisterSnapshotOk),
+    Err(TakeCanisterSnapshotError),
+}
+
+#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)]
+pub struct TakeCanisterSnapshotOk {
+    pub id: Vec<u8>,
+    pub taken_at_timestamp: u64,
+    pub total_size: u64,
+}
+
+#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)]
+pub struct TakeCanisterSnapshotError {
+    pub code: Option<i32>,
+    pub description: String,
+}

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

@@ -9,6 +9,8 @@ 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.
+
 ## Changed
 
 ## Deprecated