1use candid::{
2 CandidType, Principal, decode_args, decode_one, encode_args,
3 utils::{ArgumentDecoder, ArgumentEncoder},
4};
5use pocket_ic::{PocketIc, RejectResponse};
6use serde::de::DeserializeOwned;
7use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
8
9use super::{CandidCallContext, CandidCallError, transport};
10
11#[derive(Clone, Copy)]
12struct CallContext<'a> {
13 operation: &'static str,
14 canister_id: Principal,
15 caller: Principal,
16 method: &'a str,
17}
18
19impl CallContext<'_> {
20 fn to_error_context(self) -> CandidCallContext {
21 CandidCallContext::new(self.operation, self.canister_id, self.caller, self.method)
22 }
23}
24
25pub trait CandidCallExt {
27 fn update_candid<T, A>(
29 &self,
30 canister_id: Principal,
31 method: &str,
32 args: A,
33 ) -> Result<T, CandidCallError>
34 where
35 T: CandidType + DeserializeOwned,
36 A: ArgumentEncoder,
37 {
38 self.update_candid_as(canister_id, Principal::anonymous(), method, args)
39 }
40
41 fn update_candid_as<T, A>(
43 &self,
44 canister_id: Principal,
45 caller: Principal,
46 method: &str,
47 args: A,
48 ) -> Result<T, CandidCallError>
49 where
50 T: CandidType + DeserializeOwned,
51 A: ArgumentEncoder;
52
53 fn query_candid<T, A>(
55 &self,
56 canister_id: Principal,
57 method: &str,
58 args: A,
59 ) -> Result<T, CandidCallError>
60 where
61 T: CandidType + DeserializeOwned,
62 A: ArgumentEncoder,
63 {
64 self.query_candid_as(canister_id, Principal::anonymous(), method, args)
65 }
66
67 fn query_candid_as<T, A>(
69 &self,
70 canister_id: Principal,
71 caller: Principal,
72 method: &str,
73 args: A,
74 ) -> Result<T, CandidCallError>
75 where
76 T: CandidType + DeserializeOwned,
77 A: ArgumentEncoder;
78}
79
80impl CandidCallExt for PocketIc {
81 fn update_candid_as<T, A>(
83 &self,
84 canister_id: Principal,
85 caller: Principal,
86 method: &str,
87 args: A,
88 ) -> Result<T, CandidCallError>
89 where
90 T: CandidType + DeserializeOwned,
91 A: ArgumentEncoder,
92 {
93 let context = CallContext {
94 operation: "update_call",
95 canister_id,
96 caller,
97 method,
98 };
99 call_candid(
100 context,
101 args,
102 |bytes| Self::update_call(self, canister_id, caller, method, bytes),
103 |bytes| decode_one(bytes),
104 )
105 }
106
107 fn query_candid_as<T, A>(
109 &self,
110 canister_id: Principal,
111 caller: Principal,
112 method: &str,
113 args: A,
114 ) -> Result<T, CandidCallError>
115 where
116 T: CandidType + DeserializeOwned,
117 A: ArgumentEncoder,
118 {
119 let context = CallContext {
120 operation: "query_call",
121 canister_id,
122 caller,
123 method,
124 };
125 call_candid(
126 context,
127 args,
128 |bytes| Self::query_call(self, canister_id, caller, method, bytes),
129 |bytes| decode_one(bytes),
130 )
131 }
132}
133
134pub fn update_candid_args<T, A>(
141 pic: &PocketIc,
142 canister_id: Principal,
143 caller: Principal,
144 method: &str,
145 args: A,
146) -> Result<T, CandidCallError>
147where
148 T: for<'de> ArgumentDecoder<'de>,
149 A: ArgumentEncoder,
150{
151 let context = CallContext {
152 operation: "update_call",
153 canister_id,
154 caller,
155 method,
156 };
157 call_candid(
158 context,
159 args,
160 |bytes| pic.update_call(canister_id, caller, method, bytes),
161 |bytes| decode_args(bytes),
162 )
163}
164
165pub fn query_candid_args<T, A>(
172 pic: &PocketIc,
173 canister_id: Principal,
174 caller: Principal,
175 method: &str,
176 args: A,
177) -> Result<T, CandidCallError>
178where
179 T: for<'de> ArgumentDecoder<'de>,
180 A: ArgumentEncoder,
181{
182 let context = CallContext {
183 operation: "query_call",
184 canister_id,
185 caller,
186 method,
187 };
188 call_candid(
189 context,
190 args,
191 |bytes| pic.query_call(canister_id, caller, method, bytes),
192 |bytes| decode_args(bytes),
193 )
194}
195
196fn call_candid<T, A, F, D>(
197 context: CallContext<'_>,
198 args: A,
199 call: F,
200 decode: D,
201) -> Result<T, CandidCallError>
202where
203 A: ArgumentEncoder,
204 F: FnOnce(Vec<u8>) -> Result<Vec<u8>, RejectResponse>,
205 D: FnOnce(&[u8]) -> Result<T, candid::Error>,
206{
207 let bytes = encode_args(args).map_err(|err| CandidCallError::Encode {
208 context: Box::new(context.to_error_context()),
209 message: err.to_string(),
210 })?;
211 let result = run_raw_call(context, || call(bytes))?;
212 decode(&result).map_err(|err| CandidCallError::Decode {
213 context: Box::new(context.to_error_context()),
214 bytes: result.len(),
215 message: err.to_string(),
216 })
217}
218
219fn run_raw_call<F>(context: CallContext<'_>, call: F) -> Result<Vec<u8>, CandidCallError>
220where
221 F: FnOnce() -> Result<Vec<u8>, RejectResponse>,
222{
223 match catch_unwind(AssertUnwindSafe(call)) {
224 Ok(Ok(bytes)) => Ok(bytes),
225 Ok(Err(response)) => Err(CandidCallError::CanisterReject {
226 context: Box::new(context.to_error_context()),
227 response: Box::new(response),
228 }),
229 Err(payload) if transport::panic_is_dead_instance_transport(payload.as_ref()) => {
230 Err(CandidCallError::Transport {
231 context: Box::new(context.to_error_context()),
232 message: transport::panic_payload_to_string(payload.as_ref()),
233 })
234 }
235 Err(payload) => resume_unwind(payload),
236 }
237}
238
239#[cfg(test)]
240mod tests {
241 use candid::Principal;
242
243 use super::{CallContext, CandidCallError, call_candid, run_raw_call};
244
245 #[test]
246 fn decode_error_includes_call_context() {
247 let context = CallContext {
248 operation: "query_call",
249 canister_id: Principal::anonymous(),
250 caller: Principal::management_canister(),
251 method: "get",
252 };
253
254 let err = call_candid(
255 context,
256 (),
257 |_| Ok(vec![0xde, 0xad]),
258 |bytes| candid::decode_one::<u64>(bytes),
259 )
260 .expect_err("decode fails");
261
262 assert!(err.to_string().contains("candid decode failed"));
263 assert!(err.to_string().contains("operation=query_call"));
264 assert!(err.to_string().contains("method=get"));
265 assert!(err.to_string().contains("bytes=2"));
266 let CandidCallError::Decode { context, bytes, .. } = err else {
267 panic!("expected decode failure");
268 };
269 assert_eq!(bytes, 2);
270 assert_eq!(context.method, "get");
271 }
272
273 #[test]
274 fn reply_arguments_preserve_top_level_arity_and_wide_values() {
275 let context = CallContext {
276 operation: "update_call",
277 canister_id: Principal::anonymous(),
278 caller: Principal::management_canister(),
279 method: "measure",
280 };
281 let values = (u128::MAX, u128::from(u64::MAX) + 1);
282 let bytes = candid::encode_args(values).unwrap();
283 assert_eq!(candid::decode_args::<(u128, u128)>(&bytes).unwrap(), values);
284 assert!(candid::decode_one::<(u128, u128)>(&bytes).is_err());
285 let tuple = candid::encode_one(values).unwrap();
286 assert!(candid::decode_args::<(u128, u128)>(&tuple).is_err());
287 assert_eq!(candid::decode_one::<(u128, u128)>(&tuple).unwrap(), values);
288 let no_replies = candid::encode_args(()).unwrap();
289 candid::decode_args::<()>(&no_replies).unwrap();
290 let one_reply = candid::encode_args((7_u64,)).unwrap();
291 assert_eq!(candid::decode_args::<(u64,)>(&one_reply).unwrap(), (7,));
292 let result: (u64, bool, String) = call_candid(
293 context,
294 (42_u64, false),
295 |args| {
296 assert_eq!(
297 candid::decode_args::<(u64, bool)>(&args).unwrap(),
298 (42, false)
299 );
300 Ok(candid::encode_args((9_u64, true, "ready")).unwrap())
301 },
302 |bytes| candid::decode_args(bytes),
303 )
304 .unwrap();
305 assert_eq!(result, (9, true, "ready".into()));
306 }
307
308 #[test]
309 fn invalid_reply_arguments_keep_decode_context() {
310 let context = CallContext {
311 operation: "query_call",
312 canister_id: Principal::anonymous(),
313 caller: Principal::management_canister(),
314 method: "measure",
315 };
316 for bytes in [vec![0xde, 0xad], candid::encode_args((7_u64,)).unwrap()] {
317 let length = bytes.len();
318 let err = call_candid(
319 context,
320 (),
321 |_| Ok(bytes),
322 |bytes| candid::decode_args::<(u64, bool)>(bytes),
323 )
324 .unwrap_err();
325 assert!(err.to_string().contains(&format!("bytes={length}")));
326 let CandidCallError::Decode {
327 context: observed,
328 bytes,
329 ..
330 } = err
331 else {
332 panic!("expected decode failure");
333 };
334 assert_eq!(bytes, length);
335 assert_eq!(observed.operation, "query_call");
336 assert_eq!(observed.canister_id, context.canister_id);
337 assert_eq!(observed.caller, context.caller);
338 assert_eq!(observed.method, "measure");
339 }
340 }
341
342 #[test]
343 fn dead_instance_panic_is_classified_as_transport() {
344 let context = CallContext {
345 operation: "query_call",
346 canister_id: Principal::anonymous(),
347 caller: Principal::management_canister(),
348 method: "get",
349 };
350 let error = run_raw_call(context, || -> Result<Vec<u8>, pocket_ic::RejectResponse> {
351 panic!("HTTP failure: reqwest::Error {{ kind: Request, url: \"http://127.0.0.1:1234/instances/0/update/execute_ingress_message\", source: ConnectError(\"tcp connect error\", Os {{ code: 111, kind: ConnectionRefused, message: \"Connection refused\" }}) }}");
352 })
353 .unwrap_err();
354
355 assert!(matches!(error, CandidCallError::Transport { .. }));
356 }
357
358 #[test]
359 fn unrelated_panics_resume_with_the_original_payload() {
360 for message in [
361 "application worker channel closed",
362 "fixture expected ConnectionRefused but observed another result",
363 ] {
364 let context = CallContext {
365 operation: "query_call",
366 canister_id: Principal::anonymous(),
367 caller: Principal::management_canister(),
368 method: "get",
369 };
370 let payload = message.to_owned();
371 let original_allocation = payload.as_ptr() as usize;
372 let result = std::panic::catch_unwind(|| {
373 run_raw_call(context, || -> Result<Vec<u8>, pocket_ic::RejectResponse> {
374 std::panic::panic_any(payload);
375 })
376 });
377 let resumed = result
378 .expect_err("application panic must resume")
379 .downcast::<String>()
380 .expect("original panic type must be preserved");
381 assert_eq!(resumed.as_str(), message);
382 assert_eq!(resumed.as_ptr() as usize, original_allocation);
383 }
384 }
385}