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

sui_faucet/
local_faucet.rs
