From 4f6bf9d18d93952bc4b4626727c996239c017620 Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Fri, 19 Dec 2025 16:59:56 +0100 Subject: [PATCH 1/4] Added take_canister_snapshot to ManagementCanisterClient. --- rs/nervous_system/clients/src/lib.rs | 1 + .../clients/src/management_canister_client.rs | 64 ++++++++++++++++++- .../src/management_canister_client/tests.rs | 7 ++ .../clients/src/take_canister_snapshot.rs | 14 ++++ 4 files changed, 85 insertions(+), 1 deletion(-) create mode 100644 rs/nervous_system/clients/src/take_canister_snapshot.rs diff --git a/rs/nervous_system/clients/src/lib.rs b/rs/nervous_system/clients/src/lib.rs index 10d7db52578c..6434cd2cc570 100644 --- a/rs/nervous_system/clients/src/lib.rs +++ b/rs/nervous_system/clients/src/lib.rs @@ -5,6 +5,7 @@ pub mod delete_canister; pub mod ledger_client; pub mod management_canister_client; pub mod stop_canister; +pub mod take_canister_snapshot; pub mod update_settings; mod request; diff --git a/rs/nervous_system/clients/src/management_canister_client.rs b/rs/nervous_system/clients/src/management_canister_client.rs index 15d857bd9a05..1e1ff6780a93 100644 --- a/rs/nervous_system/clients/src/management_canister_client.rs +++ b/rs/nervous_system/clients/src/management_canister_client.rs @@ -4,13 +4,16 @@ use crate::{ canister_status::{CanisterStatusResultFromManagementCanister, canister_status}, delete_canister::delete_canister, stop_canister::stop_canister, + take_canister_snapshot::take_canister_snapshot, update_settings::{UpdateSettings, update_settings}, }; use async_trait::async_trait; use candid::Encode; use ic_base_types::PrincipalId; use ic_error_types::RejectCode; -use ic_management_canister_types_private::IC_00; +use ic_management_canister_types_private::{ + CanisterSnapshotResponse, IC_00, TakeCanisterSnapshotArgs, +}; use ic_nervous_system_proxied_canister_calls_tracker::ProxiedCanisterCallsTracker; use ic_nervous_system_runtime::Runtime; use std::{ @@ -54,6 +57,11 @@ pub trait ManagementCanisterClient { &self, canister_id_record: CanisterIdRecord, ) -> Result<(), (i32, String)>; + + async fn take_canister_snapshot( + &self, + args: TakeCanisterSnapshotArgs, + ) -> Result; } /// An example implementation of the ManagementCanisterClient trait. @@ -171,6 +179,24 @@ impl ManagementCanisterClient for ManagementCanisterClientIm delete_canister::(canister_id_record).await } + + async fn take_canister_snapshot( + &self, + args: TakeCanisterSnapshotArgs, + ) -> Result { + 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, + "take_canister_snapshot", + &encoded_args, + ) + }); + + take_canister_snapshot::(args).await + } } /// A ManagementCanisterClient that wraps another ManagementCanisterClient. @@ -288,6 +314,14 @@ where let _loan = self.try_borrow_slot()?; self.inner.delete_canister(canister_id_record).await } + + async fn take_canister_snapshot( + &self, + args: TakeCanisterSnapshotArgs, + ) -> Result { + let _loan = self.try_borrow_slot()?; + self.inner.take_canister_snapshot(args).await + } } /// Increments available_slot_count by used_slot_count when dropped. @@ -341,6 +375,7 @@ pub enum MockManagementCanisterClientCall { CanisterMetadata(PrincipalId, String), StopCanister(CanisterIdRecord), DeleteCanister(CanisterIdRecord), + TakeCanisterSnapshot(TakeCanisterSnapshotArgs), } #[derive(Clone, Eq, PartialEq, Debug)] @@ -351,6 +386,7 @@ pub enum MockManagementCanisterClientReply { CanisterMetadata(Result, (i32, String)>), StopCanister(Result<(), (i32, String)>), DeleteCanister(Result<(), (i32, String)>), + TakeCanisterSnapshot(Result), } #[async_trait] @@ -493,6 +529,32 @@ impl ManagementCanisterClient for MockManagementCanisterClient { ), } } + + async fn take_canister_snapshot( + &self, + args: TakeCanisterSnapshotArgs, + ) -> Result { + self.calls + .lock() + .unwrap() + .push_back(MockManagementCanisterClientCall::TakeCanisterSnapshot(args)); + + let reply = self + .replies + .lock() + .unwrap() + .pop_front() + .expect("Expected a MockManagementCanisterClientCall to be on the queue."); + + match reply { + MockManagementCanisterClientReply::TakeCanisterSnapshot(result) => result, + err => panic!( + "Expected MockManagementCanisterClientReply::TakeCanisterSnapshot to be at \ + the front of the queue. Had {:?}", + err + ), + } + } } impl Drop for MockManagementCanisterClient { diff --git a/rs/nervous_system/clients/src/management_canister_client/tests.rs b/rs/nervous_system/clients/src/management_canister_client/tests.rs index 3cde2358dbea..fbefab470b1f 100644 --- a/rs/nervous_system/clients/src/management_canister_client/tests.rs +++ b/rs/nervous_system/clients/src/management_canister_client/tests.rs @@ -84,6 +84,13 @@ async fn test_limit_outstanding_calls() { ) -> Result<(), (i32, String)> { unimplemented!(); } + + async fn take_canister_snapshot( + &self, + _args: TakeCanisterSnapshotArgs, + ) -> Result { + unimplemented!(); + } } impl Drop for MockManagementCanisterClient { diff --git a/rs/nervous_system/clients/src/take_canister_snapshot.rs b/rs/nervous_system/clients/src/take_canister_snapshot.rs new file mode 100644 index 000000000000..4262e23ba1c1 --- /dev/null +++ b/rs/nervous_system/clients/src/take_canister_snapshot.rs @@ -0,0 +1,14 @@ +use ic_management_canister_types_private::{ + CanisterSnapshotResponse, IC_00, TakeCanisterSnapshotArgs, +}; +use ic_nervous_system_runtime::Runtime; + +pub async fn take_canister_snapshot( + args: TakeCanisterSnapshotArgs, +) -> Result +where + Rt: Runtime, +{ + let (res,) = Rt::call_with_cleanup(IC_00, "take_canister_snapshot", (args,)).await?; + Ok(res) +} From 65f1c5c82aae32fde293392e6119c9149dea7111 Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Fri, 19 Dec 2025 17:36:35 +0100 Subject: [PATCH 2/4] Added take_canister_snapshot to Root canister. It is only callable by Governance. --- Cargo.lock | 1 + .../handlers/root/impl/canister/canister.rs | 17 ++++- rs/nns/handlers/root/impl/canister/root.did | 24 ++++++ .../root/impl/src/canister_management.rs | 76 ++++++++++++++++++- rs/nns/handlers/root/interface/BUILD.bazel | 1 + rs/nns/handlers/root/interface/Cargo.toml | 1 + rs/nns/handlers/root/interface/src/lib.rs | 25 ++++++ 7 files changed, 141 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3781688ad4fc..8f8e407838e8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12327,6 +12327,7 @@ dependencies = [ "candid", "ic-base-types", "ic-cdk", + "ic-management-canister-types-private", "ic-nervous-system-clients", "ic-nns-constants", "serde", diff --git a/rs/nns/handlers/root/impl/canister/canister.rs b/rs/nns/handlers/root/impl/canister/canister.rs index 56950c6132e6..632e42b3a131 100644 --- a/rs/nns/handlers/root/impl/canister/canister.rs +++ b/rs/nns/handlers/root/impl/canister/canister.rs @@ -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( diff --git a/rs/nns/handlers/root/impl/canister/root.did b/rs/nns/handlers/root/impl/canister/root.did index 502c9c0892c6..8ef061a7b0be 100644 --- a/rs/nns/handlers/root/impl/canister/root.did +++ b/rs/nns/handlers/root/impl/canister/root.did @@ -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, + ); } diff --git a/rs/nns/handlers/root/impl/src/canister_management.rs b/rs/nns/handlers/root/impl/src/canister_management.rs index 0cbbc800a547..34b3df0a49ea 100644 --- a/rs/nns/handlers/root/impl/src/canister_management.rs +++ b/rs/nns/handlers/root/impl/src/canister_management.rs @@ -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_resposne_to_take_canister_snapshot_ok(result); + TakeCanisterSnapshotResponse::Ok(result) + } + + Err((code, description)) => TakeCanisterSnapshotResponse::Err(TakeCanisterSnapshotError { + code: Some(code), + description, + }), + } +} + +fn convert_from_canister_snapshot_resposne_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, + } +} diff --git a/rs/nns/handlers/root/interface/BUILD.bazel b/rs/nns/handlers/root/interface/BUILD.bazel index f567f6c99f3e..1d91d0e8daf3 100644 --- a/rs/nns/handlers/root/interface/BUILD.bazel +++ b/rs/nns/handlers/root/interface/BUILD.bazel @@ -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", diff --git a/rs/nns/handlers/root/interface/Cargo.toml b/rs/nns/handlers/root/interface/Cargo.toml index 95dd4ee6ca08..3e1f39450458 100644 --- a/rs/nns/handlers/root/interface/Cargo.toml +++ b/rs/nns/handlers/root/interface/Cargo.toml @@ -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 } diff --git a/rs/nns/handlers/root/interface/src/lib.rs b/rs/nns/handlers/root/interface/src/lib.rs index 0bca606f7a08..cd7e72c2654c 100644 --- a/rs/nns/handlers/root/interface/src/lib.rs +++ b/rs/nns/handlers/root/interface/src/lib.rs @@ -82,3 +82,28 @@ pub struct UpdateCanisterSettingsError { pub code: Option, pub description: String, } + +#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] +pub struct TakeCanisterSnapshotRequest { + pub canister_id: PrincipalId, + pub replace_snapshot: Option>, +} + +#[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, + pub taken_at_timestamp: u64, + pub total_size: u64, +} + +#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] +pub struct TakeCanisterSnapshotError { + pub code: Option, + pub description: String, +} From 710aa82ca049ac17d787b960f08975f190d91aca Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Fri, 19 Dec 2025 18:04:56 +0100 Subject: [PATCH 3/4] changelog --- rs/nns/handlers/root/unreleased_changelog.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/rs/nns/handlers/root/unreleased_changelog.md b/rs/nns/handlers/root/unreleased_changelog.md index 94126a0ff421..2b7ee204e396 100644 --- a/rs/nns/handlers/root/unreleased_changelog.md +++ b/rs/nns/handlers/root/unreleased_changelog.md @@ -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 From 46816b23e052aaaa90566989c62c656ef4d66b33 Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Tue, 6 Jan 2026 09:17:14 +0100 Subject: [PATCH 4/4] typo --- rs/nns/handlers/root/impl/src/canister_management.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rs/nns/handlers/root/impl/src/canister_management.rs b/rs/nns/handlers/root/impl/src/canister_management.rs index 34b3df0a49ea..1be377c1b8a7 100644 --- a/rs/nns/handlers/root/impl/src/canister_management.rs +++ b/rs/nns/handlers/root/impl/src/canister_management.rs @@ -293,7 +293,7 @@ pub async fn take_canister_snapshot( { Ok(result) => { let result = - convert_from_canister_snapshot_resposne_to_take_canister_snapshot_ok(result); + convert_from_canister_snapshot_response_to_take_canister_snapshot_ok(result); TakeCanisterSnapshotResponse::Ok(result) } @@ -304,7 +304,7 @@ pub async fn take_canister_snapshot( } } -fn convert_from_canister_snapshot_resposne_to_take_canister_snapshot_ok( +fn convert_from_canister_snapshot_response_to_take_canister_snapshot_ok( response: CanisterSnapshotResponse, ) -> TakeCanisterSnapshotOk { let CanisterSnapshotResponse {