iota_cluster_test/
faucet.rs1use 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 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#[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
61pub 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 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
165pub 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}