sui_cluster_test/
faucet.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use super::cluster::{Cluster, new_wallet_context_from_cluster};
5use async_trait::async_trait;
6use fastcrypto::encoding::{Encoding, Hex};
7use std::collections::HashMap;
8use std::env;
9use std::sync::Arc;
10use sui_faucet::{FaucetConfig, FaucetResponse, LocalFaucet, RequestStatus};
11use sui_types::base_types::SuiAddress;
12use sui_types::crypto::KeypairTraits;
13use tracing::{Instrument, debug, info, info_span};
14
15pub struct FaucetClientFactory;
16
17impl FaucetClientFactory {
18    pub async fn new_from_cluster(
19        cluster: &(dyn Cluster + Sync + Send),
20    ) -> Arc<dyn FaucetClient + Sync + Send> {
21        match cluster.remote_faucet_url() {
22            Some(url) => Arc::new(RemoteFaucetClient::new(url.into())),
23            // If faucet_url is none, it's a local cluster
24            None => {
25                let key = cluster
26                    .local_faucet_key()
27                    .expect("Expect local faucet key for local cluster")
28                    .copy();
29                let wallet_context = new_wallet_context_from_cluster(cluster, key)
30                    .await
31                    .instrument(info_span!("init_wallet_context_for_faucet"));
32
33                let config = FaucetConfig::default();
34                let simple_faucet = LocalFaucet::new(wallet_context.into_inner(), config)
35                    .await
36                    .unwrap();
37
38                Arc::new(LocalFaucetClient::new(simple_faucet))
39            }
40        }
41    }
42}
43
44/// Faucet Client abstraction
45#[async_trait]
46pub trait FaucetClient {
47    async fn request_sui_coins(&self, request_address: SuiAddress) -> FaucetResponse;
48}
49
50/// Client for a remote faucet that is accessible by POST requests
51pub struct RemoteFaucetClient {
52    remote_url: String,
53}
54
55impl RemoteFaucetClient {
56    fn new(url: String) -> Self {
57        info!("Use remote faucet: {}", url);
58        Self { remote_url: url }
59    }
60}
61
62#[async_trait]
63impl FaucetClient for RemoteFaucetClient {
64    /// Request test SUI coins from faucet.
65    /// It also verifies the effects are observed by fullnode.
66    async fn request_sui_coins(&self, request_address: SuiAddress) -> FaucetResponse {
67        let gas_url = format!("{}/v2/gas", self.remote_url);
68        debug!("Getting coin from remote faucet {}", gas_url);
69        let data = HashMap::from([("recipient", Hex::encode(request_address))]);
70        let map = HashMap::from([("FixedAmountRequest", data)]);
71
72        let auth_header = match env::var("FAUCET_AUTH_HEADER") {
73            Ok(val) => val,
74            _ => "".to_string(),
75        };
76
77        // Remote faucets rate-limit per IP with a plain-text 429 that advises
78        // a wait (e.g. "Too Many Requests! Wait for 4s"); honor it with
79        // bounded retries instead of failing the run.
80        const MAX_ATTEMPTS: u32 = 5;
81        for attempt in 1..=MAX_ATTEMPTS {
82            let response = reqwest::Client::new()
83                .post(&gas_url)
84                .header("Authorization", auth_header.clone())
85                .json(&map)
86                .send()
87                .await
88                .unwrap_or_else(|e| {
89                    panic!("Failed to talk to remote faucet {:?}: {:?}", gas_url, e)
90                });
91            let status = response.status();
92            let retry_after_secs = response
93                .headers()
94                .get(reqwest::header::RETRY_AFTER)
95                .and_then(|v| v.to_str().ok())
96                .and_then(|s| s.parse::<u64>().ok());
97            let full_bytes = response.bytes().await.unwrap();
98
99            if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
100                let body = String::from_utf8_lossy(&full_bytes);
101                if attempt == MAX_ATTEMPTS {
102                    break;
103                }
104                // Prefer the standard Retry-After header; fall back to the
105                // advised wait in the plain-text body.
106                let wait_secs = retry_after_secs
107                    .or_else(|| {
108                        body.split_whitespace().find_map(|tok| {
109                            tok.strip_suffix('s').and_then(|n| n.parse::<u64>().ok())
110                        })
111                    })
112                    .unwrap_or(5)
113                    .clamp(1, 60);
114                info!(
115                    "Faucet rate-limited (attempt {attempt}/{MAX_ATTEMPTS}): {body}; retrying in {wait_secs}s"
116                );
117                tokio::time::sleep(std::time::Duration::from_secs(wait_secs)).await;
118                continue;
119            }
120
121            let faucet_response: FaucetResponse = serde_json::from_slice(&full_bytes)
122                .map_err(|e| {
123                    anyhow::anyhow!(
124                        "json deser failed with status {status} bytes {full_bytes:?}: {e}"
125                    )
126                })
127                .unwrap();
128
129            if let RequestStatus::Failure(error) = &faucet_response.status {
130                panic!("Failed to get gas tokens with error: {}", error)
131            };
132
133            return faucet_response;
134        }
135        panic!("Faucet {gas_url} still rate-limiting after {MAX_ATTEMPTS} attempts")
136    }
137}
138
139/// A local faucet that holds some coins since genesis
140pub struct LocalFaucetClient {
141    simple_faucet: Arc<LocalFaucet>,
142}
143
144impl LocalFaucetClient {
145    fn new(simple_faucet: Arc<LocalFaucet>) -> Self {
146        info!("Use local faucet");
147        Self { simple_faucet }
148    }
149}
150#[async_trait]
151impl FaucetClient for LocalFaucetClient {
152    async fn request_sui_coins(&self, request_address: SuiAddress) -> FaucetResponse {
153        let coins = self
154            .simple_faucet
155            .local_request_execute_tx(request_address)
156            .await
157            .unwrap_or_else(|err| panic!("Failed to get gas tokens with error: {}", err));
158
159        FaucetResponse {
160            status: RequestStatus::Success,
161            coins_sent: Some(coins),
162        }
163    }
164}
sui_cluster_test/faucet.rs

sui_cluster_test/
faucet.rs
