ic_utils/interfaces/management_canister/
builders.rs

1//! Builder interfaces for some method calls of the management canister.
2
3#[doc(inline)]
4pub use super::attributes::{
5    ComputeAllocation, FreezingThreshold, LogMemoryLimit, MemoryAllocation, ReservedCyclesLimit,
6    WasmMemoryLimit,
7};
8use super::{ChunkHash, LogVisibility, ManagementCanister};
9use crate::call::CallFuture;
10use crate::{
11    call::AsyncCall, canister::Argument, interfaces::management_canister::MgmtMethod, Canister,
12};
13use async_trait::async_trait;
14use candid::{utils::ArgumentEncoder, CandidType, Deserialize, Nat};
15use futures_util::{
16    future::ready,
17    stream::{self, FuturesUnordered},
18    FutureExt, Stream, StreamExt, TryStreamExt,
19};
20use ic_agent::{agent::CallResponse, export::Principal, AgentError};
21pub use ic_management_canister_types::{
22    CanisterInstallMode, CanisterSettings, EnvironmentVariable, InstallCodeArgs, UpgradeFlags,
23    WasmMemoryPersistence,
24};
25use sha2::{Digest, Sha256};
26use std::{
27    collections::BTreeSet,
28    convert::{From, TryInto},
29    future::IntoFuture,
30    pin::Pin,
31};
32
33/// A builder for a `create_canister` call.
34#[derive(Debug)]
35pub struct CreateCanisterBuilder<'agent, 'canister: 'agent> {
36    canister: &'canister Canister<'agent>,
37    effective_canister_id: Principal,
38    controllers: Option<Result<Vec<Principal>, AgentError>>,
39    compute_allocation: Option<Result<ComputeAllocation, AgentError>>,
40    memory_allocation: Option<Result<MemoryAllocation, AgentError>>,
41    freezing_threshold: Option<Result<FreezingThreshold, AgentError>>,
42    reserved_cycles_limit: Option<Result<ReservedCyclesLimit, AgentError>>,
43    wasm_memory_limit: Option<Result<WasmMemoryLimit, AgentError>>,
44    wasm_memory_threshold: Option<Result<WasmMemoryLimit, AgentError>>,
45    log_visibility: Option<Result<LogVisibility, AgentError>>,
46    log_memory_limit: Option<Result<LogMemoryLimit, AgentError>>,
47    environment_variables: Option<Result<Vec<EnvironmentVariable>, AgentError>>,
48    is_provisional_create: bool,
49    amount: Option<u128>,
50    specified_id: Option<Principal>,
51}
52
53impl<'agent, 'canister: 'agent> CreateCanisterBuilder<'agent, 'canister> {
54    /// Create an `CreateCanister` builder, which is also an `AsyncCall` implementation.
55    pub fn builder(canister: &'canister Canister<'agent>) -> Self {
56        Self {
57            canister,
58            effective_canister_id: Principal::management_canister(),
59            controllers: None,
60            compute_allocation: None,
61            memory_allocation: None,
62            freezing_threshold: None,
63            reserved_cycles_limit: None,
64            wasm_memory_limit: None,
65            wasm_memory_threshold: None,
66            log_visibility: None,
67            log_memory_limit: None,
68            environment_variables: None,
69            is_provisional_create: false,
70            amount: None,
71            specified_id: None,
72        }
73    }
74
75    /// Until developers can convert real ICP tokens to provision a new canister with cycles,
76    /// the system provides the `provisional_create_canister_with_cycles` method.
77    /// It behaves as `create_canister`, but initializes the canister’s balance with amount fresh cycles
78    /// (using `MAX_CANISTER_BALANCE` if amount = null, else capping the balance at `MAX_CANISTER_BALANCE`).
79    /// Cycles added to this call via `ic0.call_cycles_add` are returned to the caller.
80    /// This method is only available in local development instances, and will be removed in the future.
81    #[allow(clippy::wrong_self_convention)]
82    pub fn as_provisional_create_with_amount(self, amount: Option<u128>) -> Self {
83        Self {
84            is_provisional_create: true,
85            amount,
86            ..self
87        }
88    }
89
90    /// Specify the canister id.
91    ///
92    /// The `effective_canister_id` will also be set with the same value so that ic-ref can determine
93    /// the target subnet of this request. The replica implementation ignores it.
94    pub fn as_provisional_create_with_specified_id(self, specified_id: Principal) -> Self {
95        Self {
96            is_provisional_create: true,
97            specified_id: Some(specified_id),
98            effective_canister_id: specified_id,
99            ..self
100        }
101    }
102
103    /// Pass in an effective canister id for the update call.
104    pub fn with_effective_canister_id<C, E>(self, effective_canister_id: C) -> Self
105    where
106        E: std::fmt::Display,
107        C: TryInto<Principal, Error = E>,
108    {
109        match effective_canister_id.try_into() {
110            Ok(effective_canister_id) => Self {
111                effective_canister_id,
112                ..self
113            },
114            Err(_) => self,
115        }
116    }
117
118    /// Pass in a controller for the canister. Can be called multiple times to add multiple controllers.
119    pub fn with_controller<C, E>(self, controller: C) -> Self
120    where
121        E: std::fmt::Display,
122        C: TryInto<Principal, Error = E>,
123    {
124        let controller_result = controller
125            .try_into()
126            .map_err(|e| AgentError::MessageError(format!("{e}")));
127        let controllers = match (controller_result, self.controllers) {
128            (_, Some(Err(sticky))) => Some(Err(sticky)),
129            (Err(e), _) => Some(Err(e)),
130            (Ok(controller), Some(Ok(mut controllers))) => {
131                controllers.push(controller);
132                Some(Ok(controllers))
133            }
134            (Ok(controller), None) => Some(Ok(vec![controller])),
135        };
136        Self {
137            controllers,
138            ..self
139        }
140    }
141
142    /// Pass in a compute allocation value for the canister.
143    pub fn with_compute_allocation<C, E>(self, compute_allocation: C) -> Self
144    where
145        E: std::fmt::Display,
146        C: TryInto<ComputeAllocation, Error = E>,
147    {
148        Self {
149            compute_allocation: Some(
150                compute_allocation
151                    .try_into()
152                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
153            ),
154            ..self
155        }
156    }
157
158    /// Pass in a memory allocation value for the canister.
159    pub fn with_memory_allocation<C, E>(self, memory_allocation: C) -> Self
160    where
161        E: std::fmt::Display,
162        C: TryInto<MemoryAllocation, Error = E>,
163    {
164        Self {
165            memory_allocation: Some(
166                memory_allocation
167                    .try_into()
168                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
169            ),
170            ..self
171        }
172    }
173
174    /// Pass in a freezing threshold value for the canister.
175    pub fn with_freezing_threshold<C, E>(self, freezing_threshold: C) -> Self
176    where
177        E: std::fmt::Display,
178        C: TryInto<FreezingThreshold, Error = E>,
179    {
180        Self {
181            freezing_threshold: Some(
182                freezing_threshold
183                    .try_into()
184                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
185            ),
186            ..self
187        }
188    }
189
190    /// Pass in a reserved cycles limit value for the canister.
191    pub fn with_reserved_cycles_limit<C, E>(self, limit: C) -> Self
192    where
193        E: std::fmt::Display,
194        C: TryInto<ReservedCyclesLimit, Error = E>,
195    {
196        Self {
197            reserved_cycles_limit: Some(
198                limit
199                    .try_into()
200                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
201            ),
202            ..self
203        }
204    }
205
206    /// Pass in a Wasm memory limit value for the canister.
207    pub fn with_wasm_memory_limit<C, E>(self, wasm_memory_limit: C) -> Self
208    where
209        E: std::fmt::Display,
210        C: TryInto<WasmMemoryLimit, Error = E>,
211    {
212        Self {
213            wasm_memory_limit: Some(
214                wasm_memory_limit
215                    .try_into()
216                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
217            ),
218            ..self
219        }
220    }
221
222    /// Pass in a Wasm memory threshold value for the canister.
223    pub fn with_wasm_memory_threshold<C, E>(self, wasm_memory_threshold: C) -> Self
224    where
225        E: std::fmt::Display,
226        C: TryInto<WasmMemoryLimit, Error = E>,
227    {
228        Self {
229            wasm_memory_threshold: Some(
230                wasm_memory_threshold
231                    .try_into()
232                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
233            ),
234            ..self
235        }
236    }
237
238    /// Pass in a log memory limit value for the canister.
239    pub fn with_log_memory_limit<C, E>(self, log_memory_limit: C) -> Self
240    where
241        E: std::fmt::Display,
242        C: TryInto<LogMemoryLimit, Error = E>,
243    {
244        Self {
245            log_memory_limit: Some(
246                log_memory_limit
247                    .try_into()
248                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
249            ),
250            ..self
251        }
252    }
253
254    /// Pass in a log visibility setting for the canister.
255    pub fn with_log_visibility<C, E>(self, log_visibility: C) -> Self
256    where
257        E: std::fmt::Display,
258        C: TryInto<LogVisibility, Error = E>,
259    {
260        Self {
261            log_visibility: Some(
262                log_visibility
263                    .try_into()
264                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
265            ),
266            ..self
267        }
268    }
269
270    /// Pass in environment variables for the canister.
271    pub fn with_environment_variables<E, C>(self, environment_variables: C) -> Self
272    where
273        E: std::fmt::Display,
274        C: TryInto<Vec<EnvironmentVariable>, Error = E>,
275    {
276        Self {
277            environment_variables: Some(
278                environment_variables
279                    .try_into()
280                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
281            ),
282            ..self
283        }
284    }
285
286    /// Create an [`AsyncCall`] implementation that, when called, will create a
287    /// canister.
288    pub fn build(self) -> Result<impl 'agent + AsyncCall<Value = (Principal,)>, AgentError> {
289        let controllers = match self.controllers {
290            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
291            Some(Ok(x)) => Some(x),
292            None => None,
293        };
294        let compute_allocation = match self.compute_allocation {
295            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
296            Some(Ok(x)) => Some(Nat::from(u8::from(x))),
297            None => None,
298        };
299        let memory_allocation = match self.memory_allocation {
300            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
301            Some(Ok(x)) => Some(Nat::from(u64::from(x))),
302            None => None,
303        };
304        let freezing_threshold = match self.freezing_threshold {
305            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
306            Some(Ok(x)) => Some(Nat::from(u64::from(x))),
307            None => None,
308        };
309        let reserved_cycles_limit = match self.reserved_cycles_limit {
310            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
311            Some(Ok(x)) => Some(Nat::from(u128::from(x))),
312            None => None,
313        };
314        let wasm_memory_limit = match self.wasm_memory_limit {
315            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
316            Some(Ok(x)) => Some(Nat::from(u64::from(x))),
317            None => None,
318        };
319        let wasm_memory_threshold = match self.wasm_memory_threshold {
320            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
321            Some(Ok(x)) => Some(Nat::from(u64::from(x))),
322            None => None,
323        };
324        let log_visibility = match self.log_visibility {
325            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
326            Some(Ok(x)) => Some(x),
327            None => None,
328        };
329        let log_memory_limit = match self.log_memory_limit {
330            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
331            Some(Ok(x)) => Some(Nat::from(u64::from(x))),
332            None => None,
333        };
334        let environment_variables = match self.environment_variables {
335            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
336            Some(Ok(x)) => Some(x),
337            None => None,
338        };
339
340        #[derive(Deserialize, CandidType)]
341        struct Out {
342            canister_id: Principal,
343        }
344
345        let async_builder = if self.is_provisional_create {
346            #[derive(CandidType)]
347            struct In {
348                amount: Option<Nat>,
349                settings: CanisterSettings,
350                specified_id: Option<Principal>,
351            }
352            let in_arg = In {
353                amount: self.amount.map(Nat::from),
354                settings: CanisterSettings {
355                    controllers,
356                    compute_allocation,
357                    memory_allocation,
358                    freezing_threshold,
359                    reserved_cycles_limit,
360                    wasm_memory_limit,
361                    wasm_memory_threshold,
362                    log_visibility,
363                    log_memory_limit: log_memory_limit.clone(),
364                    environment_variables,
365                },
366                specified_id: self.specified_id,
367            };
368            self.canister
369                .update(MgmtMethod::ProvisionalCreateCanisterWithCycles.as_ref())
370                .with_arg(in_arg)
371                .with_effective_canister_id(self.effective_canister_id)
372        } else {
373            self.canister
374                .update(MgmtMethod::CreateCanister.as_ref())
375                .with_arg(CanisterSettings {
376                    controllers,
377                    compute_allocation,
378                    memory_allocation,
379                    freezing_threshold,
380                    reserved_cycles_limit,
381                    wasm_memory_limit,
382                    wasm_memory_threshold,
383                    log_visibility,
384                    log_memory_limit,
385                    environment_variables,
386                })
387                .with_effective_canister_id(self.effective_canister_id)
388        };
389
390        Ok(async_builder
391            .build()
392            .map(|result: (Out,)| (result.0.canister_id,)))
393    }
394
395    /// Make a call. This is equivalent to the [`AsyncCall::call`].
396    pub async fn call(self) -> Result<CallResponse<(Principal,)>, AgentError> {
397        self.build()?.call().await
398    }
399
400    /// Make a call. This is equivalent to the [`AsyncCall::call_and_wait`].
401    pub async fn call_and_wait(self) -> Result<(Principal,), AgentError> {
402        self.build()?.call_and_wait().await
403    }
404}
405
406#[cfg_attr(target_family = "wasm", async_trait(?Send))]
407#[cfg_attr(not(target_family = "wasm"), async_trait)]
408impl<'agent, 'canister: 'agent> AsyncCall for CreateCanisterBuilder<'agent, 'canister> {
409    type Value = (Principal,);
410
411    async fn call(self) -> Result<CallResponse<(Principal,)>, AgentError> {
412        self.build()?.call().await
413    }
414
415    async fn call_and_wait(self) -> Result<(Principal,), AgentError> {
416        self.build()?.call_and_wait().await
417    }
418}
419
420impl<'agent, 'canister: 'agent> IntoFuture for CreateCanisterBuilder<'agent, 'canister> {
421    type IntoFuture = CallFuture<'agent, (Principal,)>;
422    type Output = Result<(Principal,), AgentError>;
423
424    fn into_future(self) -> Self::IntoFuture {
425        AsyncCall::call_and_wait(self)
426    }
427}
428
429#[doc(hidden)]
430#[deprecated(since = "0.42.0", note = "Please use UpgradeFlags instead")]
431pub type CanisterUpgradeOptions = UpgradeFlags;
432
433#[doc(hidden)]
434#[deprecated(since = "0.42.0", note = "Please use CanisterInstallMode instead")]
435pub type InstallMode = CanisterInstallMode;
436
437#[doc(hidden)]
438#[deprecated(since = "0.42.0", note = "Please use InstallCodeArgs instead")]
439pub type CanisterInstall = InstallCodeArgs;
440
441/// A builder for an `install_code` call.
442#[derive(Debug)]
443pub struct InstallCodeBuilder<'agent, 'canister: 'agent> {
444    canister: &'canister Canister<'agent>,
445    canister_id: Principal,
446    wasm: &'canister [u8],
447    arg: Argument,
448    mode: Option<CanisterInstallMode>,
449}
450
451impl<'agent, 'canister: 'agent> InstallCodeBuilder<'agent, 'canister> {
452    /// Create an `InstallCode` builder, which is also an `AsyncCall` implementation.
453    pub fn builder(
454        canister: &'canister Canister<'agent>,
455        canister_id: &Principal,
456        wasm: &'canister [u8],
457    ) -> Self {
458        Self {
459            canister,
460            canister_id: *canister_id,
461            wasm,
462            arg: Default::default(),
463            mode: None,
464        }
465    }
466
467    /// Set the argument to the installation, which will be passed to the init
468    /// method of the canister. Can be called at most once.
469    pub fn with_arg<Argument: CandidType>(
470        mut self,
471        arg: Argument,
472    ) -> InstallCodeBuilder<'agent, 'canister> {
473        self.arg.set_idl_arg(arg);
474        self
475    }
476    /// Set the argument with multiple arguments as tuple to the installation,
477    /// which will be passed to the init method of the canister. Can be called at most once.
478    pub fn with_args(mut self, tuple: impl ArgumentEncoder) -> Self {
479        assert!(self.arg.0.is_none(), "argument is being set more than once");
480        self.arg = Argument::from_candid(tuple);
481        self
482    }
483    /// Set the argument passed in to the canister with raw bytes. Can be called at most once.
484    pub fn with_raw_arg(mut self, arg: Vec<u8>) -> InstallCodeBuilder<'agent, 'canister> {
485        self.arg.set_raw_arg(arg);
486        self
487    }
488
489    /// Pass in the [`CanisterInstallMode`].
490    pub fn with_mode(self, mode: CanisterInstallMode) -> Self {
491        Self {
492            mode: Some(mode),
493            ..self
494        }
495    }
496
497    /// Create an [`AsyncCall`] implementation that, when called, will install the
498    /// canister.
499    pub fn build(self) -> Result<impl 'agent + AsyncCall<Value = ()>, AgentError> {
500        Ok(self
501            .canister
502            .update(MgmtMethod::InstallCode.as_ref())
503            .with_arg(InstallCodeArgs {
504                mode: self.mode.unwrap_or(CanisterInstallMode::Install),
505                canister_id: self.canister_id,
506                wasm_module: self.wasm.to_owned(),
507                arg: self.arg.serialize()?,
508                sender_canister_version: None,
509            })
510            .with_effective_canister_id(self.canister_id)
511            .build())
512    }
513
514    /// Make a call. This is equivalent to the [`AsyncCall::call`].
515    pub async fn call(self) -> Result<CallResponse<()>, AgentError> {
516        self.build()?.call().await
517    }
518
519    /// Make a call. This is equivalent to the [`AsyncCall::call_and_wait`].
520    pub async fn call_and_wait(self) -> Result<(), AgentError> {
521        self.build()?.call_and_wait().await
522    }
523}
524
525#[cfg_attr(target_family = "wasm", async_trait(?Send))]
526#[cfg_attr(not(target_family = "wasm"), async_trait)]
527impl<'agent, 'canister: 'agent> AsyncCall for InstallCodeBuilder<'agent, 'canister> {
528    type Value = ();
529
530    async fn call(self) -> Result<CallResponse<()>, AgentError> {
531        self.build()?.call().await
532    }
533
534    async fn call_and_wait(self) -> Result<(), AgentError> {
535        self.build()?.call_and_wait().await
536    }
537}
538
539impl<'agent, 'canister: 'agent> IntoFuture for InstallCodeBuilder<'agent, 'canister> {
540    type IntoFuture = CallFuture<'agent, ()>;
541    type Output = Result<(), AgentError>;
542
543    fn into_future(self) -> Self::IntoFuture {
544        AsyncCall::call_and_wait(self)
545    }
546}
547
548/// A builder for an `install_chunked_code` call.
549#[derive(Debug)]
550pub struct InstallChunkedCodeBuilder<'agent, 'canister> {
551    canister: &'canister Canister<'agent>,
552    target_canister: Principal,
553    store_canister: Option<Principal>,
554    chunk_hashes_list: Vec<ChunkHash>,
555    wasm_module_hash: Vec<u8>,
556    arg: Argument,
557    mode: CanisterInstallMode,
558}
559
560impl<'agent: 'canister, 'canister> InstallChunkedCodeBuilder<'agent, 'canister> {
561    /// Create an `InstallChunkedCodeBuilder`.
562    pub fn builder(
563        canister: &'canister Canister<'agent>,
564        target_canister: Principal,
565        wasm_module_hash: &[u8],
566    ) -> Self {
567        Self {
568            canister,
569            target_canister,
570            wasm_module_hash: wasm_module_hash.to_vec(),
571            store_canister: None,
572            chunk_hashes_list: vec![],
573            arg: Argument::new(),
574            mode: CanisterInstallMode::Install,
575        }
576    }
577
578    /// Set the chunks to install. These must previously have been set with [`ManagementCanister::upload_chunk`].
579    pub fn with_chunk_hashes(mut self, chunk_hashes: Vec<ChunkHash>) -> Self {
580        self.chunk_hashes_list = chunk_hashes;
581        self
582    }
583
584    /// Set the canister to pull uploaded chunks from. By default this is the same as the target canister.
585    pub fn with_store_canister(mut self, store_canister: Principal) -> Self {
586        self.store_canister = Some(store_canister);
587        self
588    }
589
590    /// Set the argument to the installation, which will be passed to the init
591    /// method of the canister. Can be called at most once.
592    pub fn with_arg(mut self, argument: impl CandidType) -> Self {
593        self.arg.set_idl_arg(argument);
594        self
595    }
596
597    /// Set the argument with multiple arguments as tuple to the installation,
598    /// which will be passed to the init method of the canister. Can be called at most once.
599    pub fn with_args(mut self, argument: impl ArgumentEncoder) -> Self {
600        assert!(self.arg.0.is_none(), "argument is being set more than once");
601        self.arg = Argument::from_candid(argument);
602        self
603    }
604
605    /// Set the argument passed in to the canister with raw bytes. Can be called at most once.
606    pub fn with_raw_arg(mut self, argument: Vec<u8>) -> Self {
607        self.arg.set_raw_arg(argument);
608        self
609    }
610
611    /// Set the [`CanisterInstallMode`].
612    pub fn with_install_mode(mut self, mode: CanisterInstallMode) -> Self {
613        self.mode = mode;
614        self
615    }
616
617    /// Create an [`AsyncCall`] implementation that, when called, will install the canister.
618    pub fn build(self) -> Result<impl 'agent + AsyncCall<Value = ()>, AgentError> {
619        #[derive(CandidType)]
620        struct In {
621            mode: CanisterInstallMode,
622            target_canister: Principal,
623            store_canister: Option<Principal>,
624            chunk_hashes_list: Vec<ChunkHash>,
625            wasm_module_hash: Vec<u8>,
626            arg: Vec<u8>,
627            sender_canister_version: Option<u64>,
628        }
629        let Self {
630            mode,
631            target_canister,
632            store_canister,
633            chunk_hashes_list,
634            wasm_module_hash,
635            arg,
636            ..
637        } = self;
638        Ok(self
639            .canister
640            .update(MgmtMethod::InstallChunkedCode.as_ref())
641            .with_arg(In {
642                mode,
643                target_canister,
644                store_canister,
645                chunk_hashes_list,
646                wasm_module_hash,
647                arg: arg.serialize()?,
648                sender_canister_version: None,
649            })
650            .with_effective_canister_id(target_canister)
651            .build())
652    }
653
654    /// Make the call. This is equivalent to [`AsyncCall::call`].
655    pub async fn call(self) -> Result<CallResponse<()>, AgentError> {
656        self.build()?.call().await
657    }
658
659    /// Make the call. This is equivalent to [`AsyncCall::call_and_wait`].
660    pub async fn call_and_wait(self) -> Result<(), AgentError> {
661        self.build()?.call_and_wait().await
662    }
663}
664
665#[cfg_attr(target_family = "wasm", async_trait(?Send))]
666#[cfg_attr(not(target_family = "wasm"), async_trait)]
667impl<'agent, 'canister: 'agent> AsyncCall for InstallChunkedCodeBuilder<'agent, 'canister> {
668    type Value = ();
669
670    async fn call(self) -> Result<CallResponse<()>, AgentError> {
671        self.call().await
672    }
673
674    async fn call_and_wait(self) -> Result<(), AgentError> {
675        self.call_and_wait().await
676    }
677}
678
679impl<'agent, 'canister: 'agent> IntoFuture for InstallChunkedCodeBuilder<'agent, 'canister> {
680    type IntoFuture = CallFuture<'agent, ()>;
681    type Output = Result<(), AgentError>;
682
683    fn into_future(self) -> Self::IntoFuture {
684        AsyncCall::call_and_wait(self)
685    }
686}
687
688/// A builder for a [`ManagementCanister::install`] call. This automatically selects one-shot installation or chunked installation depending on module size.
689///
690/// # Warnings
691///
692/// This will clear chunked code storage if chunked installation is used. Do not use with canisters that you are manually uploading chunked code to.
693#[derive(Debug)]
694pub struct InstallBuilder<'agent, 'canister, 'builder> {
695    canister: &'canister ManagementCanister<'agent>,
696    canister_id: Principal,
697    // more precise lifetimes are used here at risk of annoying the user
698    // because `wasm` may be memory-mapped which is tricky to lifetime
699    wasm: &'builder [u8],
700    arg: Argument,
701    mode: CanisterInstallMode,
702}
703
704impl<'agent: 'canister, 'canister: 'builder, 'builder> InstallBuilder<'agent, 'canister, 'builder> {
705    // Messages are a maximum of 2MiB. Thus basic installation should cap the wasm and arg size at 1.85MiB, since
706    // the current API is definitely not going to produce 150KiB of framing data for it.
707    const CHUNK_CUTOFF: usize = (1.85 * 1024. * 1024.) as usize;
708
709    /// Create a canister installation builder.
710    pub fn builder(
711        canister: &'canister ManagementCanister<'agent>,
712        canister_id: &Principal,
713        wasm: &'builder [u8],
714    ) -> Self {
715        Self {
716            canister,
717            canister_id: *canister_id,
718            wasm,
719            arg: Default::default(),
720            mode: CanisterInstallMode::Install,
721        }
722    }
723
724    /// Set the argument to the installation, which will be passed to the init
725    /// method of the canister. Can be called at most once.
726    pub fn with_arg<Argument: CandidType>(mut self, arg: Argument) -> Self {
727        self.arg.set_idl_arg(arg);
728        self
729    }
730    /// Set the argument with multiple arguments as tuple to the installation,
731    /// which will be passed to the init method of the canister. Can be called at most once.
732    pub fn with_args(mut self, tuple: impl ArgumentEncoder) -> Self {
733        assert!(self.arg.0.is_none(), "argument is being set more than once");
734        self.arg = Argument::from_candid(tuple);
735        self
736    }
737    /// Set the argument passed in to the canister with raw bytes. Can be called at most once.
738    pub fn with_raw_arg(mut self, arg: Vec<u8>) -> Self {
739        self.arg.set_raw_arg(arg);
740        self
741    }
742
743    /// Pass in the [`CanisterInstallMode`].
744    pub fn with_mode(self, mode: CanisterInstallMode) -> Self {
745        Self { mode, ..self }
746    }
747
748    /// Invoke the installation process. This may result in many calls which may take several seconds;
749    /// use [`call_and_wait_with_progress`](Self::call_and_wait_with_progress) if you want progress reporting.
750    pub async fn call_and_wait(self) -> Result<(), AgentError> {
751        self.call_and_wait_with_progress()
752            .await
753            .try_for_each(|_| ready(Ok(())))
754            .await
755    }
756
757    /// Invoke the installation process. The returned stream must be iterated to completion; it is used to track progress,
758    /// as installation may take arbitrarily long, and is intended to be passed to functions like `indicatif::ProgressBar::wrap_stream`.
759    /// There are exactly [`size_hint().0`](Stream::size_hint) steps.
760    pub async fn call_and_wait_with_progress(
761        self,
762    ) -> impl Stream<Item = Result<(), AgentError>> + 'builder {
763        let stream_res = /* try { */ async move {
764            let arg = self.arg.serialize()?;
765            let stream: BoxStream<'_, _> =
766                if self.wasm.len() + arg.len() < Self::CHUNK_CUTOFF {
767                    Box::pin(
768                        async move {
769                            self.canister
770                                .install_code(&self.canister_id, self.wasm)
771                                .with_raw_arg(arg)
772                                .with_mode(self.mode)
773                                .call_and_wait()
774                                .await
775                        }
776                        .into_stream(),
777                    )
778                } else {
779                    let (existing_chunks,) = self.canister.stored_chunks(&self.canister_id).call_and_wait().await?;
780                    let existing_chunks = existing_chunks.into_iter().map(|c| c.hash).collect::<BTreeSet<_>>();
781                    let all_chunks = self.wasm.chunks(1024 * 1024).map(|x| (Sha256::digest(x).to_vec(), x)).collect::<Vec<_>>();
782                    let mut to_upload_chunks = vec![];
783                    for (hash, chunk) in &all_chunks {
784                        if !existing_chunks.contains(hash) {
785                            to_upload_chunks.push(*chunk);
786                        }
787                    }
788
789                    let upload_chunks_stream = FuturesUnordered::new();
790                    for chunk in to_upload_chunks {
791                        upload_chunks_stream.push(async move {
792                            let (_res,) = self
793                                .canister
794                                .upload_chunk(&self.canister_id, &ic_management_canister_types::UploadChunkArgs {
795                                    canister_id: self.canister_id,
796                                    chunk: chunk.to_vec(),
797                                })
798                                .call_and_wait()
799                                .await?;
800                            Ok(())
801                        });
802                    }
803                    let install_chunked_code_stream = async move {
804                        let results = all_chunks.iter().map(|(hash,_)| ChunkHash{ hash: hash.clone() }).collect();
805                        self.canister
806                            .install_chunked_code(
807                                &self.canister_id,
808                                &Sha256::digest(self.wasm),
809                            )
810                            .with_chunk_hashes(results)
811                            .with_raw_arg(arg)
812                            .with_install_mode(self.mode)
813                            .call_and_wait()
814                            .await
815                    }
816                    .into_stream();
817                    let clear_store_stream = async move {
818                        self.canister
819                            .clear_chunk_store(&self.canister_id)
820                            .call_and_wait()
821                            .await
822                    }
823                    .into_stream();
824
825                    Box::pin(
826                        upload_chunks_stream
827                            .chain(install_chunked_code_stream)
828                            .chain(clear_store_stream                        ),
829                    )
830                };
831            Ok(stream)
832        }.await;
833        match stream_res {
834            Ok(stream) => stream,
835            Err(err) => Box::pin(stream::once(async { Err(err) })),
836        }
837    }
838}
839
840impl<'agent: 'canister, 'canister: 'builder, 'builder> IntoFuture
841    for InstallBuilder<'agent, 'canister, 'builder>
842{
843    type IntoFuture = CallFuture<'builder, ()>;
844    type Output = Result<(), AgentError>;
845
846    fn into_future(self) -> Self::IntoFuture {
847        Box::pin(self.call_and_wait())
848    }
849}
850
851/// A builder for an `update_settings` call.
852#[derive(Debug)]
853pub struct UpdateSettingsBuilder<'agent, 'canister: 'agent> {
854    canister: &'canister Canister<'agent>,
855    canister_id: Principal,
856    controllers: Option<Result<Vec<Principal>, AgentError>>,
857    compute_allocation: Option<Result<ComputeAllocation, AgentError>>,
858    memory_allocation: Option<Result<MemoryAllocation, AgentError>>,
859    freezing_threshold: Option<Result<FreezingThreshold, AgentError>>,
860    reserved_cycles_limit: Option<Result<ReservedCyclesLimit, AgentError>>,
861    wasm_memory_limit: Option<Result<WasmMemoryLimit, AgentError>>,
862    wasm_memory_threshold: Option<Result<WasmMemoryLimit, AgentError>>,
863    log_visibility: Option<Result<LogVisibility, AgentError>>,
864    log_memory_limit: Option<Result<LogMemoryLimit, AgentError>>,
865    environment_variables: Option<Result<Vec<EnvironmentVariable>, AgentError>>,
866}
867
868impl<'agent, 'canister: 'agent> UpdateSettingsBuilder<'agent, 'canister> {
869    /// Create an `UpdateCanister` builder, which is also an `AsyncCall` implementation.
870    pub fn builder(canister: &'canister Canister<'agent>, canister_id: &Principal) -> Self {
871        Self {
872            canister,
873            canister_id: *canister_id,
874            controllers: None,
875            compute_allocation: None,
876            memory_allocation: None,
877            freezing_threshold: None,
878            reserved_cycles_limit: None,
879            wasm_memory_limit: None,
880            wasm_memory_threshold: None,
881            log_visibility: None,
882            log_memory_limit: None,
883            environment_variables: None,
884        }
885    }
886
887    /// Pass in a controller for the canister. Can be called multiple times to add multiple controllers.
888    pub fn with_controller<C, E>(self, controller: C) -> Self
889    where
890        E: std::fmt::Display,
891        C: TryInto<Principal, Error = E>,
892    {
893        let controller_result = controller
894            .try_into()
895            .map_err(|e| AgentError::MessageError(format!("{e}")));
896        let controllers = match (controller_result, self.controllers) {
897            (_, Some(Err(sticky))) => Some(Err(sticky)),
898            (Err(e), _) => Some(Err(e)),
899            (Ok(controller), Some(Ok(mut controllers))) => {
900                controllers.push(controller);
901                Some(Ok(controllers))
902            }
903            (Ok(controller), None) => Some(Ok(vec![controller])),
904        };
905        Self {
906            controllers,
907            ..self
908        }
909    }
910
911    /// Pass in a compute allocation value for the canister.
912    pub fn with_compute_allocation<C, E>(self, compute_allocation: C) -> Self
913    where
914        E: std::fmt::Display,
915        C: TryInto<ComputeAllocation, Error = E>,
916    {
917        Self {
918            compute_allocation: Some(
919                compute_allocation
920                    .try_into()
921                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
922            ),
923            ..self
924        }
925    }
926
927    /// Pass in a memory allocation value for the canister.
928    pub fn with_memory_allocation<C, E>(self, memory_allocation: C) -> Self
929    where
930        E: std::fmt::Display,
931        C: TryInto<MemoryAllocation, Error = E>,
932    {
933        Self {
934            memory_allocation: Some(
935                memory_allocation
936                    .try_into()
937                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
938            ),
939            ..self
940        }
941    }
942
943    /// Pass in a freezing threshold value for the canister.
944    pub fn with_freezing_threshold<C, E>(self, freezing_threshold: C) -> Self
945    where
946        E: std::fmt::Display,
947        C: TryInto<FreezingThreshold, Error = E>,
948    {
949        Self {
950            freezing_threshold: Some(
951                freezing_threshold
952                    .try_into()
953                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
954            ),
955            ..self
956        }
957    }
958
959    /// Pass in a reserved cycles limit value for the canister.
960    pub fn with_reserved_cycles_limit<C, E>(self, limit: C) -> Self
961    where
962        E: std::fmt::Display,
963        C: TryInto<ReservedCyclesLimit, Error = E>,
964    {
965        Self {
966            reserved_cycles_limit: Some(
967                limit
968                    .try_into()
969                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
970            ),
971            ..self
972        }
973    }
974
975    /// Pass in a Wasm memory limit value for the canister.
976    pub fn with_wasm_memory_limit<C, E>(self, wasm_memory_limit: C) -> Self
977    where
978        E: std::fmt::Display,
979        C: TryInto<WasmMemoryLimit, Error = E>,
980    {
981        Self {
982            wasm_memory_limit: Some(
983                wasm_memory_limit
984                    .try_into()
985                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
986            ),
987            ..self
988        }
989    }
990
991    /// Pass in a Wasm memory threshold value for the canister.
992    pub fn with_wasm_memory_threshold<C, E>(self, wasm_memory_threshold: C) -> Self
993    where
994        E: std::fmt::Display,
995        C: TryInto<WasmMemoryLimit, Error = E>,
996    {
997        Self {
998            wasm_memory_threshold: Some(
999                wasm_memory_threshold
1000                    .try_into()
1001                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
1002            ),
1003            ..self
1004        }
1005    }
1006
1007    /// Pass in a log memory limit value for the canister.
1008    pub fn with_log_memory_limit<C, E>(self, log_memory_limit: C) -> Self
1009    where
1010        E: std::fmt::Display,
1011        C: TryInto<LogMemoryLimit, Error = E>,
1012    {
1013        Self {
1014            log_memory_limit: Some(
1015                log_memory_limit
1016                    .try_into()
1017                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
1018            ),
1019            ..self
1020        }
1021    }
1022
1023    /// Pass in a log visibility setting for the canister.
1024    pub fn with_log_visibility<C, E>(self, log_visibility: C) -> Self
1025    where
1026        E: std::fmt::Display,
1027        C: TryInto<LogVisibility, Error = E>,
1028    {
1029        Self {
1030            log_visibility: Some(
1031                log_visibility
1032                    .try_into()
1033                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
1034            ),
1035            ..self
1036        }
1037    }
1038
1039    /// Pass in environment variables for the canister.
1040    pub fn with_environment_variables<C, E>(self, environment_variables: C) -> Self
1041    where
1042        E: std::fmt::Display,
1043        C: TryInto<Vec<EnvironmentVariable>, Error = E>,
1044    {
1045        Self {
1046            environment_variables: Some(
1047                environment_variables
1048                    .try_into()
1049                    .map_err(|e| AgentError::MessageError(format!("{e}"))),
1050            ),
1051            ..self
1052        }
1053    }
1054
1055    /// Create an [`AsyncCall`] implementation that, when called, will update a
1056    /// canisters settings.
1057    pub fn build(self) -> Result<impl 'agent + AsyncCall<Value = ()>, AgentError> {
1058        #[derive(CandidType)]
1059        struct In {
1060            canister_id: Principal,
1061            settings: CanisterSettings,
1062        }
1063
1064        let controllers = match self.controllers {
1065            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
1066            Some(Ok(x)) => Some(x),
1067            None => None,
1068        };
1069        let compute_allocation = match self.compute_allocation {
1070            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
1071            Some(Ok(x)) => Some(Nat::from(u8::from(x))),
1072            None => None,
1073        };
1074        let memory_allocation = match self.memory_allocation {
1075            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
1076            Some(Ok(x)) => Some(Nat::from(u64::from(x))),
1077            None => None,
1078        };
1079        let freezing_threshold = match self.freezing_threshold {
1080            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
1081            Some(Ok(x)) => Some(Nat::from(u64::from(x))),
1082            None => None,
1083        };
1084        let reserved_cycles_limit = match self.reserved_cycles_limit {
1085            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
1086            Some(Ok(x)) => Some(Nat::from(u128::from(x))),
1087            None => None,
1088        };
1089        let wasm_memory_limit = match self.wasm_memory_limit {
1090            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
1091            Some(Ok(x)) => Some(Nat::from(u64::from(x))),
1092            None => None,
1093        };
1094        let wasm_memory_threshold = match self.wasm_memory_threshold {
1095            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
1096            Some(Ok(x)) => Some(Nat::from(u64::from(x))),
1097            None => None,
1098        };
1099        let log_visibility = match self.log_visibility {
1100            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
1101            Some(Ok(x)) => Some(x),
1102            None => None,
1103        };
1104        let log_memory_limit = match self.log_memory_limit {
1105            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
1106            Some(Ok(x)) => Some(Nat::from(u64::from(x))),
1107            None => None,
1108        };
1109        let environment_variables = match self.environment_variables {
1110            Some(Err(x)) => return Err(AgentError::MessageError(format!("{x}"))),
1111            Some(Ok(x)) => Some(x),
1112            None => None,
1113        };
1114
1115        Ok(self
1116            .canister
1117            .update(MgmtMethod::UpdateSettings.as_ref())
1118            .with_arg(In {
1119                canister_id: self.canister_id,
1120                settings: CanisterSettings {
1121                    controllers,
1122                    compute_allocation,
1123                    memory_allocation,
1124                    freezing_threshold,
1125                    reserved_cycles_limit,
1126                    wasm_memory_limit,
1127                    wasm_memory_threshold,
1128                    log_visibility,
1129                    log_memory_limit,
1130                    environment_variables,
1131                },
1132            })
1133            .with_effective_canister_id(self.canister_id)
1134            .build())
1135    }
1136
1137    /// Make a call. This is equivalent to the [`AsyncCall::call`].
1138    pub async fn call(self) -> Result<CallResponse<()>, AgentError> {
1139        self.build()?.call().await
1140    }
1141
1142    /// Make a call. This is equivalent to the [`AsyncCall::call_and_wait`].
1143    pub async fn call_and_wait(self) -> Result<(), AgentError> {
1144        self.build()?.call_and_wait().await
1145    }
1146}
1147
1148#[cfg_attr(target_family = "wasm", async_trait(?Send))]
1149#[cfg_attr(not(target_family = "wasm"), async_trait)]
1150impl<'agent, 'canister: 'agent> AsyncCall for UpdateSettingsBuilder<'agent, 'canister> {
1151    type Value = ();
1152    async fn call(self) -> Result<CallResponse<()>, AgentError> {
1153        self.build()?.call().await
1154    }
1155
1156    async fn call_and_wait(self) -> Result<(), AgentError> {
1157        self.build()?.call_and_wait().await
1158    }
1159}
1160
1161impl<'agent, 'canister: 'agent> IntoFuture for UpdateSettingsBuilder<'agent, 'canister> {
1162    type IntoFuture = CallFuture<'agent, ()>;
1163    type Output = Result<(), AgentError>;
1164    fn into_future(self) -> Self::IntoFuture {
1165        AsyncCall::call_and_wait(self)
1166    }
1167}
1168
1169#[cfg(not(target_family = "wasm"))]
1170type BoxStream<'a, T> = Pin<Box<dyn Stream<Item = T> + Send + 'a>>;
1171#[cfg(target_family = "wasm")]
1172type BoxStream<'a, T> = Pin<Box<dyn Stream<Item = T> + 'a>>;
ic_utils/interfaces/management_canister/builders.rs

ic_utils/interfaces/management_canister/
builders.rs
