iota_cluster_test/
faucet.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::HashMap, env, sync::Arc};
6
7use async_trait::async_trait;
8use fastcrypto::encoding::{Encoding, Hex};
9use iota_faucet::{
10    BatchFaucetResponse, BatchStatusFaucetResponse, Faucet, FaucetConfig, FaucetResponse,
11    SimpleFaucet,
12};
13use iota_types::{base_types::IotaAddress, crypto::KeypairTraits};
14use tracing::{Instrument, debug, info, info_span};
15use uuid::Uuid;
16
17use super::cluster::{Cluster, new_wallet_context_from_cluster};
18
19pub struct FaucetClientFactory;
20
21impl FaucetClientFactory {
22    pub async fn new_from_cluster(
23        cluster: &(dyn Cluster + Sync + Send),
24    ) -> Arc<dyn FaucetClient + Sync + Send> {
25        match cluster.remote_faucet_url() {
26            Some(url) => Arc::new(RemoteFaucetClient::new(url.into())),
27            // If faucet_url is none, it's a local cluster
28            None => {
29                let key = cluster
30                    .local_faucet_key()
31                    .expect("Expect local faucet key for local cluster")
32                    .copy();
33                let wallet_context = new_wallet_context_from_cluster(cluster, key)
34                    .instrument(info_span!("init_wallet_context_for_faucet"));
35
36                let prom_registry = prometheus::Registry::new();
37                let config = FaucetConfig::default();
38                let simple_faucet = SimpleFaucet::new(
39                    wallet_context.into_inner(),
40                    &prom_registry,
41                    &cluster.config_directory().join("faucet.wal"),
42                    config,
43                )
44                .await
45                .unwrap();
46
47                Arc::new(LocalFaucetClient::new(simple_faucet))
48            }
49        }
50    }
51}
52
53/// Faucet Client abstraction
54#[async_trait]
55pub trait FaucetClient {
56    async fn request_iota_coins(&self, request_address: IotaAddress) -> FaucetResponse;
57    async fn batch_request_iota_coins(&self, request_address: IotaAddress) -> BatchFaucetResponse;
58    async fn get_batch_send_status(&self, task_id: Uuid) -> BatchStatusFaucetResponse;
59}
60
61/// Client for a remote faucet that is accessible by POST requests
62pub struct RemoteFaucetClient {
63    remote_url: String,
64}
65
66impl RemoteFaucetClient {
67    fn new(url: String) -> Self {
68        info!("Use remote faucet: {}", url);
69        Self { remote_url: url }
70    }
71}
72
73#[async_trait]
74impl FaucetClient for RemoteFaucetClient {
75    /// Request test IOTA coins from faucet.
76    /// It also verifies the effects are observed by fullnode.
77    async fn request_iota_coins(&self, request_address: IotaAddress) -> FaucetResponse {
78        let gas_url = format!("{}/gas", self.remote_url);
79        debug!("Getting coin from remote faucet {}", gas_url);
80        let data = HashMap::from([("recipient", Hex::encode(request_address))]);
81        let map = HashMap::from([("FixedAmountRequest", data)]);
82
83        let auth_header = match env::var("FAUCET_AUTH_HEADER") {
84            Ok(val) => val,
85            _ => "".to_string(),
86        };
87
88        let response = reqwest::Client::new()
89            .post(&gas_url)
90            .header("Authorization", auth_header)
91            .json(&map)
92            .send()
93            .await
94            .unwrap_or_else(|e| panic!("Failed to talk to remote faucet {:?}: {:?}", gas_url, e));
95        let full_bytes = response.bytes().await.unwrap();
96        let faucet_response: FaucetResponse = serde_json::from_slice(&full_bytes)
97            .map_err(|e| anyhow::anyhow!("json deser failed with bytes {:?}: {e}", full_bytes))
98            .unwrap();
99
100        if let Some(error) = faucet_response.error {
101            panic!("Failed to get gas tokens with error: {}", error)
102        };
103
104        faucet_response
105    }
106    async fn batch_request_iota_coins(&self, request_address: IotaAddress) -> BatchFaucetResponse {
107        let gas_url = format!("{}/v1/gas", self.remote_url);
108        debug!("Getting coin from remote faucet {}", gas_url);
109        let data = HashMap::from([("recipient", Hex::encode(request_address))]);
110        let map = HashMap::from([("FixedAmountRequest", data)]);
111
112        let auth_header = match env::var("FAUCET_AUTH_HEADER") {
113            Ok(val) => val,
114            _ => "".to_string(),
115        };
116
117        let response = reqwest::Client::new()
118            .post(&gas_url)
119            .header("Authorization", auth_header)
120            .json(&map)
121            .send()
122            .await
123            .unwrap_or_else(|e| panic!("Failed to talk to remote faucet {:?}: {:?}", gas_url, e));
124        let full_bytes = response.bytes().await.unwrap();
125        let faucet_response: BatchFaucetResponse = serde_json::from_slice(&full_bytes)
126            .map_err(|e| anyhow::anyhow!("json deser failed with bytes {:?}: {e}", full_bytes))
127            .unwrap();
128
129        if let Some(error) = faucet_response.error {
130            panic!("Failed to get gas tokens with error: {}", error)
131        };
132
133        faucet_response
134    }
135    async fn get_batch_send_status(&self, task_id: Uuid) -> BatchStatusFaucetResponse {
136        let status_url = format!("{}/v1/status/{}", self.remote_url, task_id);
137        debug!(
138            "Checking status for task {} from remote faucet {}",
139            task_id.to_string(),
140            status_url
141        );
142
143        let auth_header = match env::var("FAUCET_AUTH_HEADER") {
144            Ok(val) => val,
145            _ => "".to_string(),
146        };
147
148        let response = reqwest::Client::new()
149            .get(&status_url)
150            .header("Authorization", auth_header)
151            .send()
152            .await
153            .unwrap_or_else(|e| {
154                panic!("Failed to talk to remote faucet {:?}: {:?}", status_url, e)
155            });
156        let full_bytes = response.bytes().await.unwrap();
157        let faucet_response: BatchStatusFaucetResponse = serde_json::from_slice(&full_bytes)
158            .map_err(|e| anyhow::anyhow!("json deser failed with bytes {:?}: {e}", full_bytes))
159            .unwrap();
160
161        faucet_response
162    }
163}
164
165/// A local faucet that holds some coins since genesis
166pub struct LocalFaucetClient {
167    simple_faucet: Arc<SimpleFaucet>,
168}
169
170impl LocalFaucetClient {
171    fn new(simple_faucet: Arc<SimpleFaucet>) -> Self {
172        info!("Use local faucet");
173        Self { simple_faucet }
174    }
175}
176#[async_trait]
177impl FaucetClient for LocalFaucetClient {
178    async fn request_iota_coins(&self, request_address: IotaAddress) -> FaucetResponse {
179        let receipt = self
180            .simple_faucet
181            .send(Uuid::new_v4(), request_address, &[200_000_000_000; 5])
182            .await
183            .unwrap_or_else(|err| panic!("Failed to get gas tokens with error: {}", err));
184
185        receipt.into()
186    }
187    async fn batch_request_iota_coins(&self, request_address: IotaAddress) -> BatchFaucetResponse {
188        let receipt = self
189            .simple_faucet
190            .batch_send(Uuid::new_v4(), request_address, &[200_000_000_000; 5])
191            .await
192            .unwrap_or_else(|err| panic!("Failed to get gas tokens with error: {}", err));
193
194        receipt.into()
195    }
196    async fn get_batch_send_status(&self, task_id: Uuid) -> BatchStatusFaucetResponse {
197        let status = self
198            .simple_faucet
199            .get_batch_send_status(task_id)
200            .await
201            .unwrap_or_else(|err| panic!("Failed to get gas tokens with error: {}", err));
202
203        status.into()
204    }
205}