Limit assets per address (#47)

* limit assets per address

* evict old values

* improve testing

* make dispensed map thread local

* update rust version

* make tokio test-util a dev dep

* evict less than day

* more review comments

* add dispense tracker

* use mark_in_progress

* remove from in progress on err

* Update tests/dispense.rs

Co-authored-by: Brandon Vrooman <brandon.vrooman@gmail.com>

* only track if tx succeeds

* remove on empty wallet

* make sure to remove in prog on the right err

---------

Co-authored-by: Brandon Vrooman <brandon.vrooman@gmail.com>

Cargo.toml

@@ -38,6 +38,7 @@ fuel-types = { version = "0.43.0", features = ["random"] }
 futures = "0.3"
 insta = "1.14"
 rand = "0.8"
+tokio = { version = "1.0", features = ["test-util"] }
 
 [build-dependencies]
 minify-html = "0.8"

src/config.rs

@@ -1,7 +1,8 @@
 use crate::constants::{
-    CAPTCHA_KEY, CAPTCHA_SECRET, DEFAULT_FAUCET_DISPENSE_AMOUNT, DEFAULT_NODE_URL, DEFAULT_PORT,
-    DISPENSE_AMOUNT, FAUCET_ASSET_ID, FUEL_NODE_URL, HUMAN_LOGGING, LOG_FILTER, MIN_GAS_PRICE,
-    PUBLIC_FUEL_NODE_URL, SERVICE_PORT, TIMEOUT_SECONDS, WALLET_SECRET_KEY,
+    CAPTCHA_KEY, CAPTCHA_SECRET, DEFAULT_DISPENSE_INTERVAL, DEFAULT_FAUCET_DISPENSE_AMOUNT,
+    DEFAULT_NODE_URL, DEFAULT_PORT, DISPENSE_AMOUNT, DISPENSE_INTERVAL, FAUCET_ASSET_ID,
+    FUEL_NODE_URL, HUMAN_LOGGING, LOG_FILTER, MIN_GAS_PRICE, PUBLIC_FUEL_NODE_URL, SERVICE_PORT,
+    TIMEOUT_SECONDS, WALLET_SECRET_KEY,
 };
 use fuels_core::types::AssetId;
 use secrecy::Secret;
@@ -19,6 +20,7 @@ pub struct Config {
     pub wallet_secret_key: Option<Secret<String>>,
     pub dispense_amount: u64,
     pub dispense_asset_id: AssetId,
+    pub dispense_limit_interval: u64,
     pub min_gas_price: u64,
     pub timeout: u64,
 }
@@ -44,6 +46,10 @@ impl Default for Config {
                 .parse::<u64>()
                 .expect("expected a valid integer for DISPENSE_AMOUNT"),
             dispense_asset_id: FAUCET_ASSET_ID,
+            dispense_limit_interval: env::var(DISPENSE_INTERVAL)
+                .unwrap_or_else(|_| DEFAULT_DISPENSE_INTERVAL.to_string())
+                .parse::<u64>()
+                .expect("expected a valid integer for DISPENSE_LIMIT_INTERVAL"),
             min_gas_price: env::var(MIN_GAS_PRICE)
                 .unwrap_or_else(|_| "0".to_string())
                 .parse::<u64>()

src/constants.rs

@@ -11,6 +11,8 @@ pub const WALLET_SECRET_DEV_KEY: &str =
 pub const FUEL_NODE_URL: &str = "FUEL_NODE_URL";
 pub const DEFAULT_NODE_URL: &str = "http://127.0.0.1:4000";
 pub const DISPENSE_AMOUNT: &str = "DISPENSE_AMOUNT";
+pub const DISPENSE_INTERVAL: &str = "DISPENSE_LIMIT_INTERVAL";
+pub const DEFAULT_DISPENSE_INTERVAL: u64 = 24 * 60 * 60;
 pub const DEFAULT_FAUCET_DISPENSE_AMOUNT: u64 = 10_000_000;
 pub const FAUCET_ASSET_ID: AssetId = AssetId::new([0; 32]);
 pub const SERVICE_PORT: &str = "PORT";

src/dispense_tracker.rs

@@ -0,0 +1,100 @@
+use std::{
+    cmp::Ordering,
+    collections::{BinaryHeap, HashMap, HashSet},
+};
+
+use fuel_types::Address;
+
+#[derive(Debug, Eq, PartialEq)]
+pub struct Entry {
+    address: Address,
+    timestamp: u64,
+}
+
+impl Ord for Entry {
+    fn cmp(&self, other: &Self) -> Ordering {
+        other.timestamp.cmp(&self.timestamp)
+    }
+}
+
+impl PartialOrd for Entry {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+pub trait Clock: std::fmt::Debug + Send + Sync {
+    fn now(&self) -> u64;
+}
+
+#[derive(Debug)]
+pub struct TokioTime {}
+
+impl Clock for TokioTime {
+    fn now(&self) -> u64 {
+        tokio::time::Instant::now().elapsed().as_secs()
+    }
+}
+
+#[derive(Debug)]
+pub struct DispenseTracker {
+    tracked: HashMap<Address, u64>,
+    queue: BinaryHeap<Entry>,
+    in_progress: HashSet<Address>,
+    clock: Box<dyn Clock>,
+}
+
+impl Default for DispenseTracker {
+    fn default() -> Self {
+        Self {
+            tracked: HashMap::default(),
+            queue: BinaryHeap::default(),
+            in_progress: HashSet::default(),
+            clock: Box::new(TokioTime {}),
+        }
+    }
+}
+
+impl DispenseTracker {
+    pub fn new(clock: impl Clock + 'static) -> Self {
+        Self {
+            tracked: HashMap::new(),
+            queue: BinaryHeap::new(),
+            in_progress: HashSet::new(),
+            clock: Box::new(clock),
+        }
+    }
+
+    pub fn track(&mut self, address: Address) {
+        self.in_progress.remove(&address);
+
+        let timestamp = self.clock.now();
+        self.tracked.insert(address, timestamp);
+        self.queue.push(Entry { address, timestamp });
+    }
+
+    pub fn mark_in_progress(&mut self, address: Address) {
+        self.in_progress.insert(address);
+    }
+
+    pub fn remove_in_progress(&mut self, address: &Address) {
+        self.in_progress.remove(address);
+    }
+
+    pub fn evict_expired_entries(&mut self, eviction_duration: u64) {
+        let now = self.clock.now();
+
+        while let Some(oldest_entry) = self.queue.peek() {
+            if now - oldest_entry.timestamp > eviction_duration {
+                let removed_entry = self.queue.pop().unwrap();
+                self.tracked.remove(&removed_entry.address);
+            } else {
+                break;
+            }
+        }
+    }
+
+    pub fn has_tracked(&self, address: &Address) -> bool {
+        self.tracked.get(address).is_some() || self.in_progress.contains(address)
+    }
+}

src/lib.rs

@@ -1,6 +1,7 @@
 use crate::{
     config::Config,
     constants::{MAX_CONCURRENT_REQUESTS, WALLET_SECRET_DEV_KEY},
+    dispense_tracker::DispenseTracker,
     routes::health,
 };
 use anyhow::anyhow;
@@ -21,7 +22,7 @@ use secrecy::{ExposeSecret, Secret};
 use serde_json::json;
 use std::{
     net::{SocketAddr, TcpListener},
-    sync::Arc,
+    sync::{Arc, Mutex},
     time::Duration,
 };
 use tokio::task::JoinHandle;
@@ -37,9 +38,11 @@ pub mod config;
 pub mod models;
 
 mod constants;
+mod dispense_tracker;
 mod recaptcha;
 mod routes;
 
+pub use dispense_tracker::{Clock, TokioTime};
 pub use routes::THE_BIGGEST_AMOUNT;
 
 #[derive(Debug)]
@@ -88,9 +91,11 @@ pub type SharedFaucetState = Arc<tokio::sync::Mutex<FaucetState>>;
 pub type SharedWallet = Arc<WalletUnlocked>;
 pub type SharedConfig = Arc<Config>;
 pub type SharedNetworkConfig = Arc<NetworkConfig>;
+pub type SharedDispenseTracker = Arc<Mutex<DispenseTracker>>;
 
 pub async fn start_server(
     service_config: Config,
+    clock: impl Clock + 'static,
 ) -> (SocketAddr, JoinHandle<Result<(), anyhow::Error>>) {
     info!("{:#?}", &service_config);
 
@@ -179,6 +184,7 @@ pub async fn start_server(
                 ))))
                 .layer(Extension(Arc::new(service_config.clone())))
                 .layer(Extension(Arc::new(network_config)))
+                .layer(Extension(Arc::new(Mutex::new(DispenseTracker::new(clock)))))
                 .layer(
                     CorsLayer::new()
                         .allow_origin(Any)

src/main.rs

@@ -1,11 +1,12 @@
-use fuel_faucet::{config::Config, start_server};
+use fuel_faucet::{config::Config, start_server, TokioTime};
 use tracing_subscriber::EnvFilter;
 
 #[tokio::main]
 async fn main() {
     let config = Config::default();
     init_logger(&config);
-    let (_, task) = start_server(config).await;
+    let clock = TokioTime {};
+    let (_, task) = start_server(config, clock).await;
     let _ = task.await.unwrap();
 }