Add allowed-ip list to faucet (#16891)

This commit is contained in:
Tyera Eulberg 2021-04-27 19:31:32 -06:00 committed by GitHub
parent b468ead1b1
commit 36574c30ef
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 62 additions and 7 deletions

View File

@ -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();

View File

@ -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();