ic_testkit/pocket_ic/
snapshot.rs

1//! Bounded native snapshot copying. Callers own funding, custody and recovery.
2
3use crate::pocket_ic::{ClientError, common::rest::RawMessageId};
4use candid::Principal;
5pub use ic_management_canister_types::{
6    ReadCanisterSnapshotDataArgs, ReadCanisterSnapshotDataResult, ReadCanisterSnapshotMetadataArgs,
7    ReadCanisterSnapshotMetadataResult, SnapshotDataKind, SnapshotDataOffset,
8    UploadCanisterSnapshotDataArgs, UploadCanisterSnapshotMetadataArgs,
9    UploadCanisterSnapshotMetadataResult,
10};
11
12/// Explicit per-call bounds, including the outer PocketIC HTTP/JSON envelope.
13#[derive(Clone, Copy, Debug)]
14pub struct CallLimits {
15    pub http_bytes: usize,
16    pub candid_bytes: usize,
17    pub decoding_quota: usize,
18    pub skipping_quota: usize,
19}
20
21/// Caller-selected memory/work limits. No implicit application budget is added.
22#[derive(Clone, Copy, Debug)]
23pub struct TransferLimits {
24    pub call: CallLimits,
25    pub chunk_bytes: u64,
26    pub total_bytes: u64,
27}
28
29#[derive(Debug, thiserror::Error)]
30pub enum SnapshotError {
31    #[error("snapshot call failed: {source}")]
32    Call {
33        #[source]
34        source: ClientError,
35        /// Retained even when Candid decoding fails after accepted ingress.
36        ingress: Option<RawMessageId>,
37    },
38    #[error("invalid snapshot transfer: {0}")]
39    Invalid(&'static str),
40}
41
42/// The last attempted action. An interrupted action cannot be retried implicitly.
43#[derive(Clone, Debug)]
44pub enum TransferStatus {
45    Ready,
46    UploadingMetadata,
47    Reading(SnapshotDataKind),
48    Writing(SnapshotDataOffset),
49    Complete,
50}
51
52/// Caller-owned progress survives errors and dropped async futures.
53///
54/// First advance uploads metadata only: record `uploaded_snapshot()` in the
55/// caller's pending set before advancing again. Each later advance copies at most
56/// one bounded chunk. On error/cancellation the status retains the attempted
57/// action and further advances refuse. Reconcile explicitly using the raw methods;
58/// this object never deletes snapshots, loads them, retries or performs Drop I/O.
59#[derive(Debug)]
60pub struct SnapshotTransfer {
61    pub(crate) canister: Principal,
62    pub(crate) sender: Principal,
63    pub(crate) original: Vec<u8>,
64    pub(crate) metadata: ReadCanisterSnapshotMetadataResult,
65    pub(crate) limits: TransferLimits,
66    pub(crate) uploaded: Option<Vec<u8>>,
67    pub(crate) status: TransferStatus,
68    pub(crate) ingress: Option<RawMessageId>,
69    pub(crate) region: usize,
70    pub(crate) offset: u64,
71    pub(crate) transferred: u64,
72}
73
74impl SnapshotTransfer {
75    pub fn new(
76        canister: Principal,
77        sender: Principal,
78        original: Vec<u8>,
79        metadata: ReadCanisterSnapshotMetadataResult,
80        limits: TransferLimits,
81    ) -> Result<Self, SnapshotError> {
82        if limits.chunk_bytes == 0 || limits.chunk_bytes > 1024 * 1024 {
83            return Err(SnapshotError::Invalid("chunk size must be 1..=1 MiB"));
84        }
85        if metadata.global_timer.is_none() || metadata.on_low_wasm_memory_hook_status.is_none() {
86            return Err(SnapshotError::Invalid(
87                "snapshot omitted timer or low-memory hook state",
88            ));
89        }
90        let bytes = metadata
91            .wasm_module_size
92            .checked_add(metadata.wasm_memory_size)
93            .and_then(|bytes| bytes.checked_add(metadata.stable_memory_size))
94            .ok_or(SnapshotError::Invalid("snapshot region sizes overflow"))?;
95        if bytes > limits.total_bytes {
96            return Err(SnapshotError::Invalid("snapshot exceeds total byte budget"));
97        }
98        let mut hashes = std::collections::BTreeSet::new();
99        for chunk in &metadata.wasm_chunk_store {
100            if chunk.hash.len() != 32 || !hashes.insert(chunk.hash.clone()) {
101                return Err(SnapshotError::Invalid("invalid or duplicate chunk hash"));
102            }
103        }
104        Ok(Self {
105            canister,
106            sender,
107            original,
108            metadata,
109            limits,
110            uploaded: None,
111            status: TransferStatus::Ready,
112            ingress: None,
113            region: 0,
114            offset: 0,
115            transferred: 0,
116        })
117    }
118
119    pub fn uploaded_snapshot(&self) -> Option<&[u8]> {
120        self.uploaded.as_deref()
121    }
122    pub fn status(&self) -> &TransferStatus {
123        &self.status
124    }
125    pub fn ingress(&self) -> Option<&RawMessageId> {
126        self.ingress.as_ref()
127    }
128    pub fn transferred_bytes(&self) -> u64 {
129        self.transferred
130    }
131
132    /// Synchronous frontend for the same async implementation.
133    pub fn advance(&mut self, client: &crate::pocket_ic::PocketIc) -> Result<(), SnapshotError> {
134        client
135            .runtime
136            .block_on(self.advance_async(&client.pocket_ic))
137    }
138}
139
140impl crate::pocket_ic::PocketIc {
141    pub fn read_canister_snapshot_metadata(
142        &self,
143        sender: Principal,
144        args: ReadCanisterSnapshotMetadataArgs,
145        limits: CallLimits,
146    ) -> Result<ReadCanisterSnapshotMetadataResult, SnapshotError> {
147        self.runtime.block_on(
148            self.pocket_ic
149                .read_canister_snapshot_metadata(sender, args, limits),
150        )
151    }
152    pub fn read_canister_snapshot_data(
153        &self,
154        sender: Principal,
155        args: ReadCanisterSnapshotDataArgs,
156        limits: CallLimits,
157    ) -> Result<ReadCanisterSnapshotDataResult, SnapshotError> {
158        self.runtime.block_on(
159            self.pocket_ic
160                .read_canister_snapshot_data(sender, args, limits),
161        )
162    }
163    pub fn upload_canister_snapshot_metadata(
164        &self,
165        sender: Principal,
166        args: UploadCanisterSnapshotMetadataArgs,
167        limits: CallLimits,
168    ) -> Result<UploadCanisterSnapshotMetadataResult, SnapshotError> {
169        self.runtime.block_on(
170            self.pocket_ic
171                .upload_canister_snapshot_metadata(sender, args, limits),
172        )
173    }
174    pub fn upload_canister_snapshot_data(
175        &self,
176        sender: Principal,
177        args: UploadCanisterSnapshotDataArgs,
178        limits: CallLimits,
179    ) -> Result<(), SnapshotError> {
180        self.runtime.block_on(
181            self.pocket_ic
182                .upload_canister_snapshot_data(sender, args, limits),
183        )
184    }
185}
ic_testkit/pocket_ic/snapshot.rs

ic_testkit/pocket_ic/
snapshot.rs
