From 71d5409b3bbc6b8e2d104a0fffca526181b8a8b9 Mon Sep 17 00:00:00 2001 From: Tyera Eulberg Date: Mon, 17 Aug 2020 17:33:14 -0600 Subject: [PATCH] Faucet: Add per-request cap (#11665) * Add per-request cap; also use clap-utils * Clean up arg names and take cap inputs as SOL --- faucet/src/bin/faucet.rs | 33 +++++++++--------- faucet/src/faucet.rs | 72 ++++++++++++++++++++++++---------------- 2 files changed, 60 insertions(+), 45 deletions(-) diff --git a/faucet/src/bin/faucet.rs b/faucet/src/bin/faucet.rs index 6bb1150aae..7144704603 100644 --- a/faucet/src/bin/faucet.rs +++ b/faucet/src/bin/faucet.rs @@ -1,4 +1,5 @@ use clap::{crate_description, crate_name, App, Arg}; +use solana_clap_utils::input_parsers::{lamports_of_sol, value_of}; use solana_faucet::{ faucet::{run_faucet, Faucet, FAUCET_PORT}, socketaddr, @@ -34,36 +35,36 @@ fn main() -> Result<(), Box> { .help("Time slice over which to limit requests to faucet"), ) .arg( - Arg::with_name("cap") - .long("cap") + Arg::with_name("per_time_cap") + .long("per-time-cap") + .alias("cap") .value_name("NUM") .takes_value(true) - .help("Request limit for time slice"), + .help("Request limit for time slice, in SOL"), + ) + .arg( + Arg::with_name("per_request_cap") + .long("per-request-cap") + .value_name("NUM") + .takes_value(true) + .help("Request limit for a single request, in SOL"), ) .get_matches(); let mint_keypair = read_keypair_file(matches.value_of("keypair").unwrap()) .expect("failed to read client keypair"); - let time_slice: Option; - if let Some(secs) = matches.value_of("slice") { - time_slice = Some(secs.to_string().parse().expect("failed to parse slice")); - } else { - time_slice = None; - } - let request_cap: Option; - if let Some(c) = matches.value_of("cap") { - request_cap = Some(c.to_string().parse().expect("failed to parse cap")); - } else { - request_cap = None; - } + let time_slice = value_of(&matches, "slice"); + let per_time_cap = lamports_of_sol(&matches, "per_time_cap"); + let per_request_cap = lamports_of_sol(&matches, "per_request_cap"); let faucet_addr = socketaddr!(0, FAUCET_PORT); let faucet = Arc::new(Mutex::new(Faucet::new( mint_keypair, time_slice, - request_cap, + per_time_cap, + per_request_cap, ))); let faucet1 = faucet.clone(); diff --git a/faucet/src/faucet.rs b/faucet/src/faucet.rs index 6bf82d5222..f8187bc8a5 100644 --- a/faucet/src/faucet.rs +++ b/faucet/src/faucet.rs @@ -62,7 +62,8 @@ pub struct Faucet { mint_keypair: Keypair, ip_cache: Vec, pub time_slice: Duration, - request_cap: u64, + per_time_cap: u64, + per_request_cap: Option, pub request_current: u64, } @@ -70,27 +71,23 @@ impl Faucet { pub fn new( mint_keypair: Keypair, time_input: Option, - request_cap_input: Option, + per_time_cap: Option, + per_request_cap: Option, ) -> Faucet { - let time_slice = match time_input { - Some(time) => Duration::new(time, 0), - None => Duration::new(TIME_SLICE, 0), - }; - let request_cap = match request_cap_input { - Some(cap) => cap, - None => REQUEST_CAP, - }; + let time_slice = Duration::new(time_input.unwrap_or(TIME_SLICE), 0); + let per_time_cap = per_time_cap.unwrap_or(REQUEST_CAP); Faucet { mint_keypair, ip_cache: Vec::new(), time_slice, - request_cap, + per_time_cap, + per_request_cap, request_current: 0, } } - pub fn check_request_limit(&mut self, request_amount: u64) -> bool { - (self.request_current + request_amount) <= self.request_cap + pub fn check_time_request_limit(&mut self, request_amount: u64) -> bool { + (self.request_current + request_amount) <= self.per_time_cap } pub fn clear_request_count(&mut self) { @@ -116,7 +113,15 @@ impl Faucet { to, blockhash, } => { - if self.check_request_limit(lamports) { + if let Some(cap) = self.per_request_cap { + if lamports > cap { + return Err(Error::new( + ErrorKind::Other, + format!("request too large; req: {} cap: {}", lamports, cap), + )); + } + } + if self.check_time_request_limit(lamports) { self.request_current += lamports; datapoint_info!( "faucet-airdrop", @@ -135,7 +140,7 @@ impl Faucet { ErrorKind::Other, format!( "token limit reached; req: {} current: {} cap: {}", - lamports, self.request_current, self.request_cap + lamports, self.request_current, self.per_time_cap ), )) } @@ -248,14 +253,15 @@ pub fn request_airdrop_transaction( pub fn run_local_faucet( mint_keypair: Keypair, sender: Sender, - request_cap_input: Option, + per_time_cap: Option, ) { thread::spawn(move || { let faucet_addr = socketaddr!(0, 0); let faucet = Arc::new(Mutex::new(Faucet::new( mint_keypair, None, - request_cap_input, + per_time_cap, + None, ))); run_faucet(faucet, faucet_addr, Some(sender)); }); @@ -312,18 +318,18 @@ mod tests { use std::time::Duration; #[test] - fn test_check_request_limit() { + fn test_check_time_request_limit() { let keypair = Keypair::new(); - let mut faucet = Faucet::new(keypair, None, Some(3)); - assert!(faucet.check_request_limit(1)); + let mut faucet = Faucet::new(keypair, None, Some(3), None); + assert!(faucet.check_time_request_limit(1)); faucet.request_current = 3; - assert!(!faucet.check_request_limit(1)); + assert!(!faucet.check_time_request_limit(1)); } #[test] fn test_clear_request_count() { let keypair = Keypair::new(); - let mut faucet = Faucet::new(keypair, None, None); + let mut faucet = Faucet::new(keypair, None, None, None); faucet.request_current += 256; assert_eq!(faucet.request_current, 256); faucet.clear_request_count(); @@ -333,7 +339,7 @@ mod tests { #[test] fn test_add_ip_to_cache() { let keypair = Keypair::new(); - let mut faucet = Faucet::new(keypair, None, None); + let mut faucet = Faucet::new(keypair, None, None, None); let ip = "127.0.0.1".parse().expect("create IpAddr from string"); assert_eq!(faucet.ip_cache.len(), 0); faucet.add_ip_to_cache(ip); @@ -344,7 +350,7 @@ mod tests { #[test] fn test_clear_ip_cache() { let keypair = Keypair::new(); - let mut faucet = Faucet::new(keypair, None, None); + let mut faucet = Faucet::new(keypair, None, None, None); let ip = "127.0.0.1".parse().expect("create IpAddr from string"); assert_eq!(faucet.ip_cache.len(), 0); faucet.add_ip_to_cache(ip); @@ -359,9 +365,10 @@ mod tests { let keypair = Keypair::new(); let time_slice: Option = None; let request_cap: Option = None; - let faucet = Faucet::new(keypair, time_slice, request_cap); + let faucet = Faucet::new(keypair, time_slice, request_cap, Some(100)); assert_eq!(faucet.time_slice, Duration::new(TIME_SLICE, 0)); - assert_eq!(faucet.request_cap, REQUEST_CAP); + assert_eq!(faucet.per_time_cap, REQUEST_CAP); + assert_eq!(faucet.per_request_cap, Some(100)); } #[test] @@ -376,7 +383,7 @@ mod tests { let mint = Keypair::new(); let mint_pubkey = mint.pubkey(); - let mut faucet = Faucet::new(mint, None, None); + let mut faucet = Faucet::new(mint, None, None, None); let tx = faucet.build_airdrop_transaction(request).unwrap(); let message = tx.message(); @@ -392,8 +399,15 @@ mod tests { let instruction: SystemInstruction = deserialize(&message.instructions[0].data).unwrap(); assert_eq!(instruction, SystemInstruction::Transfer { lamports: 2 }); + // Test per-time request cap let mint = Keypair::new(); - faucet = Faucet::new(mint, None, Some(1)); + faucet = Faucet::new(mint, None, Some(1), None); + let tx = faucet.build_airdrop_transaction(request); + assert!(tx.is_err()); + + // Test per-request cap + let mint = Keypair::new(); + faucet = Faucet::new(mint, None, None, Some(1)); let tx = faucet.build_airdrop_transaction(request); assert!(tx.is_err()); } @@ -421,7 +435,7 @@ mod tests { LittleEndian::write_u16(&mut expected_vec_with_length, expected_bytes.len() as u16); expected_vec_with_length.extend_from_slice(&expected_bytes); - let mut faucet = Faucet::new(keypair, None, None); + let mut faucet = Faucet::new(keypair, None, None, None); let response = faucet.process_faucet_request(&bytes); let response_vec = response.unwrap().to_vec(); assert_eq!(expected_vec_with_length, response_vec);