1use candid::{CandidType, Principal, decode_one, encode_args, utils::ArgumentEncoder};
2use pocket_ic::{PocketIc, RejectResponse};
3use serde::de::DeserializeOwned;
4use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
5
6use super::{CandidCallContext, CandidCallError, transport};
7
8#[derive(Clone, Copy)]
9struct CallContext<'a> {
10 operation: &'static str,
11 canister_id: Principal,
12 caller: Principal,
13 method: &'a str,
14}
15
16impl CallContext<'_> {
17 fn to_error_context(self) -> CandidCallContext {
18 CandidCallContext::new(self.operation, self.canister_id, self.caller, self.method)
19 }
20}
21
22pub trait CandidCallExt {
24 fn update_candid<T, A>(
26 &self,
27 canister_id: Principal,
28 method: &str,
29 args: A,
30 ) -> Result<T, CandidCallError>
31 where
32 T: CandidType + DeserializeOwned,
33 A: ArgumentEncoder,
34 {
35 self.update_candid_as(canister_id, Principal::anonymous(), method, args)
36 }
37
38 #[track_caller]
42 fn update_candid_or_panic<T, A>(&self, canister_id: Principal, method: &str, args: A) -> T
43 where
44 T: CandidType + DeserializeOwned,
45 A: ArgumentEncoder,
46 {
47 self.update_candid(canister_id, method, args)
48 .unwrap_or_else(|err| panic!("{err}"))
49 }
50
51 fn update_candid_as<T, A>(
53 &self,
54 canister_id: Principal,
55 caller: Principal,
56 method: &str,
57 args: A,
58 ) -> Result<T, CandidCallError>
59 where
60 T: CandidType + DeserializeOwned,
61 A: ArgumentEncoder;
62
63 #[track_caller]
67 fn update_candid_as_or_panic<T, A>(
68 &self,
69 canister_id: Principal,
70 caller: Principal,
71 method: &str,
72 args: A,
73 ) -> T
74 where
75 T: CandidType + DeserializeOwned,
76 A: ArgumentEncoder,
77 {
78 self.update_candid_as(canister_id, caller, method, args)
79 .unwrap_or_else(|err| panic!("{err}"))
80 }
81
82 fn query_candid<T, A>(
84 &self,
85 canister_id: Principal,
86 method: &str,
87 args: A,
88 ) -> Result<T, CandidCallError>
89 where
90 T: CandidType + DeserializeOwned,
91 A: ArgumentEncoder,
92 {
93 self.query_candid_as(canister_id, Principal::anonymous(), method, args)
94 }
95
96 #[track_caller]
100 fn query_candid_or_panic<T, A>(&self, canister_id: Principal, method: &str, args: A) -> T
101 where
102 T: CandidType + DeserializeOwned,
103 A: ArgumentEncoder,
104 {
105 self.query_candid(canister_id, method, args)
106 .unwrap_or_else(|err| panic!("{err}"))
107 }
108
109 fn query_candid_as<T, A>(
111 &self,
112 canister_id: Principal,
113 caller: Principal,
114 method: &str,
115 args: A,
116 ) -> Result<T, CandidCallError>
117 where
118 T: CandidType + DeserializeOwned,
119 A: ArgumentEncoder;
120
121 #[track_caller]
125 fn query_candid_as_or_panic<T, A>(
126 &self,
127 canister_id: Principal,
128 caller: Principal,
129 method: &str,
130 args: A,
131 ) -> T
132 where
133 T: CandidType + DeserializeOwned,
134 A: ArgumentEncoder,
135 {
136 self.query_candid_as(canister_id, caller, method, args)
137 .unwrap_or_else(|err| panic!("{err}"))
138 }
139}
140
141impl CandidCallExt for PocketIc {
142 fn update_candid_as<T, A>(
144 &self,
145 canister_id: Principal,
146 caller: Principal,
147 method: &str,
148 args: A,
149 ) -> Result<T, CandidCallError>
150 where
151 T: CandidType + DeserializeOwned,
152 A: ArgumentEncoder,
153 {
154 let context = CallContext {
155 operation: "update_call",
156 canister_id,
157 caller,
158 method,
159 };
160 let bytes = encode_call_args(args, context)?;
161 let result = run_raw_call(context, || {
162 Self::update_call(self, canister_id, caller, method, bytes)
163 })?;
164
165 decode_call_result(&result, context)
166 }
167
168 fn query_candid_as<T, A>(
170 &self,
171 canister_id: Principal,
172 caller: Principal,
173 method: &str,
174 args: A,
175 ) -> Result<T, CandidCallError>
176 where
177 T: CandidType + DeserializeOwned,
178 A: ArgumentEncoder,
179 {
180 let context = CallContext {
181 operation: "query_call",
182 canister_id,
183 caller,
184 method,
185 };
186 let bytes = encode_call_args(args, context)?;
187 let result = run_raw_call(context, || {
188 Self::query_call(self, canister_id, caller, method, bytes)
189 })?;
190
191 decode_call_result(&result, context)
192 }
193}
194
195fn encode_call_args<A>(args: A, context: CallContext<'_>) -> Result<Vec<u8>, CandidCallError>
196where
197 A: ArgumentEncoder,
198{
199 encode_args(args).map_err(|err| CandidCallError::encode(context.to_error_context(), err))
200}
201
202fn decode_call_result<T>(result: &[u8], context: CallContext<'_>) -> Result<T, CandidCallError>
203where
204 T: CandidType + DeserializeOwned,
205{
206 decode_one(result)
207 .map_err(|err| CandidCallError::decode(context.to_error_context(), result.len(), err))
208}
209
210fn run_raw_call<F>(context: CallContext<'_>, call: F) -> Result<Vec<u8>, CandidCallError>
211where
212 F: FnOnce() -> Result<Vec<u8>, RejectResponse>,
213{
214 match catch_unwind(AssertUnwindSafe(call)) {
215 Ok(Ok(bytes)) => Ok(bytes),
216 Ok(Err(response)) => Err(CandidCallError::canister_reject(
217 context.to_error_context(),
218 response,
219 )),
220 Err(payload) if transport::panic_is_dead_instance_transport(payload.as_ref()) => {
221 Err(CandidCallError::transport(
222 context.to_error_context(),
223 transport::panic_payload_to_string(payload.as_ref()),
224 ))
225 }
226 Err(payload) => resume_unwind(payload),
227 }
228}
229
230#[cfg(test)]
231mod tests {
232 use candid::Principal;
233
234 use crate::pic::CandidCallErrorKind;
235
236 use super::{CallContext, decode_call_result, run_raw_call};
237
238 #[test]
239 fn decode_error_includes_call_context() {
240 let context = CallContext {
241 operation: "query_call",
242 canister_id: Principal::anonymous(),
243 caller: Principal::management_canister(),
244 method: "get",
245 };
246
247 let err = decode_call_result::<u64>(&[0xde, 0xad], context).expect_err("decode fails");
248
249 assert!(err.message().contains("candid decode_one failed"));
250 assert!(err.message().contains("operation=query_call"));
251 assert!(err.message().contains("method=get"));
252 assert!(err.message().contains("bytes=2"));
253 assert_eq!(err.kind(), CandidCallErrorKind::Decode);
254 assert_eq!(err.context().expect("decode error context").method(), "get");
255 }
256
257 #[test]
258 fn dead_instance_panic_is_classified_as_transport() {
259 let context = CallContext {
260 operation: "query_call",
261 canister_id: Principal::anonymous(),
262 caller: Principal::management_canister(),
263 method: "get",
264 };
265 let error = run_raw_call(context, || -> Result<Vec<u8>, pocket_ic::RejectResponse> {
266 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\" }}) }}");
267 })
268 .unwrap_err();
269
270 assert_eq!(error.kind(), crate::pic::CandidCallErrorKind::Transport);
271 assert!(error.reject_response().is_none());
272 }
273
274 #[test]
275 fn unrelated_panics_resume_with_the_original_payload() {
276 for message in [
277 "application worker channel closed",
278 "fixture expected ConnectionRefused but observed another result",
279 ] {
280 let context = CallContext {
281 operation: "query_call",
282 canister_id: Principal::anonymous(),
283 caller: Principal::management_canister(),
284 method: "get",
285 };
286 let payload = message.to_owned();
287 let original_allocation = payload.as_ptr() as usize;
288 let result = std::panic::catch_unwind(|| {
289 run_raw_call(context, || -> Result<Vec<u8>, pocket_ic::RejectResponse> {
290 std::panic::panic_any(payload);
291 })
292 });
293 let resumed = result
294 .expect_err("application panic must resume")
295 .downcast::<String>()
296 .expect("original panic type must be preserved");
297 assert_eq!(resumed.as_str(), message);
298 assert_eq!(resumed.as_ptr() as usize, original_allocation);
299 }
300 }
301}