iota_sdk/
wallet_context.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{collections::BTreeSet, path::Path, sync::Arc};
6
7use anyhow::{anyhow, bail, ensure};
8use colored::Colorize;
9use futures::{StreamExt, TryStreamExt, future};
10use getset::{Getters, MutGetters};
11use iota_config::{Config, PersistedConfig};
12use iota_json_rpc_types::{
13    IotaObjectData, IotaObjectDataFilter, IotaObjectDataOptions, IotaObjectResponseQuery,
14    IotaTransactionBlockResponse, IotaTransactionBlockResponseOptions,
15};
16use iota_keys::keystore::{AccountKeystore, Keystore};
17use iota_sdk_crypto::simple::SimpleKeypair;
18use iota_sdk_types::{Address, ObjectId, ObjectReference, StructTag, Transaction, crypto::Intent};
19use iota_types::{
20    gas_coin::GasCoin,
21    transaction::{TransactionAPI, TransactionEnvelope},
22};
23use tokio::sync::RwLock;
24use tracing::warn;
25
26use crate::{
27    IotaClient, PagedFn,
28    iota_client_config::{IotaClientConfig, IotaEnv},
29};
30
31/// Wallet for managing accounts, objects, and interact with client APIs.
32// Mainly used in the CLI and tests.
33#[derive(Getters, MutGetters)]
34#[getset(get = "pub", get_mut = "pub")]
35pub struct WalletContext {
36    config: PersistedConfig<IotaClientConfig>,
37    request_timeout: Option<std::time::Duration>,
38    client: Arc<RwLock<Option<IotaClient>>>,
39    grpc_client: Arc<RwLock<Option<iota_grpc_client::Client>>>,
40    max_concurrent_requests: Option<u64>,
41    env_override: Option<String>,
42}
43
44impl WalletContext {
45    /// Create a new [`WalletContext`] with the config path to an existing
46    /// [`IotaClientConfig`] and optional parameters for the client.
47    pub fn new(config_path: &Path) -> Result<Self, anyhow::Error> {
48        let config: IotaClientConfig = PersistedConfig::read(config_path).map_err(|err| {
49            anyhow!("Cannot open wallet config file at {config_path:?}. Err: {err}")
50        })?;
51
52        if let Some(active_address) = &config.active_address {
53            let addresses = match &config.keystore {
54                Keystore::File(file) => file.addresses(),
55                Keystore::InMem(mem) => mem.addresses(),
56            };
57            ensure!(
58                addresses.contains(active_address),
59                "error in '{}': active address not found in the keystore",
60                config_path.display()
61            );
62        }
63
64        if let Some(active_env) = &config.active_env {
65            ensure!(
66                config.get_env(active_env).is_some(),
67                "error in '{}': active environment not found in the envs list",
68                config_path.display()
69            );
70        }
71
72        let config = config.persisted(config_path);
73        let context = Self {
74            config,
75            request_timeout: None,
76            client: Default::default(),
77            grpc_client: Default::default(),
78            max_concurrent_requests: None,
79            env_override: None,
80        };
81        Ok(context)
82    }
83
84    pub fn with_request_timeout(mut self, request_timeout: std::time::Duration) -> Self {
85        self.request_timeout = Some(request_timeout);
86        self
87    }
88
89    pub fn with_max_concurrent_requests(mut self, max_concurrent_requests: u64) -> Self {
90        self.max_concurrent_requests = Some(max_concurrent_requests);
91        self
92    }
93
94    pub fn with_env_override(mut self, env_override: String) -> Self {
95        self.env_override = Some(env_override);
96        self
97    }
98
99    /// Get all addresses from the keystore.
100    pub fn get_addresses(&self) -> Vec<Address> {
101        self.config.keystore.addresses()
102    }
103
104    pub fn get_env_override(&self) -> Option<String> {
105        self.env_override.clone()
106    }
107
108    /// Get the configured [`IotaClient`].
109    pub async fn get_client(&self) -> Result<IotaClient, anyhow::Error> {
110        let read = self.client.read().await;
111
112        Ok(if let Some(client) = read.as_ref() {
113            client.clone()
114        } else {
115            drop(read);
116            let client = self
117                .active_env()?
118                .create_rpc_client(self.request_timeout, self.max_concurrent_requests)
119                .await?;
120            if let Err(e) = client.check_api_version() {
121                warn!("{e}");
122                eprintln!("{}", format!("[warn] {e}").yellow().bold());
123            }
124            self.client.write().await.insert(client).clone()
125        })
126    }
127
128    /// Get the configured gRPC client, creating and caching it on first use.
129    /// Errors if the active env has no `grpc` URL configured.
130    pub async fn get_grpc_client(&self) -> Result<iota_grpc_client::Client, anyhow::Error> {
131        let read = self.grpc_client.read().await;
132
133        Ok(if let Some(client) = read.as_ref() {
134            client.clone()
135        } else {
136            drop(read);
137            let client = self.active_env()?.create_grpc_client()?;
138            self.grpc_client.write().await.insert(client).clone()
139        })
140    }
141
142    /// Get the active [`Address`].
143    /// If not set, defaults to the first address in the keystore.
144    pub fn active_address(&self) -> Result<Address, anyhow::Error> {
145        if self.config.keystore.addresses().is_empty() {
146            bail!("No managed addresses. Create new address with the `new-address` command.");
147        }
148
149        Ok(if let Some(addr) = self.config.active_address() {
150            *addr
151        } else {
152            self.config.keystore().addresses()[0]
153        })
154    }
155
156    /// Get the active [`IotaEnv`].
157    /// If not set, defaults to the first environment in the config.
158    pub fn active_env(&self) -> Result<&IotaEnv, anyhow::Error> {
159        if self.config.envs.is_empty() {
160            bail!("No managed environments. Create new environment with the `new-env` command.");
161        }
162
163        if let Some(env_override) = &self.env_override {
164            self.config.get_env(env_override).ok_or_else(|| {
165                anyhow!("Environment configuration not found for env [{env_override}]")
166            })
167        } else {
168            Ok(if self.config.active_env().is_some() {
169                self.config.get_active_env()?
170            } else {
171                &self.config.envs()[0]
172            })
173        }
174    }
175
176    /// Get the latest object reference given a object id.
177    pub async fn get_object_ref(
178        &self,
179        object_id: ObjectId,
180    ) -> Result<ObjectReference, anyhow::Error> {
181        let client = self.get_client().await?;
182        Ok(client
183            .read_api()
184            .get_object_with_options(object_id, IotaObjectDataOptions::new())
185            .await?
186            .into_object()?
187            .object_ref())
188    }
189
190    /// Get all the gas objects (and conveniently, gas amounts) for the address.
191    pub async fn gas_objects(
192        &self,
193        address: Address,
194    ) -> Result<Vec<(u64, IotaObjectData)>, anyhow::Error> {
195        let client = self.get_client().await?;
196
197        let values_objects = PagedFn::stream(async |cursor| {
198            client
199                .read_api()
200                .get_owned_objects(
201                    address,
202                    IotaObjectResponseQuery::new(
203                        Some(IotaObjectDataFilter::StructType(StructTag::new_gas_coin())),
204                        Some(IotaObjectDataOptions::full_content()),
205                    ),
206                    cursor,
207                    None,
208                )
209                .await
210        })
211        .filter_map(|res| async {
212            match res {
213                Ok(res) => {
214                    if let Some(o) = res.data {
215                        match GasCoin::try_from(&o) {
216                            Ok(gas_coin) => Some(Ok((gas_coin.value(), o))),
217                            Err(e) => Some(Err(anyhow!("{e}"))),
218                        }
219                    } else {
220                        None
221                    }
222                }
223                Err(e) => Some(Err(anyhow!("{e}"))),
224            }
225        })
226        .try_collect::<Vec<_>>()
227        .await?;
228
229        Ok(values_objects)
230    }
231
232    /// Get the address that owns the object of the provided [`ObjectId`].
233    pub async fn get_object_owner(&self, id: &ObjectId) -> Result<Address, anyhow::Error> {
234        let client = self.get_client().await?;
235        let object = client
236            .read_api()
237            .get_object_with_options(*id, IotaObjectDataOptions::new().with_owner())
238            .await?
239            .into_object()?;
240        Ok(*object
241            .owner
242            .ok_or_else(|| anyhow!("Owner field is None"))?
243            .address_or_object()
244            .ok_or_else(|| anyhow::anyhow!("not an address or object owner"))?)
245    }
246
247    /// Get the address that owns the object, if an [`ObjectId`] is provided.
248    pub async fn try_get_object_owner(
249        &self,
250        id: &Option<ObjectId>,
251    ) -> Result<Option<Address>, anyhow::Error> {
252        if let Some(id) = id {
253            Ok(Some(self.get_object_owner(id).await?))
254        } else {
255            Ok(None)
256        }
257    }
258
259    /// Infer the sender of a transaction based on the gas objects provided. If
260    /// no gas objects are provided, assume the active address is the
261    /// sender.
262    pub async fn infer_sender(&mut self, gas: &[ObjectId]) -> Result<Address, anyhow::Error> {
263        if gas.is_empty() {
264            return self.active_address();
265        }
266
267        // Find the owners of all supplied object IDs
268        let owners = future::try_join_all(gas.iter().map(|id| self.get_object_owner(id))).await?;
269
270        // SAFETY `gas` is non-empty.
271        let owner = owners[0];
272
273        ensure!(
274            owners.iter().all(|o| o == &owner),
275            "Cannot infer sender, not all gas objects have the same owner."
276        );
277
278        Ok(owner)
279    }
280
281    /// Find a gas object which fits the budget.
282    pub async fn gas_for_owner_budget(
283        &self,
284        address: Address,
285        budget: u64,
286        forbidden_gas_objects: BTreeSet<ObjectId>,
287    ) -> Result<(u64, IotaObjectData), anyhow::Error> {
288        for o in self.gas_objects(address).await? {
289            if o.0 >= budget && !forbidden_gas_objects.contains(&o.1.object_id) {
290                return Ok((o.0, o.1));
291            }
292        }
293        bail!(
294            "No non-argument gas objects found for this address with value >= budget {budget}. Run iota client gas to check for gas objects."
295        )
296    }
297
298    /// Get the [`ObjectReference`] for gas objects owned by the provided
299    /// address. Maximum is RPC_QUERY_MAX_RESULT_LIMIT (50 by default).
300    pub async fn get_all_gas_objects_owned_by_address(
301        &self,
302        address: Address,
303    ) -> anyhow::Result<Vec<ObjectReference>> {
304        self.get_gas_objects_owned_by_address(address, None).await
305    }
306
307    /// Get a limited amount of [`ObjectReference`]s for gas objects owned by
308    /// the provided address. Max limit is RPC_QUERY_MAX_RESULT_LIMIT (50 by
309    /// default).
310    pub async fn get_gas_objects_owned_by_address(
311        &self,
312        address: Address,
313        limit: impl Into<Option<usize>>,
314    ) -> anyhow::Result<Vec<ObjectReference>> {
315        let client = self.get_client().await?;
316        let results: Vec<_> = client
317            .read_api()
318            .get_owned_objects(
319                address,
320                IotaObjectResponseQuery::new(
321                    Some(IotaObjectDataFilter::StructType(StructTag::new_gas_coin())),
322                    Some(IotaObjectDataOptions::full_content()),
323                ),
324                None,
325                limit,
326            )
327            .await?
328            .data
329            .into_iter()
330            .filter_map(|r| r.data.map(|o| o.object_ref()))
331            .collect();
332        Ok(results)
333    }
334
335    /// Given an address, return one gas object owned by this address.
336    /// The actual implementation just returns the first one returned by the
337    /// read api.
338    pub async fn get_one_gas_object_owned_by_address(
339        &self,
340        address: Address,
341    ) -> anyhow::Result<Option<ObjectReference>> {
342        Ok(self
343            .get_gas_objects_owned_by_address(address, 1)
344            .await?
345            .pop())
346    }
347
348    /// Return one address and all gas objects owned by that address.
349    pub async fn get_one_account(&self) -> anyhow::Result<(Address, Vec<ObjectReference>)> {
350        let address = self.get_addresses().pop().unwrap();
351        Ok((
352            address,
353            self.get_all_gas_objects_owned_by_address(address).await?,
354        ))
355    }
356
357    /// Return a gas object owned by an arbitrary address managed by the wallet.
358    pub async fn get_one_gas_object(&self) -> anyhow::Result<Option<(Address, ObjectReference)>> {
359        for address in self.get_addresses() {
360            if let Some(gas_object) = self.get_one_gas_object_owned_by_address(address).await? {
361                return Ok(Some((address, gas_object)));
362            }
363        }
364        Ok(None)
365    }
366
367    /// Return all the account addresses managed by the wallet and their owned
368    /// gas objects.
369    pub async fn get_all_accounts_and_gas_objects(
370        &self,
371    ) -> anyhow::Result<Vec<(Address, Vec<ObjectReference>)>> {
372        let mut result = vec![];
373        for address in self.get_addresses() {
374            let objects = self
375                .gas_objects(address)
376                .await?
377                .into_iter()
378                .map(|(_, o)| o.object_ref())
379                .collect();
380            result.push((address, objects));
381        }
382        Ok(result)
383    }
384
385    pub async fn get_reference_gas_price(&self) -> Result<u64, anyhow::Error> {
386        let client = self.get_client().await?;
387        let gas_price = client.governance_api().get_reference_gas_price().await?;
388        Ok(gas_price)
389    }
390
391    /// Add an account.
392    pub fn add_account(&mut self, alias: impl Into<Option<String>>, keypair: SimpleKeypair) {
393        self.config.keystore.add_key(alias.into(), keypair).unwrap();
394    }
395
396    /// Sign a transaction with a key currently managed by the WalletContext.
397    pub fn sign_transaction(&self, tx: &Transaction) -> TransactionEnvelope {
398        let sig = self
399            .config
400            .keystore
401            .sign_secure(&tx.sender(), tx, Intent::iota_transaction())
402            .unwrap();
403        // TODO: To support sponsored transaction, we should also look at the gas owner.
404        TransactionEnvelope::from_data(tx.clone(), vec![sig])
405    }
406
407    /// Execute a transaction and wait for it to be locally executed on the
408    /// fullnode. Also expects the effects status to be
409    /// ExecutionStatus::Success.
410    pub async fn execute_transaction_must_succeed(
411        &self,
412        tx: TransactionEnvelope,
413    ) -> IotaTransactionBlockResponse {
414        tracing::debug!("Executing transaction: {:?}", tx);
415        let response = self.execute_transaction_may_fail(tx).await.unwrap();
416        assert!(
417            response.status_ok().unwrap(),
418            "Transaction failed: {response:?}"
419        );
420        response
421    }
422
423    /// Execute a transaction and wait for it to be locally executed on the
424    /// fullnode. The transaction execution is not guaranteed to succeed and
425    /// may fail. This is usually only needed in non-test environment or the
426    /// caller is explicitly testing some failure behavior.
427    pub async fn execute_transaction_may_fail(
428        &self,
429        tx: TransactionEnvelope,
430    ) -> anyhow::Result<IotaTransactionBlockResponse> {
431        let client = self.get_client().await?;
432        Ok(client
433            .quorum_driver_api()
434            .execute_transaction_block(
435                tx,
436                IotaTransactionBlockResponseOptions::new()
437                    .with_effects()
438                    .with_input()
439                    .with_events()
440                    .with_object_changes()
441                    .with_balance_changes(),
442                iota_types::quorum_driver_types::ExecuteTransactionRequestType::WaitForLocalExecution,
443            )
444            .await?)
445    }
446}
iota_sdk/wallet_context.rs

iota_sdk/
wallet_context.rs
