ic_testkit/pocket_ic/nonblocking/
snapshot.rs

1use super::*;
2use crate::pocket_ic::snapshot::*;
3use candid::{CandidType, DecoderConfig};
4
5impl PocketIc {
6    async fn snapshot_call<A: CandidType, R: for<'de> ArgumentDecoder<'de>>(
7        &self,
8        canister: Principal,
9        sender: Principal,
10        method: &str,
11        args: A,
12        limits: CallLimits,
13        ingress: &mut Option<RawMessageId>,
14    ) -> Result<R, SnapshotError> {
15        let result = async {
16            let payload = candid::encode_one(args)?;
17            if payload.len() > limits.candid_bytes {
18                return Err(ClientError::Protocol(
19                    "snapshot request exceeds Candid byte budget".into(),
20                ));
21            }
22            let message: Result<RawMessageId, RejectResponse> = self
23                .request_limited(
24                    HttpMethod::Post,
25                    "update/submit_ingress_message",
26                    RawCanisterCall {
27                        canister_id: Principal::management_canister().as_slice().to_vec(),
28                        effective_principal: RawEffectivePrincipal::CanisterId(
29                            canister.as_slice().to_vec(),
30                        ),
31                        sender: sender.as_slice().to_vec(),
32                        method: method.into(),
33                        payload,
34                        sender_info: None,
35                    },
36                    Some(limits.http_bytes),
37                )
38                .await?;
39            let message = message.map_err(ClientError::from)?;
40            *ingress = Some(message.clone());
41            let result: RawCanisterResult = self
42                .request_limited(
43                    HttpMethod::Post,
44                    "update/await_ingress_message",
45                    &message,
46                    Some(limits.http_bytes),
47                )
48                .await
49                .map_err(|error| error.with_ingress(message))?;
50            let reply =
51                Result::<Vec<u8>, RejectResponse>::from(result).map_err(ClientError::from)?;
52            if reply.len() > limits.candid_bytes {
53                return Err(ClientError::Protocol(
54                    "snapshot reply exceeds Candid byte budget".into(),
55                ));
56            }
57            let mut config = DecoderConfig::new();
58            config.set_full_error_message(false);
59            config.set_decoding_quota(limits.decoding_quota);
60            config.set_skipping_quota(limits.skipping_quota);
61            candid::decode_args_with_config(&reply, &config).map_err(ClientError::from)
62        }
63        .await;
64        result.map_err(|source| SnapshotError::Call {
65            source,
66            ingress: ingress.clone(),
67        })
68    }
69
70    pub async fn read_canister_snapshot_metadata(
71        &self,
72        sender: Principal,
73        args: ReadCanisterSnapshotMetadataArgs,
74        limits: CallLimits,
75    ) -> Result<ReadCanisterSnapshotMetadataResult, SnapshotError> {
76        Ok(self
77            .snapshot_call::<_, (_,)>(
78                args.canister_id,
79                sender,
80                "read_canister_snapshot_metadata",
81                args,
82                limits,
83                &mut None,
84            )
85            .await?
86            .0)
87    }
88    pub async fn read_canister_snapshot_data(
89        &self,
90        sender: Principal,
91        args: ReadCanisterSnapshotDataArgs,
92        limits: CallLimits,
93    ) -> Result<ReadCanisterSnapshotDataResult, SnapshotError> {
94        Ok(self
95            .snapshot_call::<_, (_,)>(
96                args.canister_id,
97                sender,
98                "read_canister_snapshot_data",
99                args,
100                limits,
101                &mut None,
102            )
103            .await?
104            .0)
105    }
106    pub async fn upload_canister_snapshot_metadata(
107        &self,
108        sender: Principal,
109        args: UploadCanisterSnapshotMetadataArgs,
110        limits: CallLimits,
111    ) -> Result<UploadCanisterSnapshotMetadataResult, SnapshotError> {
112        Ok(self
113            .snapshot_call::<_, (_,)>(
114                args.canister_id,
115                sender,
116                "upload_canister_snapshot_metadata",
117                args,
118                limits,
119                &mut None,
120            )
121            .await?
122            .0)
123    }
124    pub async fn upload_canister_snapshot_data(
125        &self,
126        sender: Principal,
127        args: UploadCanisterSnapshotDataArgs,
128        limits: CallLimits,
129    ) -> Result<(), SnapshotError> {
130        self.snapshot_upload_data(sender, args, limits, &mut None)
131            .await
132    }
133
134    async fn snapshot_upload_data(
135        &self,
136        sender: Principal,
137        args: UploadCanisterSnapshotDataArgs,
138        limits: CallLimits,
139        ingress: &mut Option<RawMessageId>,
140    ) -> Result<(), SnapshotError> {
141        self.snapshot_call(
142            args.canister_id,
143            sender,
144            "upload_canister_snapshot_data",
145            args,
146            limits,
147            ingress,
148        )
149        .await
150    }
151}
152
153impl SnapshotTransfer {
154    /// Copy one chunk, or create the upload and return its ID before copying data.
155    /// Dropping this future leaves the attempted status and accepted ingress in self.
156    pub async fn advance_async(&mut self, client: &PocketIc) -> Result<(), SnapshotError> {
157        if matches!(self.status, TransferStatus::Complete) {
158            return Ok(());
159        }
160        if !matches!(self.status, TransferStatus::Ready) {
161            return Err(SnapshotError::Invalid(
162                "interrupted transfer requires explicit reconciliation",
163            ));
164        }
165        self.ingress = None;
166        if self.uploaded.is_none() {
167            self.status = TransferStatus::UploadingMetadata;
168            let (uploaded,): (UploadCanisterSnapshotMetadataResult,) = client
169                .snapshot_call(
170                    self.canister,
171                    self.sender,
172                    "upload_canister_snapshot_metadata",
173                    UploadCanisterSnapshotMetadataArgs {
174                        canister_id: self.canister,
175                        replace_snapshot: None,
176                        wasm_module_size: self.metadata.wasm_module_size,
177                        globals: self.metadata.globals.iter().flatten().cloned().collect(),
178                        wasm_memory_size: self.metadata.wasm_memory_size,
179                        stable_memory_size: self.metadata.stable_memory_size,
180                        certified_data: self.metadata.certified_data.clone(),
181                        global_timer: self.metadata.global_timer.clone(),
182                        on_low_wasm_memory_hook_status: self
183                            .metadata
184                            .on_low_wasm_memory_hook_status
185                            .clone(),
186                    },
187                    self.limits.call,
188                    &mut self.ingress,
189                )
190                .await?;
191            self.uploaded = Some(uploaded.snapshot_id);
192            self.status = TransferStatus::Ready;
193            return Ok(());
194        }
195        let sizes = [
196            self.metadata.wasm_module_size,
197            self.metadata.wasm_memory_size,
198            self.metadata.stable_memory_size,
199        ];
200        while self.region < 3 && self.offset == sizes[self.region] {
201            self.region += 1;
202            self.offset = 0;
203        }
204        let (kind, destination, expected) = if self.region < 3 {
205            let size = self
206                .limits
207                .chunk_bytes
208                .min(sizes[self.region] - self.offset);
209            let offset = self.offset;
210            match self.region {
211                0 => (
212                    SnapshotDataKind::WasmModule { offset, size },
213                    SnapshotDataOffset::WasmModule { offset },
214                    Some(size),
215                ),
216                1 => (
217                    SnapshotDataKind::WasmMemory { offset, size },
218                    SnapshotDataOffset::WasmMemory { offset },
219                    Some(size),
220                ),
221                _ => (
222                    SnapshotDataKind::StableMemory { offset, size },
223                    SnapshotDataOffset::StableMemory { offset },
224                    Some(size),
225                ),
226            }
227        } else if let Some(chunk) = self.metadata.wasm_chunk_store.get(self.region - 3) {
228            (
229                SnapshotDataKind::WasmChunk {
230                    hash: chunk.hash.clone(),
231                },
232                SnapshotDataOffset::WasmChunk,
233                None,
234            )
235        } else {
236            self.status = TransferStatus::Complete;
237            return Ok(());
238        };
239        self.status = TransferStatus::Reading(kind.clone());
240        let (data,): (ReadCanisterSnapshotDataResult,) = client
241            .snapshot_call(
242                self.canister,
243                self.sender,
244                "read_canister_snapshot_data",
245                ReadCanisterSnapshotDataArgs {
246                    canister_id: self.canister,
247                    snapshot_id: self.original.clone(),
248                    kind: kind.clone(),
249                },
250                self.limits.call,
251                &mut self.ingress,
252            )
253            .await?;
254        let bytes = data.chunk.len() as u64;
255        if bytes > self.limits.chunk_bytes || expected.is_some_and(|size| bytes != size) {
256            return Err(SnapshotError::Invalid(
257                "snapshot returned an incomplete or oversized chunk",
258            ));
259        }
260        if let SnapshotDataKind::WasmChunk { hash } = kind
261            && Sha256::digest(&data.chunk)[..] != hash
262        {
263            return Err(SnapshotError::Invalid("snapshot chunk hash mismatch"));
264        }
265        let transferred = self
266            .transferred
267            .checked_add(bytes)
268            .filter(|bytes| *bytes <= self.limits.total_bytes)
269            .ok_or(SnapshotError::Invalid("snapshot exceeds total byte budget"))?;
270        self.status = TransferStatus::Writing(destination.clone());
271        self.ingress = None;
272        client
273            .snapshot_upload_data(
274                self.sender,
275                UploadCanisterSnapshotDataArgs {
276                    canister_id: self.canister,
277                    snapshot_id: self.uploaded.clone().unwrap(),
278                    kind: destination,
279                    chunk: data.chunk,
280                },
281                self.limits.call,
282                &mut self.ingress,
283            )
284            .await?;
285        self.transferred = transferred;
286        if self.region < 3 {
287            self.offset += bytes;
288        } else {
289            self.region += 1;
290        }
291        self.status = TransferStatus::Ready;
292        Ok(())
293    }
294}
ic_testkit/pocket_ic/nonblocking/snapshot.rs

ic_testkit/pocket_ic/nonblocking/
snapshot.rs
