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

iota_sdk/
wallet_context.rs
