[CLI] Add support for setting faucet auth token

crates/aptos-rest-client/src/faucet.rs

@@ -17,30 +17,34 @@ pub struct FaucetClient {
 
 impl FaucetClient {
     pub fn new(faucet_url: Url, rest_url: Url) -> Self {
+        Self::new_from_rest_client(faucet_url, Client::new(rest_url))
+    }
+
+    pub fn new_for_testing(faucet_url: Url, rest_url: Url) -> Self {
         Self {
             faucet_url,
             inner: ReqwestClient::builder()
                 .timeout(Duration::from_secs(10))
                 .build()
                 .unwrap(),
-            rest_client: Client::new(rest_url),
+            rest_client: Client::new(rest_url)
+                // By default the path is prefixed with the version, e.g. `v1`.
+                // The fake API used in the faucet tests doesn't have a
+                // versioned API however, so we just set it to `/`.
+                .version_path_base("/".to_string())
+                .unwrap(),
             token: None,
         }
     }
 
-    pub fn new_for_testing(faucet_url: Url, rest_url: Url) -> Self {
+    pub fn new_from_rest_client(faucet_url: Url, rest_client: Client) -> Self {
         Self {
             faucet_url,
             inner: ReqwestClient::builder()
                 .timeout(Duration::from_secs(10))
                 .build()
                 .unwrap(),
-            rest_client: Client::new(rest_url)
-                // By default the path is prefixed with the version, e.g. `v1`.
-                // The fake API used in the faucet tests doesn't have a
-                // versioned API however, so we just set it to `/`.
-                .version_path_base("/".to_string())
-                .unwrap(),
+            rest_client,
             token: None,
         }
     }

crates/aptos/Cargo.toml

@@ -47,7 +47,7 @@ async-trait = { workspace = true }
 base64 = { workspace = true }
 bcs = { workspace = true }
 chrono = { workspace = true }
-clap = { workspace = true, features = ["unstable-styles"] }
+clap = { workspace = true, features = ["env", "unstable-styles"] }
 clap_complete = { workspace = true }
 codespan-reporting = { workspace = true }
 dirs = { workspace = true }

crates/aptos/src/account/fund.rs

@@ -3,10 +3,7 @@
 
 use crate::{
     account::create::DEFAULT_FUNDED_COINS,
-    common::{
-        types::{CliCommand, CliTypedResult, FaucetOptions, ProfileOptions, RestOptions},
-        utils::{fund_account, wait_for_transactions},
-    },
+    common::types::{CliCommand, CliTypedResult, FaucetOptions, ProfileOptions, RestOptions},
 };
 use aptos_types::account_address::AccountAddress;
 use async_trait::async_trait;
@@ -51,14 +48,10 @@ impl CliCommand<String> for FundWithFaucet {
         } else {
             self.profile_options.account_address()?
         };
-        let hashes = fund_account(
-            self.faucet_options.faucet_url(&self.profile_options)?,
-            self.amount,
-            address,
-        )
-        .await?;
         let client = self.rest_options.client(&self.profile_options)?;
-        wait_for_transactions(&client, hashes).await?;
+        self.faucet_options
+            .fund_account(client, &self.profile_options, self.amount, address)
+            .await?;
         return Ok(format!(
             "Added {} Octas to account {}",
             self.amount, address

crates/aptos/src/common/init.rs

@@ -9,7 +9,7 @@ use crate::{
             ConfigSearchMode, EncodingOptions, PrivateKeyInputOptions, ProfileConfig,
             ProfileOptions, PromptOptions, RngArgs, DEFAULT_PROFILE,
         },
-        utils::{fund_account, prompt_yes_with_override, read_line, wait_for_transactions},
+        utils::{fund_account, prompt_yes_with_override, read_line},
     },
 };
 use aptos_crypto::{ed25519::Ed25519PrivateKey, PrivateKey, ValidCryptoMaterialStringExt};
@@ -45,6 +45,11 @@ pub struct InitTool {
     #[clap(long)]
     pub faucet_url: Option<Url>,
 
+    /// Auth token, if we're using the faucet. This is only used this time, we don't
+    /// store it.
+    #[clap(long, env)]
+    pub faucet_auth_token: Option<String>,
+
     /// Whether to skip the faucet for a non-faucet endpoint
     #[clap(long)]
     pub skip_faucet: bool,
@@ -159,15 +164,14 @@ impl CliCommand<()> for InitTool {
         };
         let public_key = private_key.public_key();
 
-        let client = aptos_rest_client::Client::new(
-            Url::parse(
-                profile_config
-                    .rest_url
-                    .as_ref()
-                    .expect("Must have rest client as created above"),
-            )
-            .map_err(|err| CliError::UnableToParse("rest_url", err.to_string()))?,
-        );
+        let rest_url = Url::parse(
+            profile_config
+                .rest_url
+                .as_ref()
+                .expect("Must have rest client as created above"),
+        )
+        .map_err(|err| CliError::UnableToParse("rest_url", err.to_string()))?;
+        let client = aptos_rest_client::Client::new(rest_url);
 
         // lookup the address from onchain instead of deriving it
         // if this is the rotated key, deriving it will outputs an incorrect address
@@ -225,14 +229,15 @@ impl CliCommand<()> for InitTool {
                     "Account {} doesn't exist, creating it and funding it with {} Octas",
                     address, NUM_DEFAULT_OCTAS
                 );
-                let hashes = fund_account(
+                fund_account(
+                    client,
                     Url::parse(faucet_url)
                         .map_err(|err| CliError::UnableToParse("rest_url", err.to_string()))?,
-                    NUM_DEFAULT_OCTAS,
+                    self.faucet_auth_token.as_deref(),
                     address,
+                    NUM_DEFAULT_OCTAS,
                 )
                 .await?;
-                wait_for_transactions(&client, hashes).await?;
                 eprintln!("Account {} funded successfully", address);
             }
         } else if account_exists {

crates/aptos/src/common/types.rs

@@ -1,6 +1,7 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
+use super::utils::fund_account;
 use crate::{
     common::{
         init::Network,
@@ -1268,14 +1269,22 @@ pub struct FaucetOptions {
     /// URL for the faucet endpoint e.g. `https://faucet.devnet.aptoslabs.com`
     #[clap(long)]
     faucet_url: Option<reqwest::Url>,
+
+    /// Auth token to bypass faucet ratelimits. You can also set this as an environment
+    /// variable with FAUCET_AUTH_TOKEN.
+    #[clap(long, env)]
+    faucet_auth_token: Option<String>,
 }
 
 impl FaucetOptions {
-    pub fn new(faucet_url: Option<reqwest::Url>) -> Self {
-        FaucetOptions { faucet_url }
+    pub fn new(faucet_url: Option<reqwest::Url>, faucet_auth_token: Option<String>) -> Self {
+        FaucetOptions {
+            faucet_url,
+            faucet_auth_token,
+        }
     }
 
-    pub fn faucet_url(&self, profile: &ProfileOptions) -> CliTypedResult<reqwest::Url> {
+    fn faucet_url(&self, profile: &ProfileOptions) -> CliTypedResult<reqwest::Url> {
         if let Some(ref faucet_url) = self.faucet_url {
             Ok(faucet_url.clone())
         } else if let Some(Some(url)) = CliConfig::load_profile(
@@ -1287,9 +1296,27 @@ impl FaucetOptions {
             reqwest::Url::parse(&url)
                 .map_err(|err| CliError::UnableToParse("config faucet_url", err.to_string()))
         } else {
-            Err(CliError::CommandArgumentError("No faucet given.  Please add --faucet-url or add a faucet URL to the .aptos/config.yaml for the current profile".to_string()))
+            Err(CliError::CommandArgumentError("No faucet given. Please add --faucet-url or add a faucet URL to the .aptos/config.yaml for the current profile".to_string()))
         }
     }
+
+    /// Fund an account with the faucet.
+    pub async fn fund_account(
+        &self,
+        rest_client: Client,
+        profile: &ProfileOptions,
+        num_octas: u64,
+        address: AccountAddress,
+    ) -> CliTypedResult<()> {
+        fund_account(
+            rest_client,
+            self.faucet_url(profile)?,
+            self.faucet_auth_token.as_deref(),
+            address,
+            num_octas,
+        )
+        .await
+    }
 }
 
 /// Gas price options for manipulating how to prioritize transactions

crates/aptos/src/common/utils.rs

@@ -13,7 +13,7 @@ use aptos_build_info::build_information;
 use aptos_crypto::ed25519::{Ed25519PrivateKey, Ed25519PublicKey};
 use aptos_keygen::KeyGen;
 use aptos_logger::{debug, Level};
-use aptos_rest_client::{aptos_api_types::HashValue, Account, Client, State};
+use aptos_rest_client::{aptos_api_types::HashValue, Account, Client, FaucetClient, State};
 use aptos_telemetry::service::telemetry_is_disabled;
 use aptos_types::{
     account_address::create_multisig_account_address,
@@ -416,35 +416,23 @@ pub fn read_line(input_name: &'static str) -> CliTypedResult<String> {
     Ok(input_buf)
 }
 
-/// Fund account (and possibly create it) from a faucet
+/// Fund account (and possibly create it) from a faucet. This function waits for the
+/// transaction on behalf of the caller.
 pub async fn fund_account(
+    rest_client: Client,
     faucet_url: Url,
-    num_octas: u64,
+    faucet_auth_token: Option<&str>,
     address: AccountAddress,
-) -> CliTypedResult<Vec<HashValue>> {
-    let response = reqwest::Client::new()
-        .post(format!(
-            "{}mint?amount={}&auth_key={}",
-            faucet_url, num_octas, address
-        ))
-        .body("{}")
-        .send()
-        .await
-        .map_err(|err| {
-            CliError::ApiError(format!("Failed to fund account with faucet: {:#}", err))
-        })?;
-    if response.status() == 200 {
-        let hashes: Vec<HashValue> = response
-            .json()
-            .await
-            .map_err(|err| CliError::UnexpectedError(err.to_string()))?;
-        Ok(hashes)
-    } else {
-        Err(CliError::ApiError(format!(
-            "Faucet issue: {}",
-            response.status()
-        )))
+    num_octas: u64,
+) -> CliTypedResult<()> {
+    let mut client = FaucetClient::new_from_rest_client(faucet_url, rest_client);
+    if let Some(token) = faucet_auth_token {
+        client = client.with_auth_token(token.to_string());
     }
+    client
+        .fund(address, num_octas)
+        .await
+        .map_err(|err| CliError::ApiError(format!("Faucet issue: {:#}", err)))
 }
 
 /// Wait for transactions, returning an error if any of them fail.

crates/aptos/src/test/mod.rs

@@ -541,6 +541,7 @@ impl CliTestFramework {
             network: Some(Network::Custom),
             rest_url: Some(self.endpoint.clone()),
             faucet_url: Some(self.faucet_endpoint.clone()),
+            faucet_auth_token: None,
             rng_args: RngArgs::from_seed([0; 32]),
             private_key_options: PrivateKeyInputOptions::from_private_key(private_key)?,
             profile_options: Default::default(),
@@ -1079,7 +1080,7 @@ impl CliTestFramework {
     }
 
     pub fn faucet_options(&self) -> FaucetOptions {
-        FaucetOptions::new(Some(self.faucet_endpoint.clone()))
+        FaucetOptions::new(Some(self.faucet_endpoint.clone()), None)
     }
 
     fn transaction_options(