From d1920bbd098ef2d84c085fda5d59e4dd363c24c7 Mon Sep 17 00:00:00 2001 From: Severin Siffert Date: Fri, 17 May 2024 13:25:31 +0200 Subject: [PATCH] first version --- ic-utils/src/interfaces/wallet.rs | 131 +++++++++++++++++++++++++++++- ref-tests/tests/integration.rs | 14 +++- 2 files changed, 140 insertions(+), 5 deletions(-) diff --git a/ic-utils/src/interfaces/wallet.rs b/ic-utils/src/interfaces/wallet.rs index 3a5c0733..e39bac8b 100644 --- a/ic-utils/src/interfaces/wallet.rs +++ b/ic-utils/src/interfaces/wallet.rs @@ -14,7 +14,7 @@ use crate::{ Canister, }; use async_trait::async_trait; -use candid::{decode_args, utils::ArgumentDecoder, CandidType, Deserialize, Nat}; +use candid::{decode_args, utils::ArgumentDecoder, CandidType, Decode, Deserialize, Nat}; use ic_agent::{export::Principal, Agent, AgentError, RequestId}; use once_cell::sync::Lazy; use semver::{Version, VersionReq}; @@ -37,6 +37,19 @@ where phantom_out: std::marker::PhantomData, } +/// An interface for forwarding a canister method call through the wallet canister via `wallet_canister_call_with_max_cycles`. +#[derive(Debug)] +pub struct MaxCyclesCallForwarder<'agent, 'canister, Out> +where + Out: CandidType + for<'de> candid::Deserialize<'de> + Send + Sync, +{ + wallet: &'canister WalletCanister<'agent>, + destination: Principal, + method_name: String, + arg: Argument, + phantom_out: std::marker::PhantomData, +} + /// A canister's settings. Similar to the canister settings struct from [`management_canister`](super::management_canister), /// but the management canister may evolve to have more settings without the wallet canister evolving to recognize them. #[derive(Debug, Clone, CandidType, Deserialize)] @@ -145,6 +158,92 @@ where } } +impl<'agent: 'canister, 'canister, Out> MaxCyclesCallForwarder<'agent, 'canister, Out> +where + Out: CandidType + for<'de> candid::Deserialize<'de> + Send + Sync, +{ + /// Set the argument with candid argument. Can be called at most once. + pub fn with_arg(mut self, arg: Argument) -> Self + where + Argument: CandidType + Sync + Send, + { + self.arg.set_idl_arg(arg); + self + } + /// Set the argument with multiple arguments as tuple. Can be called at most once. + pub fn with_args(mut self, tuple: impl candid::utils::ArgumentEncoder) -> Self { + if self.arg.0.is_some() { + panic!("argument is being set more than once"); + } + self.arg = Argument::from_candid(tuple); + self + } + + /// Set the argument with raw argument bytes. Can be called at most once. + pub fn with_arg_raw(mut self, arg: Vec) -> Self { + self.arg.set_raw_arg(arg); + self + } + + /// Creates an [`AsyncCall`] implementation that, when called, will forward the specified canister call. + pub fn build<'out>(self) -> Result, AgentError> + where + Out: 'out, + 'agent: 'out, + { + #[derive(CandidType, Deserialize)] + struct In { + canister: Principal, + method_name: String, + #[serde(with = "serde_bytes")] + args: Vec, + } + Ok(self + .wallet + .update("wallet_call_with_max_cycles") + .with_arg(In { + canister: self.destination, + method_name: self.method_name, + args: self.arg.serialize()?.to_vec(), + }) + .build() + .and_then( + |(result,): (Result,)| async move { + let result = result.map_err(AgentError::WalletCallFailed)?; + let out = Decode!(result.r#return.as_slice(), Out) + .map_err(|e| AgentError::CandidError(Box::new(e)))?; + Ok((out, result.attached_cycles)) + }, + )) + } + + /// Calls the forwarded canister call on the wallet canister. Equivalent to `.build().call()`. + pub async fn call(self) -> Result { + self.build()?.call().await + } + + /// Calls the forwarded canister call on the wallet canister, and waits for the result. Equivalent to `.build().call_and_wait()`. + pub async fn call_and_wait(self) -> Result<(Out, candid::Nat), AgentError> { + self.build()?.call_and_wait().await + } +} + +#[cfg_attr(target_family = "wasm", async_trait(?Send))] +#[cfg_attr(not(target_family = "wasm"), async_trait)] +impl<'agent: 'canister, 'canister, Out> AsyncCall<(Out, candid::Nat)> + for MaxCyclesCallForwarder<'agent, 'canister, Out> +where + Out: CandidType + for<'de> candid::Deserialize<'de> + Send + Sync, +{ + async fn call(self) -> Result { + self.call().await + } + + async fn call_and_wait(self) -> Result<(Out, candid::Nat), AgentError> { + self.call_and_wait().await + } +} + /// A wallet canister interface, for the standard wallet provided by DFINITY. /// This interface implements most methods conveniently for the user. #[derive(Debug, Clone)] @@ -414,6 +513,16 @@ pub struct CallResult { pub r#return: Vec, } +/// The result of a call forwarding request via `wallet_call_with_max_cycles`. +#[derive(Debug, Clone, CandidType, Deserialize)] +pub struct MaxCyclesCallResult { + /// The encoded return value blob of the canister method. + #[serde(with = "serde_bytes")] + pub r#return: Vec, + /// How many cycles the call was performed with. + pub attached_cycles: candid::Nat, +} + impl<'agent> WalletCanister<'agent> { /// Create an instance of a `WalletCanister` interface pointing to the given Canister ID. Fails if it cannot learn the wallet's version. pub async fn create( @@ -1038,6 +1147,26 @@ impl<'agent> WalletCanister<'agent> { } } + /// Forward a call to another canister, including an amount of cycles + /// from the wallet, using the 128-bit API. + pub fn call_with_max_cycles<'canister, Out, M: Into>( + &'canister self, + destination: Principal, + method_name: M, + arg: Argument, + ) -> MaxCyclesCallForwarder<'agent, 'canister, Out> + where + Out: CandidType + for<'de> candid::Deserialize<'de> + Send + Sync, + { + MaxCyclesCallForwarder { + wallet: self, + destination, + method_name: method_name.into(), + arg, + phantom_out: std::marker::PhantomData, + } + } + /// Forward a call to another canister, including an amount of cycles /// from the wallet. pub fn call<'canister, Out, M: Into>( diff --git a/ref-tests/tests/integration.rs b/ref-tests/tests/integration.rs index 663d432e..ff6be206 100644 --- a/ref-tests/tests/integration.rs +++ b/ref-tests/tests/integration.rs @@ -2,7 +2,7 @@ //! //! Contrary to ic-ref.rs, these tests are not meant to match any other tests. They're //! integration tests with a running IC-Ref. -use candid::CandidType; +use candid::{CandidType, Nat}; use ic_agent::{ agent::{agent_error::HttpErrorPayload, Envelope, EnvelopeContent, RejectCode, RejectResponse}, export::Principal, @@ -220,15 +220,21 @@ fn wallet_canister_forward() { .reply_data(b"DIDL\0\x01\x71\x0bHello World") .build(); - let args = Argument::from_raw(arg); - let (result,): (String,) = wallet - .call(universal_id, "update", args, 0) + .call(universal_id, "update", Argument::from_raw(arg.clone()), 0) .call_and_wait() .await .unwrap(); assert_eq!(result, "Hello World"); + + let (result_2, cycles): (String, Nat) = wallet + .call_with_max_cycles(universal_id, "update", Argument::from_raw(arg)) + .call_and_wait() + .await + .unwrap(); + assert_eq!(result_2, "Hello World"); + assert!(cycles > 0_u8); Ok(()) }); }