sui_faucet/
local_faucet.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::fmt;
5use std::sync::Arc;
6
7use anyhow::bail;
8use sui_rpc_api::client::ExecutedTransaction;
9use sui_sdk::types::effects::TransactionEffectsAPI;
10use tokio::sync::Mutex;
11use tokio::time::Duration;
12use tracing::info;
13
14use crate::FaucetConfig;
15use crate::FaucetError;
16
17use crate::CoinInfo;
18use shared_crypto::intent::Intent;
19use sui_keys::keystore::AccountKeystore;
20use sui_sdk::types::programmable_transaction_builder::ProgrammableTransactionBuilder;
21use sui_sdk::types::{
22    base_types::{ObjectID, SuiAddress},
23    gas_coin::GasCoin,
24    transaction::{Transaction, TransactionData},
25};
26use sui_sdk::wallet_context::WalletContext;
27
28const GAS_BUDGET: u64 = 10_000_000;
29const NUM_RETRIES: u8 = 2;
30
31pub struct LocalFaucet {
32    wallet: WalletContext,
33    active_address: SuiAddress,
34    coin_id: Arc<Mutex<ObjectID>>,
35    coin_amount: u64,
36    num_coins: usize,
37}
38
39/// We do not just derive(Debug) because WalletContext and the WriteAheadLog do not implement Debug / are also hard
40/// to implement Debug.
41impl fmt::Debug for LocalFaucet {
42    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
43        f.debug_struct("SimpleFaucet")
44            .field("faucet_wallet", &self.active_address)
45            .field("coin_amount", &self.coin_amount)
46            .finish()
47    }
48}
49
50impl LocalFaucet {
51    pub async fn new(
52        mut wallet: WalletContext,
53        config: FaucetConfig,
54    ) -> Result<Arc<Self>, FaucetError> {
55        let (coins, active_address) = find_gas_coins_and_address(&mut wallet, &config).await?;
56        info!("Starting faucet with address: {:?}", active_address);
57
58        Ok(Arc::new(LocalFaucet {
59            wallet,
60            active_address,
61            coin_id: Arc::new(Mutex::new(*coins[0].id())),
62            coin_amount: config.amount,
63            num_coins: config.num_coins,
64        }))
65    }
66
67    /// Make transaction and execute it.
68    pub async fn local_request_execute_tx(
69        &self,
70        recipient: SuiAddress,
71    ) -> Result<Vec<CoinInfo>, FaucetError> {
72        let gas_price = self
73            .wallet
74            .get_reference_gas_price()
75            .await
76            .map_err(|e| FaucetError::internal(format!("Failed to get gas price: {}", e)))?;
77
78        let mut ptb = ProgrammableTransactionBuilder::new();
79        let recipients = vec![recipient; self.num_coins];
80        let amounts = vec![self.coin_amount; self.num_coins];
81        ptb.pay_sui(recipients, amounts)
82            .map_err(FaucetError::internal)?;
83
84        let ptb = ptb.finish();
85
86        let coin_id = self.coin_id.lock().await;
87        let coin_id_ref = self
88            .wallet
89            .get_object_ref(*coin_id)
90            .await
91            .map_err(|e| FaucetError::internal(format!("Failed to get object ref: {}", e)))?;
92        let tx_data = TransactionData::new_programmable(
93            self.active_address,
94            vec![coin_id_ref],
95            ptb,
96            GAS_BUDGET,
97            gas_price,
98        );
99
100        let tx = self
101            .execute_txn_with_retries(tx_data, *coin_id, NUM_RETRIES)
102            .await
103            .map_err(FaucetError::internal)?;
104
105        let coins: Vec<CoinInfo> = tx
106            .effects
107            .created()
108            .into_iter()
109            .map(|o| CoinInfo {
110                amount: self.coin_amount,
111                id: o.0.0,
112                transfer_tx_digest: *tx.effects.transaction_digest(),
113            })
114            .collect();
115
116        Ok(coins)
117    }
118
119    async fn execute_txn(
120        &self,
121        tx_data: &TransactionData,
122        coin_id: ObjectID,
123    ) -> Result<ExecutedTransaction, anyhow::Error> {
124        let signature = self
125            .wallet
126            .config
127            .keystore
128            .sign_secure(&self.active_address, &tx_data, Intent::sui_transaction())
129            .await
130            .map_err(FaucetError::internal)?;
131        let tx = Transaction::from_data(tx_data.clone(), vec![signature]);
132
133        let client = self.wallet.grpc_client()?;
134
135        client
136            .execute_transaction_and_wait_for_checkpoint(&tx)
137            .await
138            .map_err(|e| {
139                FaucetError::internal(format!(
140                    "Failed to execute PaySui transaction for coin {:?}, with err {:?}",
141                    coin_id, e
142                ))
143            })
144            .map_err(Into::into)
145    }
146
147    async fn execute_txn_with_retries(
148        &self,
149        tx: TransactionData,
150        coin_id: ObjectID,
151        num_retries: u8,
152    ) -> Result<ExecutedTransaction, anyhow::Error> {
153        let mut retry_delay = Duration::from_millis(500);
154        let mut i = 0;
155
156        loop {
157            if i == num_retries {
158                bail!("Failed to execute transaction after {num_retries} retries",);
159            }
160            let res = self.execute_txn(&tx, coin_id).await;
161
162            if res.is_ok() {
163                return res;
164            }
165            i += 1;
166            tokio::time::sleep(retry_delay).await;
167            retry_delay *= 2;
168        }
169    }
170
171    pub fn get_coin_amount(&self) -> u64 {
172        self.coin_amount
173    }
174}
175
176/// Finds gas coins with sufficient balance and returns the address to use as the active address
177/// for the faucet. If the initial active address in the wallet does not have enough gas coins,
178/// it will iterate through the addresses to find one with sufficient gas coins.
179async fn find_gas_coins_and_address(
180    wallet: &mut WalletContext,
181    config: &FaucetConfig,
182) -> Result<(Vec<GasCoin>, SuiAddress), FaucetError> {
183    let active_address = wallet
184        .active_address()
185        .map_err(|e| FaucetError::Wallet(e.to_string()))?;
186
187    for address in std::iter::once(active_address).chain(wallet.get_addresses().into_iter()) {
188        let coins: Vec<_> = wallet
189            .gas_objects(address)
190            .await
191            .map_err(|e| FaucetError::Wallet(e.to_string()))?
192            .iter()
193            .filter_map(|(balance, obj)| {
194                if *balance >= config.amount {
195                    GasCoin::try_from(obj).ok()
196                } else {
197                    None
198                }
199            })
200            .collect();
201
202        if !coins.is_empty() {
203            return Ok((coins, address));
204        }
205    }
206
207    Err(FaucetError::Wallet(
208        "No address found with sufficient coins".to_string(),
209    ))
210}
211
212#[cfg(test)]
213mod tests {
214
215    use super::*;
216    use test_cluster::TestClusterBuilder;
217
218    #[tokio::test]
219    async fn test_local_faucet_execute_txn() {
220        // Setup test cluster
221        let cluster = TestClusterBuilder::new().build().await;
222        let client = cluster.grpc_client();
223
224        let config = FaucetConfig::default();
225        let local_faucet = LocalFaucet::new(cluster.wallet, config).await.unwrap();
226
227        // Test execute_txn
228        let recipient = SuiAddress::random_for_testing_only();
229        let tx = local_faucet.local_request_execute_tx(recipient).await;
230
231        assert!(tx.is_ok());
232
233        let coins = client
234            .get_owned_objects(recipient, None, None, None)
235            .await
236            .unwrap();
237
238        assert_eq!(coins.items.len(), local_faucet.num_coins);
239
240        let tx = local_faucet.local_request_execute_tx(recipient).await;
241        assert!(tx.is_ok());
242        let coins = client
243            .get_owned_objects(recipient, None, None, None)
244            .await
245            .unwrap();
246
247        assert_eq!(coins.items.len(), 2 * local_faucet.num_coins);
248    }
249
250    #[tokio::test]
251    async fn test_find_gas_coins_and_address() {
252        let mut cluster = TestClusterBuilder::new().build().await;
253        let wallet = cluster.wallet_mut();
254        let config = FaucetConfig::default();
255
256        // Test find_gas_coins_and_address
257        let result = find_gas_coins_and_address(wallet, &config).await;
258        assert!(result.is_ok());
259
260        let (coins, _) = result.unwrap();
261        assert!(!coins.is_empty());
262        assert!(coins.iter().map(|c| c.value()).sum::<u64>() >= config.amount);
263    }
264}
sui_faucet/local_faucet.rs

sui_faucet/
local_faucet.rs
