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 fn update_candid_as<T, A>(
40 &self,
41 canister_id: Principal,
42 caller: Principal,
43 method: &str,
44 args: A,
45 ) -> Result<T, CandidCallError>
46 where
47 T: CandidType + DeserializeOwned,
48 A: ArgumentEncoder;
49
50 fn query_candid<T, A>(
52 &self,
53 canister_id: Principal,
54 method: &str,
55 args: A,
56 ) -> Result<T, CandidCallError>
57 where
58 T: CandidType + DeserializeOwned,
59 A: ArgumentEncoder,
60 {
61 self.query_candid_as(canister_id, Principal::anonymous(), method, args)
62 }
63
64 fn query_candid_as<T, A>(
66 &self,
67 canister_id: Principal,
68 caller: Principal,
69 method: &str,
70 args: A,
71 ) -> Result<T, CandidCallError>
72 where
73 T: CandidType + DeserializeOwned,
74 A: ArgumentEncoder;
75}
76
77impl CandidCallExt for PocketIc {
78 fn update_candid_as<T, A>(
80 &self,
81 canister_id: Principal,
82 caller: Principal,
83 method: &str,
84 args: A,
85 ) -> Result<T, CandidCallError>
86 where
87 T: CandidType + DeserializeOwned,
88 A: ArgumentEncoder,
89 {
90 let context = CallContext {
91 operation: "update_call",
92 canister_id,
93 caller,
94 method,
95 };
96 let bytes = encode_call_args(args, context)?;
97 let result = run_raw_call(context, || {
98 Self::update_call(self, canister_id, caller, method, bytes)
99 })?;
100
101 decode_call_result(&result, context)
102 }
103
104 fn query_candid_as<T, A>(
106 &self,
107 canister_id: Principal,
108 caller: Principal,
109 method: &str,
110 args: A,
111 ) -> Result<T, CandidCallError>
112 where
113 T: CandidType + DeserializeOwned,
114 A: ArgumentEncoder,
115 {
116 let context = CallContext {
117 operation: "query_call",
118 canister_id,
119 caller,
120 method,
121 };
122 let bytes = encode_call_args(args, context)?;
123 let result = run_raw_call(context, || {
124 Self::query_call(self, canister_id, caller, method, bytes)
125 })?;
126
127 decode_call_result(&result, context)
128 }
129}
130
131fn encode_call_args<A>(args: A, context: CallContext<'_>) -> Result<Vec<u8>, CandidCallError>
132where
133 A: ArgumentEncoder,
134{
135 encode_args(args).map_err(|err| CandidCallError::encode(context.to_error_context(), err))
136}
137
138fn decode_call_result<T>(result: &[u8], context: CallContext<'_>) -> Result<T, CandidCallError>
139where
140 T: CandidType + DeserializeOwned,
141{
142 decode_one(result)
143 .map_err(|err| CandidCallError::decode(context.to_error_context(), result.len(), err))
144}
145
146fn run_raw_call<F>(context: CallContext<'_>, call: F) -> Result<Vec<u8>, CandidCallError>
147where
148 F: FnOnce() -> Result<Vec<u8>, RejectResponse>,
149{
150 match catch_unwind(AssertUnwindSafe(call)) {
151 Ok(Ok(bytes)) => Ok(bytes),
152 Ok(Err(response)) => Err(CandidCallError::canister_reject(
153 context.to_error_context(),
154 response,
155 )),
156 Err(payload) if transport::panic_is_dead_instance_transport(payload.as_ref()) => {
157 Err(CandidCallError::transport(
158 context.to_error_context(),
159 transport::panic_payload_to_string(payload.as_ref()),
160 ))
161 }
162 Err(payload) => resume_unwind(payload),
163 }
164}
165
166#[cfg(test)]
167mod tests {
168 use candid::Principal;
169
170 use crate::pic::CandidCallErrorKind;
171
172 use super::{CallContext, decode_call_result, run_raw_call};
173
174 #[test]
175 fn decode_error_includes_call_context() {
176 let context = CallContext {
177 operation: "query_call",
178 canister_id: Principal::anonymous(),
179 caller: Principal::management_canister(),
180 method: "get",
181 };
182
183 let err = decode_call_result::<u64>(&[0xde, 0xad], context).expect_err("decode fails");
184
185 assert!(err.message.contains("candid decode_one failed"));
186 assert!(err.message.contains("operation=query_call"));
187 assert!(err.message.contains("method=get"));
188 assert!(err.message.contains("bytes=2"));
189 assert_eq!(err.kind, CandidCallErrorKind::Decode);
190 assert_eq!(
191 err.context.as_deref().expect("decode error context").method,
192 "get"
193 );
194 }
195
196 #[test]
197 fn dead_instance_panic_is_classified_as_transport() {
198 let context = CallContext {
199 operation: "query_call",
200 canister_id: Principal::anonymous(),
201 caller: Principal::management_canister(),
202 method: "get",
203 };
204 let error = run_raw_call(context, || -> Result<Vec<u8>, pocket_ic::RejectResponse> {
205 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\" }}) }}");
206 })
207 .unwrap_err();
208
209 assert_eq!(error.kind, crate::pic::CandidCallErrorKind::Transport);
210 assert!(error.reject_response.is_none());
211 }
212
213 #[test]
214 fn unrelated_panics_resume_with_the_original_payload() {
215 for message in [
216 "application worker channel closed",
217 "fixture expected ConnectionRefused but observed another result",
218 ] {
219 let context = CallContext {
220 operation: "query_call",
221 canister_id: Principal::anonymous(),
222 caller: Principal::management_canister(),
223 method: "get",
224 };
225 let payload = message.to_owned();
226 let original_allocation = payload.as_ptr() as usize;
227 let result = std::panic::catch_unwind(|| {
228 run_raw_call(context, || -> Result<Vec<u8>, pocket_ic::RejectResponse> {
229 std::panic::panic_any(payload);
230 })
231 });
232 let resumed = result
233 .expect_err("application panic must resume")
234 .downcast::<String>()
235 .expect("original panic type must be preserved");
236 assert_eq!(resumed.as_str(), message);
237 assert_eq!(resumed.as_ptr() as usize, original_allocation);
238 }
239 }
240}