feat: [EXC-1955] Read canister snapshot metadata (#4514)

This PR implements the download of snapshot metadata. 

The feature is currently behind a feature flag in the execution config.
When disabled, the mgmt canister endpoint returns an empty blob.

---------

Co-authored-by: Stefan Schneider <31004026+schneiderstefan@users.noreply.github.com>

Author: michael-weigelt

rs/execution_environment/src/canister_manager.rs

@@ -20,10 +20,11 @@ use ic_interfaces::execution_environment::{IngressHistoryWriter, SubnetAvailable
 use ic_logger::{error, fatal, info, ReplicaLogger};
 use ic_management_canister_types_private::{
     CanisterChangeDetails, CanisterChangeOrigin, CanisterInstallModeV2, CanisterSnapshotResponse,
-    CanisterStatusResultV2, CanisterStatusType, ChunkHash, Method as Ic00Method, StoredChunksReply,
-    UploadChunkReply,
+    CanisterStatusResultV2, CanisterStatusType, ChunkHash, Method as Ic00Method,
+    ReadCanisterSnapshotMetadataResponse, StoredChunksReply, UploadChunkReply,
 };
 use ic_registry_provisional_whitelist::ProvisionalWhitelist;
+use ic_replicated_state::canister_state::WASM_PAGE_SIZE_IN_BYTES;
 use ic_replicated_state::{
     canister_snapshots::CanisterSnapshot,
     canister_state::{
@@ -1981,6 +1982,54 @@ impl CanisterManager {
             WasmExecutionMode::Wasm64 => self.config.max_canister_memory_size_wasm64,
         }
     }
+
+    pub(crate) fn read_snapshot_metadata(
+        &self,
+        sender: PrincipalId,
+        snapshot_id: SnapshotId,
+        canister: &CanisterState,
+        state: &ReplicatedState,
+    ) -> Result<ReadCanisterSnapshotMetadataResponse, CanisterManagerError> {
+        // Check sender is a controller.
+        validate_controller(canister, &sender)?;
+        // If not found, the operation fails due to invalid parameters.
+        let Some(snapshot) = state.canister_snapshots.get(snapshot_id) else {
+            return Err(CanisterManagerError::CanisterSnapshotNotFound {
+                canister_id: canister.canister_id(),
+                snapshot_id,
+            });
+        };
+        // Verify the provided `snapshot_id` belongs to this canister.
+        if snapshot.canister_id() != canister.canister_id() {
+            return Err(CanisterManagerError::CanisterSnapshotInvalidOwnership {
+                canister_id: canister.canister_id(),
+                snapshot_id,
+            });
+        }
+
+        Ok(ReadCanisterSnapshotMetadataResponse {
+            source: snapshot.source(),
+            taken_at_timestamp: snapshot.taken_at_timestamp().as_nanos_since_unix_epoch(),
+            wasm_module_size: snapshot.execution_snapshot().wasm_binary.len() as u64,
+            exported_globals: snapshot.exported_globals().clone(),
+            wasm_memory_size: snapshot.execution_snapshot().wasm_memory.size.get() as u64
+                * WASM_PAGE_SIZE_IN_BYTES as u64,
+            stable_memory_size: snapshot.execution_snapshot().stable_memory.size.get() as u64
+                * WASM_PAGE_SIZE_IN_BYTES as u64,
+            wasm_chunk_store: snapshot
+                .chunk_store()
+                .keys()
+                .cloned()
+                .map(|x| ChunkHash { hash: x.to_vec() })
+                .collect(),
+            canister_version: snapshot.canister_version(),
+            certified_data: snapshot.certified_data().clone(),
+            global_timer: snapshot.execution_snapshot().global_timer.into(),
+            on_low_wasm_memory_hook_status: snapshot
+                .execution_snapshot()
+                .on_low_wasm_memory_hook_status,
+        })
+    }
 }
 
 /// Uninstalls a canister.

rs/execution_environment/src/execution_environment.rs

@@ -1657,10 +1657,16 @@ impl ExecutionEnvironment {
             }
 
             Ok(Ic00Method::ReadCanisterSnapshotMetadata) => {
-                // TODO: EXC-1955
-                #[allow(clippy::bind_instead_of_map)]
-                let res = ReadCanisterSnapshotMetadataArgs::decode(payload)
-                    .and_then(|_args| Ok((vec![], None)));
+                let res = ReadCanisterSnapshotMetadataArgs::decode(payload).and_then(|args| {
+                    match self.config.canister_snapshot_download {
+                        FlagStatus::Disabled => Ok((vec![], None)),
+                        FlagStatus::Enabled => {
+                            let canister_id = args.get_canister_id();
+                            self.read_canister_snapshot_metadata(*msg.sender(), &state, args)
+                                .map(|x| (x, Some(canister_id)))
+                        }
+                    }
+                });
                 ExecuteSubnetMessageResult::Finished {
                     response: res,
                     refund: msg.take_cycles(),
@@ -1670,8 +1676,13 @@ impl ExecutionEnvironment {
             Ok(Ic00Method::ReadCanisterSnapshotData) => {
                 // TODO: EXC-1957
                 #[allow(clippy::bind_instead_of_map)]
-                let res = ReadCanisterSnapshotDataArgs::decode(payload)
-                    .and_then(|_args| Ok((vec![], None)));
+                let res =
+                    ReadCanisterSnapshotDataArgs::decode(payload).and_then(|_args| {
+                        match self.config.canister_snapshot_download {
+                            FlagStatus::Disabled => Ok((vec![], None)),
+                            FlagStatus::Enabled => Ok((vec![], None)),
+                        }
+                    });
                 ExecuteSubnetMessageResult::Finished {
                     response: res,
                     refund: msg.take_cycles(),
@@ -1681,8 +1692,13 @@ impl ExecutionEnvironment {
             Ok(Ic00Method::UploadCanisterSnapshotMetadata) => {
                 // TODO: EXC-1959
                 #[allow(clippy::bind_instead_of_map)]
-                let res = UploadCanisterSnapshotMetadataArgs::decode(payload)
-                    .and_then(|_args| Ok((vec![], None)));
+                let res =
+                    UploadCanisterSnapshotMetadataArgs::decode(payload).and_then(
+                        |_args| match self.config.canister_snapshot_upload {
+                            FlagStatus::Disabled => Ok((vec![], None)),
+                            FlagStatus::Enabled => Ok((vec![], None)),
+                        },
+                    );
                 ExecuteSubnetMessageResult::Finished {
                     response: res,
                     refund: msg.take_cycles(),
@@ -1692,8 +1708,12 @@ impl ExecutionEnvironment {
             Ok(Ic00Method::UploadCanisterSnapshotData) => {
                 // TODO: EXC-1960
                 #[allow(clippy::bind_instead_of_map)]
-                let res = UploadCanisterSnapshotDataArgs::decode(payload)
-                    .and_then(|_args| Ok((vec![], None)));
+                let res = UploadCanisterSnapshotDataArgs::decode(payload).and_then(|_args| {
+                    match self.config.canister_snapshot_upload {
+                        FlagStatus::Disabled => Ok((vec![], None)),
+                        FlagStatus::Enabled => Ok((vec![], None)),
+                    }
+                });
                 ExecuteSubnetMessageResult::Finished {
                     response: res,
                     refund: msg.take_cycles(),
@@ -2428,6 +2448,20 @@ impl ExecutionEnvironment {
         result
     }
 
+    fn read_canister_snapshot_metadata(
+        &self,
+        sender: PrincipalId,
+        state: &ReplicatedState,
+        args: ReadCanisterSnapshotMetadataArgs,
+    ) -> Result<Vec<u8>, UserError> {
+        let canister = get_canister(args.get_canister_id(), state)?;
+        let snapshot_id = args.get_snapshot_id();
+        self.canister_manager
+            .read_snapshot_metadata(sender, snapshot_id, canister, state)
+            .map(|res| Encode!(&res).unwrap())
+            .map_err(|e| e.into())
+    }
+
     fn node_metrics_history(
         &self,
         state: &ReplicatedState,

rs/execution_environment/src/execution_environment/tests/canister_snapshots.rs

@@ -6,16 +6,18 @@ use ic_cycles_account_manager::ResourceSaturation;
 use ic_error_types::{ErrorCode, RejectCode};
 use ic_management_canister_types_private::{
     self as ic00, CanisterChange, CanisterChangeDetails, CanisterSettingsArgsBuilder,
-    CanisterSnapshotResponse, ClearChunkStoreArgs, DeleteCanisterSnapshotArgs,
+    CanisterSnapshotResponse, ClearChunkStoreArgs, DeleteCanisterSnapshotArgs, GlobalTimer,
     ListCanisterSnapshotArgs, LoadCanisterSnapshotArgs, Method, OnLowWasmMemoryHookStatus,
-    Payload as Ic00Payload, TakeCanisterSnapshotArgs, UpdateSettingsArgs, UploadChunkArgs,
+    Payload as Ic00Payload, ReadCanisterSnapshotMetadataArgs, ReadCanisterSnapshotMetadataResponse,
+    SnapshotSource, TakeCanisterSnapshotArgs, UpdateSettingsArgs, UploadChunkArgs,
 };
 use ic_registry_subnet_type::SubnetType;
 use ic_replicated_state::{
     canister_snapshots::SnapshotOperation,
     canister_state::{
         execution_state::{WasmBinary, WasmExecutionMode},
         system_state::CyclesUseCase,
+        WASM_PAGE_SIZE_IN_BYTES,
     },
     CanisterState, ExecutionState, SchedulerState,
 };
@@ -30,6 +32,7 @@ use ic_types::{
     time::UNIX_EPOCH,
     CanisterId, Cycles, NumInstructions, SnapshotId,
 };
+use ic_types_test_utils::ids::user_test_id;
 use ic_universal_canister::{wasm, UNIVERSAL_CANISTER_WASM};
 use more_asserts::assert_gt;
 use serde_bytes::ByteBuf;
@@ -1995,6 +1998,155 @@ fn snapshot_must_include_globals() {
     assert_eq!(result, WasmResult::Reply(vec![1, 0, 0, 0]));
 }
 
+#[test]
+fn read_canister_snapshot_metadata_succeeds() {
+    let own_subnet = subnet_test_id(1);
+    let caller_canister = canister_test_id(1);
+    let mut test = ExecutionTestBuilder::new()
+        .with_snapshot_metadata_download()
+        .with_own_subnet_id(own_subnet)
+        .with_caller(own_subnet, caller_canister)
+        .build();
+    // Create new canister.
+    let uni_canister_wasm = UNIVERSAL_CANISTER_WASM.to_vec();
+    let canister_id = test
+        .canister_from_cycles_and_binary(
+            Cycles::new(1_000_000_000_000_000),
+            uni_canister_wasm.clone(),
+        )
+        .unwrap();
+
+    // Upload chunk.
+    let chunk = vec![1, 2, 3, 4, 5];
+    let upload_args = UploadChunkArgs {
+        canister_id: canister_id.into(),
+        chunk,
+    };
+    let result = test.subnet_message("upload_chunk", upload_args.encode());
+    assert!(result.is_ok());
+    // Grow the stable memory
+    let stable_pages = 13;
+    let payload = wasm().stable64_grow(stable_pages).reply().build();
+    let _res = test.ingress(canister_id, "update", payload).unwrap();
+    // Set some cert data
+    let cert_data = [42];
+    let payload = wasm().certified_data_set(&cert_data).reply().build();
+    let _res = test.ingress(canister_id, "update", payload).unwrap();
+    // Set a global timer
+    let timestamp = 43;
+    let payload = wasm().api_global_timer_set(timestamp).reply().build();
+    let _res = test.ingress(canister_id, "update", payload).unwrap();
+
+    // Take a snapshot of the canister.
+    let args: TakeCanisterSnapshotArgs = TakeCanisterSnapshotArgs::new(canister_id, None);
+    let result = test.subnet_message("take_canister_snapshot", args.encode());
+    let snapshot_id = CanisterSnapshotResponse::decode(&result.unwrap().bytes())
+        .unwrap()
+        .snapshot_id();
+
+    // Get the metadata
+    let args = ReadCanisterSnapshotMetadataArgs::new(canister_id, snapshot_id);
+    let WasmResult::Reply(bytes) = test
+        .subnet_message("read_canister_snapshot_metadata", args.encode())
+        .unwrap()
+    else {
+        panic!("expected WasmResult::Reply")
+    };
+    let metadata = Decode!(&bytes, ReadCanisterSnapshotMetadataResponse).unwrap();
+    assert_eq!(metadata.source, SnapshotSource::TakenFromCanister);
+    assert_eq!(
+        metadata.stable_memory_size,
+        WASM_PAGE_SIZE_IN_BYTES as u64 * stable_pages
+    );
+    assert_eq!(metadata.wasm_module_size, uni_canister_wasm.len() as u64);
+    assert_eq!(metadata.wasm_chunk_store.len(), 1);
+    assert_eq!(metadata.certified_data, cert_data);
+    assert_eq!(metadata.global_timer, GlobalTimer::Active(timestamp));
+    assert_eq!(metadata.canister_version, 4);
+}
+
+#[test]
+fn read_canister_snapshot_metadata_fails_canister_and_snapshot_must_match() {
+    let own_subnet = subnet_test_id(1);
+    let caller_canister = canister_test_id(1);
+    let mut test = ExecutionTestBuilder::new()
+        .with_snapshot_metadata_download()
+        .with_own_subnet_id(own_subnet)
+        .with_manual_execution()
+        .with_caller(own_subnet, caller_canister)
+        .build();
+
+    // Create canister
+    let canister_id = test
+        .canister_from_cycles_and_binary(
+            Cycles::new(1_000_000_000_000_000),
+            UNIVERSAL_CANISTER_WASM.to_vec(),
+        )
+        .unwrap();
+
+    // Create other canister.
+    let other_canister_id = test
+        .canister_from_cycles_and_binary(
+            Cycles::new(1_000_000_000_000_000),
+            UNIVERSAL_CANISTER_WASM.to_vec(),
+        )
+        .unwrap();
+
+    // Take a snapshot of the first canister.
+    let args: TakeCanisterSnapshotArgs = TakeCanisterSnapshotArgs::new(canister_id, None);
+    let result = test.subnet_message("take_canister_snapshot", args.encode());
+    let snapshot_id = CanisterSnapshotResponse::decode(&result.unwrap().bytes())
+        .unwrap()
+        .snapshot_id();
+
+    // Try to access metadata via the wrong canister id
+    let args = ReadCanisterSnapshotMetadataArgs::new(other_canister_id, snapshot_id);
+    let error = test
+        .subnet_message("read_canister_snapshot_metadata", args.encode())
+        .unwrap_err();
+    assert_eq!(error.code(), ErrorCode::CanisterRejectedMessage);
+
+    // Try to access metadata via a bad snapshot id
+    let args = ReadCanisterSnapshotMetadataArgs::new(canister_id, (canister_id, 42).into());
+    let error = test
+        .subnet_message("read_canister_snapshot_metadata", args.encode())
+        .unwrap_err();
+    assert_eq!(error.code(), ErrorCode::CanisterSnapshotNotFound);
+}
+
+#[test]
+fn read_canister_snapshot_metadata_fails_invalid_controller() {
+    let own_subnet = subnet_test_id(1);
+    let mut test = ExecutionTestBuilder::new()
+        .with_snapshot_metadata_download()
+        .with_own_subnet_id(own_subnet)
+        .with_manual_execution()
+        .build();
+    // Create new canister.
+    let uni_canister_wasm = UNIVERSAL_CANISTER_WASM.to_vec();
+    let canister_id = test
+        .canister_from_cycles_and_binary(
+            Cycles::new(1_000_000_000_000_000),
+            uni_canister_wasm.clone(),
+        )
+        .unwrap();
+
+    // Take a snapshot of the canister.
+    let args: TakeCanisterSnapshotArgs = TakeCanisterSnapshotArgs::new(canister_id, None);
+    let result = test.subnet_message("take_canister_snapshot", args.encode());
+    let snapshot_id = CanisterSnapshotResponse::decode(&result.unwrap().bytes())
+        .unwrap()
+        .snapshot_id();
+
+    // Non-controller user tries to get metadata
+    test.set_user_id(user_test_id(42));
+    let args = ReadCanisterSnapshotMetadataArgs::new(canister_id, snapshot_id);
+    let error = test
+        .subnet_message("read_canister_snapshot_metadata", args.encode())
+        .unwrap_err();
+    assert_eq!(error.code(), ErrorCode::CanisterInvalidController);
+}
+
 /// Early warning system / stumbling block forcing the authors of changes adding
 /// or removing canister state fields to think about and/or ask the Execution
 /// team to think about any repercussions to the canister snapshot logic.

rs/execution_environment/tests/dts.rs

@@ -1231,14 +1231,17 @@ fn dts_aborted_execution_does_not_block_subnet_messages() {
                 (method, call_args().other_side(args))
             }),
             Method::ReadCanisterSnapshotMetadata => test_supported(|aborted_canister_id| {
-                let args =
-                    ReadCanisterSnapshotMetadataArgs::new(aborted_canister_id, vec![]).encode();
+                let args = ReadCanisterSnapshotMetadataArgs::new(
+                    aborted_canister_id,
+                    (aborted_canister_id, 0).into(),
+                )
+                .encode();
                 (method, call_args().other_side(args))
             }),
             Method::ReadCanisterSnapshotData => test_supported(|aborted_canister_id| {
                 let args = ReadCanisterSnapshotDataArgs::new(
                     aborted_canister_id,
-                    vec![],
+                    (aborted_canister_id, 0).into(),
                     CanisterSnapshotDataKind::WasmModule { size: 0, offset: 0 },
                 )
                 .encode();

rs/test_utilities/execution_environment/src/lib.rs

@@ -2204,6 +2204,16 @@ impl ExecutionTestBuilder {
         self
     }
 
+    pub fn with_snapshot_metadata_download(mut self) -> Self {
+        self.execution_config.canister_snapshot_download = FlagStatus::Enabled;
+        self
+    }
+
+    pub fn with_snapshot_metadata_upload(mut self) -> Self {
+        self.execution_config.canister_snapshot_upload = FlagStatus::Enabled;
+        self
+    }
+
     pub fn build(self) -> ExecutionTest {
         let own_range = CanisterIdRange {
             start: CanisterId::from(CANISTER_IDS_PER_SUBNET),