From 36574c30efe244898b7c3877966d04b781db6cce Mon Sep 17 00:00:00 2001
From: Tyera Eulberg <teulberg@gmail.com>
Date: Tue, 27 Apr 2021 19:31:32 -0600
Subject: [PATCH] Add allowed-ip list to faucet (#16891)

---
 faucet/src/bin/faucet.rs | 24 ++++++++++++++++++---
 faucet/src/faucet.rs     | 45 ++++++++++++++++++++++++++++++++++++----
 2 files changed, 62 insertions(+), 7 deletions(-)

diff --git a/faucet/src/bin/faucet.rs b/faucet/src/bin/faucet.rs
index 322a2d33f..cd79d4651 100644
--- a/faucet/src/bin/faucet.rs
+++ b/faucet/src/bin/faucet.rs
@@ -1,5 +1,5 @@
 use {
-    clap::{crate_description, crate_name, App, Arg},
+    clap::{crate_description, crate_name, values_t, App, Arg},
     log::*,
     solana_clap_utils::input_parsers::{lamports_of_sol, value_of},
     solana_faucet::{
@@ -8,7 +8,8 @@ use {
     },
     solana_sdk::signature::read_keypair_file,
     std::{
-        net::{Ipv4Addr, SocketAddr},
+        collections::HashSet,
+        net::{IpAddr, Ipv4Addr, SocketAddr},
         sync::{Arc, Mutex},
         thread,
     },
@@ -55,6 +56,17 @@ async fn main() {
                 .takes_value(true)
                 .help("Request limit for a single request, in SOL"),
         )
+        .arg(
+            Arg::with_name("allowed_ip")
+                .long("allow-ip")
+                .value_name("IP_ADDRESS")
+                .takes_value(true)
+                .multiple(true)
+                .help(
+                    "Allow requests from a particular IP address without request limit; \
+                    recipient address will be used to check request limits instead",
+                ),
+        )
         .get_matches();
 
     let faucet_keypair = read_keypair_file(matches.value_of("keypair").unwrap())
@@ -64,13 +76,19 @@ async fn main() {
     let per_time_cap = lamports_of_sol(&matches, "per_time_cap");
     let per_request_cap = lamports_of_sol(&matches, "per_request_cap");
 
+    let allowed_ips: HashSet<_> = values_t!(matches.values_of("allowed_ip"), IpAddr)
+        .unwrap_or_default()
+        .into_iter()
+        .collect();
+
     let faucet_addr = socketaddr!(0, FAUCET_PORT);
 
-    let faucet = Arc::new(Mutex::new(Faucet::new(
+    let faucet = Arc::new(Mutex::new(Faucet::new_with_allowed_ips(
         faucet_keypair,
         time_slice,
         per_time_cap,
         per_request_cap,
+        allowed_ips,
     )));
 
     let faucet1 = faucet.clone();
diff --git a/faucet/src/faucet.rs b/faucet/src/faucet.rs
index ed0be44cf..c2831ee62 100644
--- a/faucet/src/faucet.rs
+++ b/faucet/src/faucet.rs
@@ -22,7 +22,7 @@ use {
         transaction::Transaction,
     },
     std::{
-        collections::HashMap,
+        collections::{HashMap, HashSet},
         io::{Read, Write},
         net::{IpAddr, Ipv4Addr, SocketAddr, TcpStream},
         sync::{mpsc::Sender, Arc, Mutex},
@@ -106,6 +106,7 @@ pub struct Faucet {
     pub time_slice: Duration,
     per_time_cap: Option<u64>,
     per_request_cap: Option<u64>,
+    allowed_ips: HashSet<IpAddr>,
 }
 
 impl Faucet {
@@ -114,7 +115,23 @@ impl Faucet {
         time_input: Option<u64>,
         per_time_cap: Option<u64>,
         per_request_cap: Option<u64>,
-    ) -> Faucet {
+    ) -> Self {
+        Self::new_with_allowed_ips(
+            faucet_keypair,
+            time_input,
+            per_time_cap,
+            per_request_cap,
+            HashSet::new(),
+        )
+    }
+
+    pub fn new_with_allowed_ips(
+        faucet_keypair: Keypair,
+        time_input: Option<u64>,
+        per_time_cap: Option<u64>,
+        per_request_cap: Option<u64>,
+        allowed_ips: HashSet<IpAddr>,
+    ) -> Self {
         let time_slice = Duration::new(time_input.unwrap_or(TIME_SLICE), 0);
         if let Some((per_request_cap, per_time_cap)) = per_request_cap.zip(per_time_cap) {
             if per_time_cap < per_request_cap {
@@ -126,13 +143,14 @@ impl Faucet {
                 );
             }
         }
-        Faucet {
+        Self {
             faucet_keypair,
             ip_cache: HashMap::new(),
             address_cache: HashMap::new(),
             time_slice,
             per_time_cap,
             per_request_cap,
+            allowed_ips,
         }
     }
 
@@ -205,7 +223,7 @@ impl Faucet {
                         )));
                     }
                 }
-                if !ip.is_loopback() {
+                if !ip.is_loopback() && !self.allowed_ips.contains(&ip) {
                     self.check_time_request_limit(lamports, ip)?;
                 }
                 self.check_time_request_limit(lamports, to)?;
@@ -587,6 +605,25 @@ mod tests {
         let tx1 = faucet.build_airdrop_transaction(request1, ip);
         assert!(tx1.is_err());
 
+        // Test multiple requests from allowed ip with different addresses succeed
+        let mint = Keypair::new();
+        let ip = socketaddr!([203, 0, 113, 1], 0).ip();
+        let mut allowed_ips = HashSet::new();
+        allowed_ips.insert(ip);
+        faucet = Faucet::new_with_allowed_ips(mint, None, Some(2), None, allowed_ips);
+        let other = Pubkey::new_unique();
+        let _tx0 = faucet.build_airdrop_transaction(request, ip).unwrap(); // first request succeeds
+        let request1 = FaucetRequest::GetAirdrop {
+            lamports: 2,
+            to: other,
+            blockhash,
+        };
+        let _tx1 = faucet.build_airdrop_transaction(request1, ip).unwrap(); // first request succeeds
+        let tx0 = faucet.build_airdrop_transaction(request, ip);
+        assert!(tx0.is_err());
+        let tx1 = faucet.build_airdrop_transaction(request1, ip);
+        assert!(tx1.is_err());
+
         // Test per-request cap
         let mint = Keypair::new();
         let mint_pubkey = mint.pubkey();