ic_test_state_machine_client/
lib.rs

1use crate::management_canister::{
2    CanisterId, CanisterIdRecord, CanisterInstallMode, CreateCanisterArgument, InstallCodeArgument,
3};
4use candid::utils::{ArgumentDecoder, ArgumentEncoder};
5use candid::{decode_args, encode_args, Principal};
6use ciborium::de::from_reader;
7use serde::de::DeserializeOwned;
8use serde::{Deserialize, Serialize};
9use serde_bytes::ByteBuf;
10use std::cell::RefCell;
11use std::fmt;
12use std::io::{Read, Write};
13use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
14use std::time::{Duration, SystemTime};
15
16mod management_canister;
17
18pub use management_canister::CanisterSettings;
19
20#[derive(Serialize, Deserialize)]
21pub enum Request {
22    RootKey,
23    Time,
24    SetTime(SystemTime),
25    AdvanceTime(Duration),
26    CanisterUpdateCall(CanisterCall),
27    CanisterQueryCall(CanisterCall),
28    CanisterExists(RawCanisterId),
29    CyclesBalance(RawCanisterId),
30    AddCycles(AddCyclesArg),
31    SetStableMemory(SetStableMemoryArg),
32    ReadStableMemory(RawCanisterId),
33    Tick,
34    RunUntilCompletion(RunUntilCompletionArg),
35    VerifyCanisterSig(VerifyCanisterSigArg),
36}
37
38#[derive(Serialize, Deserialize)]
39pub struct VerifyCanisterSigArg {
40    pub msg: Vec<u8>,
41    pub sig: Vec<u8>,
42    pub pubkey: Vec<u8>,
43    pub root_pubkey: Vec<u8>,
44}
45
46#[derive(Serialize, Deserialize)]
47pub struct RunUntilCompletionArg {
48    // max_ticks until completion must be reached
49    pub max_ticks: u64,
50}
51
52#[derive(Serialize, Deserialize)]
53pub struct AddCyclesArg {
54    // raw bytes of the principal
55    pub canister_id: Vec<u8>,
56    pub amount: u128,
57}
58
59#[derive(Serialize, Deserialize)]
60pub struct SetStableMemoryArg {
61    // raw bytes of the principal
62    pub canister_id: Vec<u8>,
63    pub data: ByteBuf,
64}
65
66#[derive(Serialize, Deserialize)]
67pub struct RawCanisterId {
68    // raw bytes of the principal
69    pub canister_id: Vec<u8>,
70}
71
72impl From<Principal> for RawCanisterId {
73    fn from(principal: Principal) -> Self {
74        Self {
75            canister_id: principal.as_slice().to_vec(),
76        }
77    }
78}
79
80#[derive(Serialize, Deserialize)]
81pub struct CanisterCall {
82    pub sender: Vec<u8>,
83    pub canister_id: Vec<u8>,
84    pub method: String,
85    pub arg: Vec<u8>,
86}
87
88pub struct StateMachine {
89    proc: Child,
90    child_in: RefCell<ChildStdin>,
91    child_out: RefCell<ChildStdout>,
92}
93
94impl StateMachine {
95    pub fn new(binary_path: &str, debug: bool) -> Self {
96        let mut command = Command::new(binary_path);
97        command
98            .env("LOG_TO_STDERR", "1")
99            .stdin(Stdio::piped())
100            .stdout(Stdio::piped());
101
102        if debug {
103            command.arg("--debug");
104        }
105
106        let mut child = command.spawn().unwrap_or_else(|err| {
107            panic!(
108                "failed to start test state machine at path {}: {:?}",
109                binary_path, err
110            )
111        });
112
113        let child_in = child.stdin.take().unwrap();
114        let child_out = child.stdout.take().unwrap();
115        Self {
116            proc: child,
117            child_in: RefCell::new(child_in),
118            child_out: RefCell::new(child_out),
119        }
120    }
121
122    pub fn update_call(
123        &self,
124        canister_id: Principal,
125        sender: Principal,
126        method: &str,
127        arg: Vec<u8>,
128    ) -> Result<WasmResult, UserError> {
129        self.call_state_machine(Request::CanisterUpdateCall(CanisterCall {
130            sender: sender.as_slice().to_vec(),
131            canister_id: canister_id.as_slice().to_vec(),
132            method: method.to_string(),
133            arg,
134        }))
135    }
136
137    pub fn query_call(
138        &self,
139        canister_id: Principal,
140        sender: Principal,
141        method: &str,
142        arg: Vec<u8>,
143    ) -> Result<WasmResult, UserError> {
144        self.call_state_machine(Request::CanisterQueryCall(CanisterCall {
145            sender: sender.as_slice().to_vec(),
146            canister_id: canister_id.as_slice().to_vec(),
147            method: method.to_string(),
148            arg,
149        }))
150    }
151
152    pub fn root_key(&self) -> Vec<u8> {
153        self.call_state_machine(Request::RootKey)
154    }
155
156    pub fn create_canister(&self, sender: Option<Principal>) -> CanisterId {
157        let CanisterIdRecord { canister_id } = call_candid_as(
158            self,
159            Principal::management_canister(),
160            sender.unwrap_or(Principal::anonymous()),
161            "create_canister",
162            (CreateCanisterArgument { settings: None },),
163        )
164        .map(|(x,)| x)
165        .unwrap();
166        canister_id
167    }
168
169    pub fn create_canister_with_settings(
170        &self,
171        settings: Option<CanisterSettings>,
172        sender: Option<Principal>,
173    ) -> CanisterId {
174        let CanisterIdRecord { canister_id } = call_candid_as(
175            self,
176            Principal::management_canister(),
177            sender.unwrap_or(Principal::anonymous()),
178            "create_canister",
179            (CreateCanisterArgument { settings },),
180        )
181        .map(|(x,)| x)
182        .unwrap();
183        canister_id
184    }
185
186    pub fn install_canister(
187        &self,
188        canister_id: CanisterId,
189        wasm_module: Vec<u8>,
190        arg: Vec<u8>,
191        sender: Option<Principal>,
192    ) {
193        call_candid_as::<(InstallCodeArgument,), ()>(
194            self,
195            Principal::management_canister(),
196            sender.unwrap_or(Principal::anonymous()),
197            "install_code",
198            (InstallCodeArgument {
199                mode: CanisterInstallMode::Install,
200                canister_id,
201                wasm_module,
202                arg,
203            },),
204        )
205        .unwrap();
206    }
207
208    pub fn upgrade_canister(
209        &self,
210        canister_id: CanisterId,
211        wasm_module: Vec<u8>,
212        arg: Vec<u8>,
213        sender: Option<Principal>,
214    ) -> Result<(), CallError> {
215        call_candid_as::<(InstallCodeArgument,), ()>(
216            self,
217            Principal::management_canister(),
218            sender.unwrap_or(Principal::anonymous()),
219            "install_code",
220            (InstallCodeArgument {
221                mode: CanisterInstallMode::Upgrade,
222                canister_id,
223                wasm_module,
224                arg,
225            },),
226        )
227    }
228
229    pub fn reinstall_canister(
230        &self,
231        canister_id: CanisterId,
232        wasm_module: Vec<u8>,
233        arg: Vec<u8>,
234        sender: Option<Principal>,
235    ) -> Result<(), CallError> {
236        call_candid_as::<(InstallCodeArgument,), ()>(
237            self,
238            Principal::management_canister(),
239            sender.unwrap_or(Principal::anonymous()),
240            "install_code",
241            (InstallCodeArgument {
242                mode: CanisterInstallMode::Reinstall,
243                canister_id,
244                wasm_module,
245                arg,
246            },),
247        )
248    }
249
250    pub fn start_canister(
251        &self,
252        canister_id: CanisterId,
253        sender: Option<Principal>,
254    ) -> Result<(), CallError> {
255        call_candid_as::<(CanisterIdRecord,), ()>(
256            self,
257            Principal::management_canister(),
258            sender.unwrap_or(Principal::anonymous()),
259            "start_canister",
260            (CanisterIdRecord { canister_id },),
261        )
262    }
263
264    pub fn stop_canister(
265        &self,
266        canister_id: CanisterId,
267        sender: Option<Principal>,
268    ) -> Result<(), CallError> {
269        call_candid_as::<(CanisterIdRecord,), ()>(
270            self,
271            Principal::management_canister(),
272            sender.unwrap_or(Principal::anonymous()),
273            "stop_canister",
274            (CanisterIdRecord { canister_id },),
275        )
276    }
277
278    pub fn delete_canister(
279        &self,
280        canister_id: CanisterId,
281        sender: Option<Principal>,
282    ) -> Result<(), CallError> {
283        call_candid_as::<(CanisterIdRecord,), ()>(
284            self,
285            Principal::management_canister(),
286            sender.unwrap_or(Principal::anonymous()),
287            "delete_canister",
288            (CanisterIdRecord { canister_id },),
289        )
290    }
291
292    pub fn canister_exists(&self, canister_id: Principal) -> bool {
293        self.call_state_machine(Request::CanisterExists(RawCanisterId::from(canister_id)))
294    }
295
296    pub fn time(&self) -> SystemTime {
297        self.call_state_machine(Request::Time)
298    }
299
300    pub fn set_time(&self, time: SystemTime) {
301        self.call_state_machine(Request::SetTime(time))
302    }
303
304    pub fn advance_time(&self, duration: Duration) {
305        self.call_state_machine(Request::AdvanceTime(duration))
306    }
307
308    pub fn tick(&self) {
309        self.call_state_machine(Request::Tick)
310    }
311
312    pub fn run_until_completion(&self, max_ticks: u64) {
313        self.call_state_machine(Request::RunUntilCompletion(RunUntilCompletionArg {
314            max_ticks,
315        }))
316    }
317
318    pub fn stable_memory(&self, canister_id: Principal) -> Vec<u8> {
319        self.call_state_machine(Request::ReadStableMemory(RawCanisterId::from(canister_id)))
320    }
321
322    pub fn set_stable_memory(&self, canister_id: Principal, data: ByteBuf) {
323        self.call_state_machine(Request::SetStableMemory(SetStableMemoryArg {
324            canister_id: canister_id.as_slice().to_vec(),
325            data,
326        }))
327    }
328
329    pub fn cycle_balance(&self, canister_id: Principal) -> u128 {
330        self.call_state_machine(Request::CyclesBalance(RawCanisterId::from(canister_id)))
331    }
332
333    pub fn add_cycles(&self, canister_id: Principal, amount: u128) -> u128 {
334        self.call_state_machine(Request::AddCycles(AddCyclesArg {
335            canister_id: canister_id.as_slice().to_vec(),
336            amount,
337        }))
338    }
339
340    /// Verifies a canister signature. Returns Ok(()) if the signature is valid.
341    /// On error, returns a string describing the error.
342    pub fn verify_canister_signature(
343        &self,
344        msg: Vec<u8>,
345        sig: Vec<u8>,
346        pubkey: Vec<u8>,
347        root_pubkey: Vec<u8>,
348    ) -> Result<(), String> {
349        self.call_state_machine(Request::VerifyCanisterSig(VerifyCanisterSigArg {
350            msg,
351            sig,
352            pubkey,
353            root_pubkey,
354        }))
355    }
356
357    fn call_state_machine<T: DeserializeOwned>(&self, request: Request) -> T {
358        self.send_request(request);
359        self.read_response()
360    }
361
362    fn send_request(&self, request: Request) {
363        let mut cbor = vec![];
364        ciborium::ser::into_writer(&request, &mut cbor).expect("failed to serialize request");
365        let mut child_in = self.child_in.borrow_mut();
366        child_in
367            .write_all(&(cbor.len() as u64).to_le_bytes())
368            .expect("failed to send request length");
369        child_in
370            .write_all(cbor.as_slice())
371            .expect("failed to send request data");
372        child_in.flush().expect("failed to flush child stdin");
373    }
374
375    fn read_response<T: DeserializeOwned>(&self) -> T {
376        let vec = self.read_bytes(8);
377        let size = usize::from_le_bytes(TryFrom::try_from(vec).expect("failed to read data size"));
378        from_reader(&self.read_bytes(size)[..]).expect("failed to deserialize response")
379    }
380
381    fn read_bytes(&self, num_bytes: usize) -> Vec<u8> {
382        let mut buf = vec![0u8; num_bytes];
383        self.child_out
384            .borrow_mut()
385            .read_exact(&mut buf)
386            .expect("failed to read from child_stdout");
387        buf
388    }
389}
390
391impl Drop for StateMachine {
392    fn drop(&mut self) {
393        self.proc
394            .kill()
395            .expect("failed to kill state machine process")
396    }
397}
398
399/// Call a canister candid query method, anonymous.
400pub fn query_candid<Input, Output>(
401    env: &StateMachine,
402    canister_id: Principal,
403    method: &str,
404    input: Input,
405) -> Result<Output, CallError>
406where
407    Input: ArgumentEncoder,
408    Output: for<'a> ArgumentDecoder<'a>,
409{
410    query_candid_as(env, canister_id, Principal::anonymous(), method, input)
411}
412
413/// Call a canister candid query method, authenticated.
414pub fn query_candid_as<Input, Output>(
415    env: &StateMachine,
416    canister_id: Principal,
417    sender: Principal,
418    method: &str,
419    input: Input,
420) -> Result<Output, CallError>
421where
422    Input: ArgumentEncoder,
423    Output: for<'a> ArgumentDecoder<'a>,
424{
425    with_candid(input, |bytes| {
426        env.query_call(canister_id, sender, method, bytes)
427    })
428}
429
430/// Call a canister candid method, authenticated.
431/// The state machine executes update calls synchronously, so there is no need to poll for the result.
432pub fn call_candid_as<Input, Output>(
433    env: &StateMachine,
434    canister_id: Principal,
435    sender: Principal,
436    method: &str,
437    input: Input,
438) -> Result<Output, CallError>
439where
440    Input: ArgumentEncoder,
441    Output: for<'a> ArgumentDecoder<'a>,
442{
443    with_candid(input, |bytes| {
444        env.update_call(canister_id, sender, method, bytes)
445    })
446}
447
448/// Call a canister candid method, anonymous.
449/// The state machine executes update calls synchronously, so there is no need to poll for the result.
450pub fn call_candid<Input, Output>(
451    env: &StateMachine,
452    canister_id: Principal,
453    method: &str,
454    input: Input,
455) -> Result<Output, CallError>
456where
457    Input: ArgumentEncoder,
458    Output: for<'a> ArgumentDecoder<'a>,
459{
460    call_candid_as(env, canister_id, Principal::anonymous(), method, input)
461}
462
463/// User-facing error codes.
464///
465/// The error codes are currently assigned using an HTTP-like
466/// convention: the most significant digit is the corresponding reject
467/// code and the rest is just a sequentially assigned two-digit
468/// number.
469#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
470pub enum ErrorCode {
471    SubnetOversubscribed = 101,
472    MaxNumberOfCanistersReached = 102,
473    CanisterOutputQueueFull = 201,
474    IngressMessageTimeout = 202,
475    CanisterQueueNotEmpty = 203,
476    CanisterNotFound = 301,
477    CanisterMethodNotFound = 302,
478    CanisterAlreadyInstalled = 303,
479    CanisterWasmModuleNotFound = 304,
480    InsufficientMemoryAllocation = 402,
481    InsufficientCyclesForCreateCanister = 403,
482    SubnetNotFound = 404,
483    CanisterNotHostedBySubnet = 405,
484    CanisterOutOfCycles = 501,
485    CanisterTrapped = 502,
486    CanisterCalledTrap = 503,
487    CanisterContractViolation = 504,
488    CanisterInvalidWasm = 505,
489    CanisterDidNotReply = 506,
490    CanisterOutOfMemory = 507,
491    CanisterStopped = 508,
492    CanisterStopping = 509,
493    CanisterNotStopped = 510,
494    CanisterStoppingCancelled = 511,
495    CanisterInvalidController = 512,
496    CanisterFunctionNotFound = 513,
497    CanisterNonEmpty = 514,
498    CertifiedStateUnavailable = 515,
499    CanisterRejectedMessage = 516,
500    QueryCallGraphLoopDetected = 517,
501    UnknownManagementMessage = 518,
502    InvalidManagementPayload = 519,
503    InsufficientCyclesInCall = 520,
504    CanisterWasmEngineError = 521,
505    CanisterInstructionLimitExceeded = 522,
506    CanisterInstallCodeRateLimited = 523,
507    CanisterMemoryAccessLimitExceeded = 524,
508    QueryCallGraphTooDeep = 525,
509    QueryCallGraphTotalInstructionLimitExceeded = 526,
510    CompositeQueryCalledInReplicatedMode = 527,
511}
512
513impl fmt::Display for ErrorCode {
514    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
515        // E.g. "IC0301"
516        write!(f, "IC{:04}", *self as i32)
517    }
518}
519
520/// The error that is sent back to users of IC if something goes
521/// wrong. It's designed to be copyable and serializable so that we
522/// can persist it in ingress history.
523#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
524pub struct UserError {
525    pub code: ErrorCode,
526    pub description: String,
527}
528
529impl fmt::Display for UserError {
530    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
531        // E.g. "IC0301: Canister 42 not found"
532        write!(f, "{}: {}", self.code, self.description)
533    }
534}
535
536#[derive(Debug)]
537pub enum CallError {
538    Reject(String),
539    UserError(UserError),
540}
541
542/// This struct describes the different types that executing a Wasm function in
543/// a canister can produce
544#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
545pub enum WasmResult {
546    /// Raw response, returned in a "happy" case
547    Reply(#[serde(with = "serde_bytes")] Vec<u8>),
548    /// Returned with an error message when the canister decides to reject the
549    /// message
550    Reject(String),
551}
552
553/// A helper function that we use to implement both [`call_candid`] and
554/// [`query_candid`].
555pub fn with_candid<Input, Output>(
556    input: Input,
557    f: impl FnOnce(Vec<u8>) -> Result<WasmResult, UserError>,
558) -> Result<Output, CallError>
559where
560    Input: ArgumentEncoder,
561    Output: for<'a> ArgumentDecoder<'a>,
562{
563    let in_bytes = encode_args(input).expect("failed to encode args");
564    match f(in_bytes) {
565        Ok(WasmResult::Reply(out_bytes)) => Ok(decode_args(&out_bytes).unwrap_or_else(|e| {
566            panic!(
567                "Failed to decode response as candid type {}:\nerror: {}\nbytes: {:?}\nutf8: {}",
568                std::any::type_name::<Output>(),
569                e,
570                out_bytes,
571                String::from_utf8_lossy(&out_bytes),
572            )
573        })),
574        Ok(WasmResult::Reject(message)) => Err(CallError::Reject(message)),
575        Err(user_error) => Err(CallError::UserError(user_error)),
576    }
577}
ic_test_state_machine_client/lib.rs

ic_test_state_machine_client/
lib.rs
