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

iota_sdk/
wallet_context.rs
