sui_sdk/
wallet_context.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::sui_client_config::{SuiClientConfig, SuiEnv};
5use anyhow::{anyhow, ensure};
6use futures::future;
7use futures::stream::TryStreamExt;
8use shared_crypto::intent::Intent;
9use std::collections::BTreeSet;
10use std::path::{Path, PathBuf};
11use sui_config::{Config, PersistedConfig};
12use sui_keys::key_identity::KeyIdentity;
13use sui_keys::keystore::{AccountKeystore, Alias, Keystore};
14use sui_rpc_api::client::ExecutedTransaction;
15use sui_types::base_types::{FullObjectRef, ObjectID, ObjectRef, SuiAddress};
16use sui_types::crypto::{Signature, SuiKeyPair};
17use sui_types::effects::TransactionEffectsAPI;
18use sui_types::object::Object;
19
20use std::sync::OnceLock;
21use sui_rpc_api::Client;
22use sui_types::gas_coin::GasCoin;
23use sui_types::transaction::{Transaction, TransactionData, TransactionDataAPI};
24use tracing::info;
25
26pub struct WalletContext {
27    pub config: PersistedConfig<SuiClientConfig>,
28    request_timeout: Option<std::time::Duration>,
29    grpc: OnceLock<Client>,
30    max_concurrent_requests: Option<u64>,
31    env_override: Option<String>,
32}
33
34impl WalletContext {
35    pub fn new(config_path: &Path) -> Result<Self, anyhow::Error> {
36        let config: SuiClientConfig = PersistedConfig::read(config_path).map_err(|err| {
37            anyhow!(
38                "Cannot open wallet config file at {:?}. Err: {err}",
39                config_path
40            )
41        })?;
42
43        let config = config.persisted(config_path);
44        let context = Self {
45            config,
46            request_timeout: None,
47            grpc: OnceLock::new(),
48            max_concurrent_requests: None,
49            env_override: None,
50        };
51        Ok(context)
52    }
53
54    pub fn new_for_tests(
55        keystore: Keystore,
56        external: Option<Keystore>,
57        path: Option<PathBuf>,
58    ) -> Self {
59        let mut config = SuiClientConfig::new(keystore)
60            .persisted(&path.unwrap_or(PathBuf::from("test_config.yaml")));
61        config.external_keys = external;
62        Self {
63            config,
64            request_timeout: None,
65            grpc: OnceLock::new(),
66            max_concurrent_requests: None,
67            env_override: None,
68        }
69    }
70
71    pub fn with_request_timeout(mut self, request_timeout: std::time::Duration) -> Self {
72        self.request_timeout = Some(request_timeout);
73        self
74    }
75
76    pub fn with_max_concurrent_requests(mut self, max_concurrent_requests: u64) -> Self {
77        self.max_concurrent_requests = Some(max_concurrent_requests);
78        self
79    }
80
81    pub fn with_env_override(mut self, env_override: String) -> Self {
82        self.env_override = Some(env_override);
83        self
84    }
85
86    pub fn get_addresses(&self) -> Vec<SuiAddress> {
87        let mut addresses = self.config.keystore.addresses();
88
89        if let Some(external_keys) = &self.config.external_keys {
90            addresses.extend(external_keys.addresses());
91        }
92
93        addresses
94    }
95
96    pub fn addresses_with_alias(&self) -> Vec<(&SuiAddress, &Alias)> {
97        let mut addresses = self.config.keystore.addresses_with_alias();
98
99        if let Some(external_keys) = &self.config.external_keys {
100            addresses.extend(external_keys.addresses_with_alias());
101        }
102
103        addresses
104    }
105
106    pub fn get_env_override(&self) -> Option<String> {
107        self.env_override.clone()
108    }
109
110    pub fn get_identity_address(
111        &mut self,
112        input: Option<KeyIdentity>,
113    ) -> Result<SuiAddress, anyhow::Error> {
114        if let Some(key_identity) = input {
115            if let Ok(address) = self.config.keystore.get_by_identity(&key_identity) {
116                return Ok(address);
117            }
118            if let Some(address) = self
119                .config
120                .external_keys
121                .as_ref()
122                .and_then(|external_keys| external_keys.get_by_identity(&key_identity).ok())
123            {
124                return Ok(address);
125            }
126
127            Err(anyhow!(
128                "No address found for the provided key identity: {key_identity}"
129            ))
130        } else {
131            self.active_address()
132        }
133    }
134
135    pub fn grpc_client(&self) -> Result<Client, anyhow::Error> {
136        if let Some(client) = self.grpc.get() {
137            Ok(client.clone())
138        } else {
139            let client = self.get_active_env()?.create_grpc_client()?;
140            Ok(self.grpc.get_or_init(move || client).clone())
141        }
142    }
143
144    /// Load the chain ID corresponding to the active environment, or fetch and cache it if not
145    /// present.
146    ///
147    /// The chain ID is cached in the `client.yaml` file to avoid redundant network requests.
148    pub async fn load_or_cache_chain_id(&self) -> Result<String, anyhow::Error> {
149        self.internal_load_or_cache_chain_id(false).await
150    }
151
152    /// Try to load the cached chain ID for the active environment.
153    pub async fn try_load_chain_id_from_cache(
154        &self,
155        env: Option<String>,
156    ) -> Result<String, anyhow::Error> {
157        let env = if let Some(env) = env {
158            self.config
159                .get_env(&Some(env.to_string()))
160                .ok_or_else(|| anyhow!("Environment configuration not found for env [{}]", env))?
161        } else {
162            self.get_active_env()?
163        };
164        if let Some(chain_id) = &env.chain_id {
165            Ok(chain_id.clone())
166        } else {
167            Err(anyhow!(
168                "No cached chain ID found for env {}. Please pass `-e env_name` to your command",
169                env.alias
170            ))
171        }
172    }
173
174    /// Cache (or recache) chain ID for the active environment by fetching it from the
175    /// network
176    pub async fn cache_chain_id(&self) -> Result<String, anyhow::Error> {
177        self.internal_load_or_cache_chain_id(true).await
178    }
179
180    async fn internal_load_or_cache_chain_id(
181        &self,
182        force_recache: bool,
183    ) -> Result<String, anyhow::Error> {
184        let env = self.get_active_env()?;
185        if !force_recache && env.chain_id.is_some() {
186            let chain_id = env.chain_id.as_ref().unwrap();
187            info!("Found cached chain ID for env {}: {}", env.alias, chain_id);
188            return Ok(chain_id.clone());
189        }
190        let chain_id = self.grpc_client()?.get_chain_identifier().await?;
191        let path = self.config.path();
192        let mut config_result = SuiClientConfig::load_with_lock(path)?;
193
194        config_result.update_env_chain_id(&env.alias, chain_id.to_string())?;
195        config_result.save_with_lock(path)?;
196        Ok(chain_id.to_string())
197    }
198
199    pub fn get_active_env(&self) -> Result<&SuiEnv, anyhow::Error> {
200        if self.env_override.is_some() {
201            self.config.get_env(&self.env_override).ok_or_else(|| {
202                anyhow!(
203                    "Environment configuration not found for env [{}]",
204                    self.env_override.as_deref().unwrap_or("None")
205                )
206            })
207        } else {
208            self.config.get_active_env()
209        }
210    }
211
212    // TODO: Ger rid of mut
213    pub fn active_address(&mut self) -> Result<SuiAddress, anyhow::Error> {
214        if self.config.keystore.entries().is_empty() {
215            return Err(anyhow!(
216                "No managed addresses. Create new address with `new-address` command."
217            ));
218        }
219
220        // Ok to unwrap because we checked that config addresses not empty
221        // Set it if not exists
222        self.config.active_address = Some(
223            self.config
224                .active_address
225                .unwrap_or(*self.config.keystore.addresses().first().unwrap()),
226        );
227
228        Ok(self.config.active_address.unwrap())
229    }
230
231    /// Get the latest object reference given a object id
232    pub async fn get_object_ref(&self, object_id: ObjectID) -> Result<ObjectRef, anyhow::Error> {
233        Ok(self
234            .grpc_client()?
235            .get_object(object_id)
236            .await?
237            .compute_object_reference())
238    }
239
240    /// Get the latest full object reference given a object id
241    pub async fn get_full_object_ref(
242        &self,
243        object_id: ObjectID,
244    ) -> Result<FullObjectRef, anyhow::Error> {
245        Ok(self
246            .grpc_client()?
247            .get_object(object_id)
248            .await?
249            .compute_full_object_reference())
250    }
251
252    /// Get all the gas objects (and conveniently, gas amounts) for the address
253    pub async fn gas_objects(
254        &self,
255        owner: SuiAddress,
256    ) -> Result<Vec<(u64, Object)>, anyhow::Error> {
257        let client = self.grpc_client()?;
258
259        client
260            .list_owned_objects(owner, Some(GasCoin::type_()))
261            .map_err(Into::into)
262            .and_then(|object| async move {
263                let gas_coin = GasCoin::try_from(&object)?;
264
265                Ok((gas_coin.value(), object))
266            })
267            .try_collect()
268            .await
269    }
270
271    pub async fn get_object_owner(&self, id: &ObjectID) -> Result<SuiAddress, anyhow::Error> {
272        self.grpc_client()?
273            .get_object(*id)
274            .await?
275            .owner()
276            .get_owner_address()
277            .map_err(Into::into)
278    }
279
280    pub async fn try_get_object_owner(
281        &self,
282        id: &Option<ObjectID>,
283    ) -> Result<Option<SuiAddress>, anyhow::Error> {
284        if let Some(id) = id {
285            Ok(Some(self.get_object_owner(id).await?))
286        } else {
287            Ok(None)
288        }
289    }
290
291    /// Infer the sender of a transaction based on the gas objects provided. If no gas objects are
292    /// provided, assume the active address is the sender.
293    pub async fn infer_sender(&mut self, gas: &[ObjectID]) -> Result<SuiAddress, anyhow::Error> {
294        if gas.is_empty() {
295            return self.active_address();
296        }
297
298        // Find the owners of all supplied object IDs
299        let owners = future::try_join_all(gas.iter().map(|id| self.get_object_owner(id))).await?;
300
301        // SAFETY `gas` is non-empty.
302        let owner = owners.first().copied().unwrap();
303
304        ensure!(
305            owners.iter().all(|o| o == &owner),
306            "Cannot infer sender, not all gas objects have the same owner."
307        );
308
309        Ok(owner)
310    }
311
312    /// Find a gas object which fits the budget
313    pub async fn gas_for_owner_budget(
314        &self,
315        address: SuiAddress,
316        budget: u64,
317        forbidden_gas_objects: BTreeSet<ObjectID>,
318    ) -> Result<(u64, Object), anyhow::Error> {
319        for o in self.gas_objects(address).await? {
320            if o.0 >= budget && !forbidden_gas_objects.contains(&o.1.id()) {
321                return Ok((o.0, o.1));
322            }
323        }
324        Err(anyhow!(
325            "No non-argument gas objects found for this address with value >= budget {budget}. Run sui client gas to check for gas objects."
326        ))
327    }
328
329    pub async fn get_all_gas_objects_owned_by_address(
330        &self,
331        address: SuiAddress,
332    ) -> anyhow::Result<Vec<ObjectRef>> {
333        self.get_gas_objects_owned_by_address(address, None).await
334    }
335
336    pub async fn get_gas_objects_owned_by_address(
337        &self,
338        owner: SuiAddress,
339        page_size: Option<u32>,
340    ) -> anyhow::Result<Vec<ObjectRef>> {
341        let page = self
342            .grpc_client()?
343            .get_owned_objects(owner, Some(GasCoin::type_()), page_size, None)
344            .await?;
345
346        Ok(page
347            .items
348            .into_iter()
349            .map(|o| o.compute_object_reference())
350            .collect())
351    }
352
353    /// Given an address, return one gas object owned by this address.
354    /// The actual implementation just returns the first one returned by the read api.
355    pub async fn get_one_gas_object_owned_by_address(
356        &self,
357        address: SuiAddress,
358    ) -> anyhow::Result<Option<ObjectRef>> {
359        Ok(self
360            .get_gas_objects_owned_by_address(address, Some(1))
361            .await?
362            .pop())
363    }
364
365    /// Returns one address and all gas objects owned by that address.
366    pub async fn get_one_account(&self) -> anyhow::Result<(SuiAddress, Vec<ObjectRef>)> {
367        let address = self.get_addresses().pop().unwrap();
368        Ok((
369            address,
370            self.get_all_gas_objects_owned_by_address(address).await?,
371        ))
372    }
373
374    /// Return a gas object owned by an arbitrary address managed by the wallet.
375    pub async fn get_one_gas_object(&self) -> anyhow::Result<Option<(SuiAddress, ObjectRef)>> {
376        for address in self.get_addresses() {
377            if let Some(gas_object) = self.get_one_gas_object_owned_by_address(address).await? {
378                return Ok(Some((address, gas_object)));
379            }
380        }
381        Ok(None)
382    }
383
384    /// Returns all the account addresses managed by the wallet and their owned gas objects.
385    pub async fn get_all_accounts_and_gas_objects(
386        &self,
387    ) -> anyhow::Result<Vec<(SuiAddress, Vec<ObjectRef>)>> {
388        let mut result = vec![];
389        for address in self.get_addresses() {
390            let objects = self
391                .gas_objects(address)
392                .await?
393                .into_iter()
394                .map(|(_, o)| o.compute_object_reference())
395                .collect();
396            result.push((address, objects));
397        }
398        Ok(result)
399    }
400
401    pub async fn get_reference_gas_price(&self) -> Result<u64, anyhow::Error> {
402        self.grpc_client()?
403            .get_reference_gas_price()
404            .await
405            .map_err(Into::into)
406    }
407
408    /// Add an account
409    pub async fn add_account(&mut self, alias: Option<String>, keypair: SuiKeyPair) {
410        self.config.keystore.import(alias, keypair).await.unwrap();
411    }
412
413    pub fn get_keystore_by_identity(
414        &self,
415        key_identity: &KeyIdentity,
416    ) -> Result<&Keystore, anyhow::Error> {
417        if self
418            .config
419            .keystore
420            .get_by_identity(key_identity)
421            .map(|address| self.config.keystore.addresses().contains(&address))
422            .unwrap_or(false)
423        {
424            return Ok(&self.config.keystore);
425        }
426
427        if let Some(external_keys) = self.config.external_keys.as_ref()
428            && external_keys
429                .get_by_identity(key_identity)
430                .map(|address| external_keys.addresses().contains(&address))
431                .unwrap_or(false)
432        {
433            return Ok(external_keys);
434        }
435
436        Err(anyhow!(
437            "No keystore found for the provided key identity: {key_identity}"
438        ))
439    }
440
441    pub fn get_keystore_by_identity_mut(
442        &mut self,
443        key_identity: &KeyIdentity,
444    ) -> Result<&mut Keystore, anyhow::Error> {
445        if self.config.keystore.get_by_identity(key_identity).is_ok() {
446            return Ok(&mut self.config.keystore);
447        }
448
449        if let Some(external_keys) = self.config.external_keys.as_mut()
450            && external_keys.get_by_identity(key_identity).is_ok()
451        {
452            return Ok(external_keys);
453        }
454
455        Err(anyhow!(
456            "No keystore found for the provided key identity: {key_identity}"
457        ))
458    }
459
460    pub async fn sign_secure(
461        &self,
462        key_identity: &KeyIdentity,
463        data: &TransactionData,
464        intent: Intent,
465    ) -> Result<Signature, anyhow::Error> {
466        let keystore = self.get_keystore_by_identity(key_identity)?;
467        let address = keystore.get_by_identity(key_identity).map_err(|_| {
468            anyhow!("No address found for the provided key identity: {key_identity}")
469        })?;
470        let sig = keystore.sign_secure(&address, data, intent).await?;
471        Ok(sig)
472    }
473
474    /// Sign a transaction with a key currently managed by the WalletContext
475    pub async fn sign_transaction(&self, data: &TransactionData) -> Transaction {
476        let sig = self
477            .config
478            .keystore
479            .sign_secure(&data.sender(), data, Intent::sui_transaction())
480            .await
481            .unwrap();
482        // TODO: To support sponsored transaction, we should also look at the gas owner.
483        Transaction::from_data(data.clone(), vec![sig])
484    }
485
486    /// Execute a transaction and wait for it to be locally executed on the fullnode.
487    /// Also expects the effects status to be ExecutionStatus::Success.
488    pub async fn execute_transaction_must_succeed(&self, tx: Transaction) -> ExecutedTransaction {
489        tracing::debug!("Executing transaction: {:?}", tx);
490        let response = self.execute_transaction_may_fail(tx).await.unwrap();
491        assert!(
492            response.effects.status().is_ok(),
493            "Transaction failed: {:?}",
494            response
495        );
496        response
497    }
498
499    /// Execute a transaction and wait for it to be locally executed on the fullnode.
500    /// The transaction execution is not guaranteed to succeed and may fail. This is usually only
501    /// needed in non-test environment or the caller is explicitly testing some failure behavior.
502    pub async fn execute_transaction_may_fail(
503        &self,
504        tx: Transaction,
505    ) -> anyhow::Result<ExecutedTransaction> {
506        self.grpc_client()?
507            .execute_transaction_and_wait_for_checkpoint(&tx)
508            .await
509            .map_err(Into::into)
510    }
511}
sui_sdk/wallet_context.rs

sui_sdk/
wallet_context.rs
