feat(faucet): rate limit based upon wallet locks

sn_faucet/Cargo.toml

@@ -42,6 +42,7 @@ sn_transfers = { path = "../sn_transfers", version = "0.16.5" }
 tokio = { version = "1.32.0", features = ["parking_lot", "rt"] }
 tracing = { version = "~0.1.26" }
 url = "2.5.0"
+fs2 = "0.4.3"
 
 [lints]
 workspace = true

sn_faucet/src/faucet_server.rs

@@ -10,12 +10,18 @@
 use crate::token_distribution;
 use crate::{claim_genesis, send_tokens};
 use color_eyre::eyre::Result;
+use fs2::FileExt;
 use sn_client::Client;
-use sn_transfers::{get_faucet_data_dir, HotWallet, NanoTokens};
+use sn_transfers::{
+    get_faucet_data_dir, wallet_lockfile_name, HotWallet, NanoTokens, WALLET_DIR_NAME,
+};
 use std::collections::HashMap;
 use std::path::Path;
-use tracing::{debug, error};
-use warp::{http::Response, Filter, Reply};
+use tracing::{debug, error, info, warn};
+use warp::{
+    http::{Response, StatusCode},
+    Filter, Reply,
+};
 
 /// Run the faucet server.
 ///
@@ -64,6 +70,17 @@ async fn respond_to_distribution_request(
     query: HashMap<String, String>,
     balances: HashMap<String, NanoTokens>,
 ) -> std::result::Result<impl Reply, std::convert::Infallible> {
+    // some rate limiting
+    if is_wallet_locked() {
+        warn!("Rate limited request due to locked wallet");
+
+        let mut response = Response::new("Rate limited".to_string());
+        *response.status_mut() = StatusCode::TOO_MANY_REQUESTS;
+
+        // Either opening the file or locking it failed, indicating rate limiting should occur
+        return Ok(response);
+    }
+
     let r =
         match token_distribution::handle_distribution_req(&client, query, balances.clone()).await {
             Ok(distribution) => Response::new(distribution.to_string()),
@@ -77,10 +94,42 @@ async fn respond_to_distribution_request(
     Ok(r)
 }
 
+fn is_wallet_locked() -> bool {
+    info!("Checking if wallet is locked");
+    let root_dir = get_faucet_data_dir();
+
+    let wallet_dir = root_dir.join(WALLET_DIR_NAME);
+    let wallet_lockfile_name = wallet_lockfile_name(&wallet_dir);
+    let file_result = std::fs::OpenOptions::new()
+        .create(true)
+        .write(true)
+        .truncate(true)
+        .open(wallet_lockfile_name)
+        .and_then(|file| file.try_lock_exclusive());
+    info!("After if wallet is locked");
+
+    if file_result.is_err() {
+        // Either opening the file or locking it failed, indicating rate limiting should occur
+        return true;
+    }
+
+    false
+}
+
 async fn respond_to_gift_request(
     client: Client,
     key: String,
 ) -> std::result::Result<impl Reply, std::convert::Infallible> {
+    // some rate limiting
+    if is_wallet_locked() {
+        warn!("Rate limited request due to locked wallet");
+        let mut response = Response::new("Rate limited".to_string());
+        *response.status_mut() = StatusCode::TOO_MANY_REQUESTS;
+
+        // Either opening the file or locking it failed, indicating rate limiting should occur
+        return Ok(response);
+    }
+
     const GIFT_AMOUNT_SNT: &str = "1";
     match send_tokens(&client, GIFT_AMOUNT_SNT, &key).await {
         Ok(transfer) => {

sn_transfers/src/lib.rs

@@ -33,8 +33,9 @@ pub use genesis::{
     TOTAL_SUPPLY,
 };
 pub use wallet::{
-    bls_secret_from_hex, Error as WalletError, HotWallet, Payment, PaymentQuote, QuotingMetrics,
-    Result as WalletResult, WatchOnlyWallet, QUOTE_EXPIRATION_SECS,
+    bls_secret_from_hex, wallet_lockfile_name, Error as WalletError, HotWallet, Payment,
+    PaymentQuote, QuotingMetrics, Result as WalletResult, WatchOnlyWallet, QUOTE_EXPIRATION_SECS,
+    WALLET_DIR_NAME,
 };
 
 // re-export crates used in our public API

sn_transfers/src/wallet.rs

@@ -59,17 +59,18 @@ mod keys;
 mod wallet_file;
 mod watch_only;
 
-use crate::{NanoTokens, UniquePubkey};
-use wallet_file::wallet_file_name;
-
 pub use self::{
     data_payments::{Payment, PaymentQuote, QuotingMetrics, QUOTE_EXPIRATION_SECS},
     error::{Error, Result},
     hot_wallet::HotWallet,
     keys::bls_secret_from_hex,
     watch_only::WatchOnlyWallet,
 };
+use crate::{NanoTokens, UniquePubkey};
+pub use hot_wallet::WALLET_DIR_NAME;
 pub(crate) use keys::store_new_keypair;
+use wallet_file::wallet_file_name;
+pub use wallet_file::wallet_lockfile_name;
 
 use serde::{Deserialize, Serialize};
 use std::{collections::BTreeMap, fs, path::Path};

sn_transfers/src/wallet/hot_wallet.rs

@@ -34,7 +34,7 @@ use std::{
     time::Instant,
 };
 
-const WALLET_DIR_NAME: &str = "wallet";
+pub const WALLET_DIR_NAME: &str = "wallet";
 
 /// A locked file handle, that when dropped releases the lock.
 pub type WalletExclusiveAccess = File;
@@ -76,6 +76,12 @@ impl HotWallet {
         self.watchonly_wallet.lock()
     }
 
+    /// Tries to locks the wallet and returns exclusive access to the wallet
+    /// Bails if unable to lock the wallet
+    pub fn try_lock(&self) -> Result<WalletExclusiveAccess> {
+        self.watchonly_wallet.try_lock()
+    }
+
     /// Stores the given cash_notes to the `created cash_notes dir` in the wallet dir.
     /// These can then be sent to the recipients out of band, over any channel preferred.
     pub fn store_cash_notes_to_disk<'a, T>(&self, cash_notes: T) -> Result<()>

sn_transfers/src/wallet/wallet_file.rs

@@ -39,7 +39,7 @@ pub(super) fn wallet_file_name(wallet_dir: &Path) -> PathBuf {
 }
 
 /// Returns the wallet lockfile filename
-pub(super) fn wallet_lockfile_name(wallet_dir: &Path) -> PathBuf {
+pub fn wallet_lockfile_name(wallet_dir: &Path) -> PathBuf {
     wallet_dir.join(WALLET_LOCK_FILE_NAME)
 }
 

sn_transfers/src/wallet/watch_only.rs

@@ -337,6 +337,22 @@ impl WatchOnlyWallet {
         file.lock_exclusive()?;
         Ok(file)
     }
+
+    // Tries to locks the wallet and returns exclusive access to the wallet
+    // This lock prevents any other process from locking the wallet dir, effectively acts as a mutex for the wallet
+    // This does not wait to get the lock, choosing to bail instead
+    pub(super) fn try_lock(&self) -> Result<WalletExclusiveAccess> {
+        let lock = wallet_lockfile_name(&self.wallet_dir);
+        let file = OpenOptions::new()
+            .create(true)
+            .write(true)
+            .truncate(true)
+            .open(lock)?;
+
+        #[cfg(not(target_arch = "wasm32"))]
+        file.try_lock_exclusive()?;
+        Ok(file)
+    }
 }
 
 #[cfg(test)]