From 73b0ff5b55ec5bd28bc6aec8cded04cdb6a9abb2 Mon Sep 17 00:00:00 2001 From: Tyera Eulberg Date: Wed, 20 Jun 2018 00:51:04 -0600 Subject: [PATCH] Add request-count check and tests; fmt --- src/drone.rs | 166 ++++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 130 insertions(+), 36 deletions(-) diff --git a/src/drone.rs b/src/drone.rs index b273025185..e50cd10a68 100644 --- a/src/drone.rs +++ b/src/drone.rs @@ -1,18 +1,28 @@ +//! The `drone` module provides an object for launching a Solana Drone, +//! which is the custodian of any remaining tokens in a mint. +//! The Solana Drone builds and send airdrop transactions, +//! checking requests against a request cap for a given time time_slice +//! and (to come) an IP rate limit. + use signature::{KeyPair, PublicKey}; use std::io; +use std::io::{Error, ErrorKind}; use std::net::{IpAddr, SocketAddr, UdpSocket}; use std::time::Duration; use thin_client::ThinClient; use transaction::Transaction; -pub const TIME_SLICE: u64 = 1; +pub const TIME_SLICE: u64 = 60; pub const REQUEST_CAP: u64 = 500; pub const SMALL_BATCH: i64 = 50; pub const TPS_BATCH: i64 = 5_000_000; #[derive(Serialize, Deserialize, Debug)] pub enum DroneRequest { - GetAirdrop { request_type: DroneRequestType, client_public_key: PublicKey }, + GetAirdrop { + request_type: DroneRequestType, + client_public_key: PublicKey, + }, } #[derive(Serialize, Deserialize, Debug)] @@ -24,18 +34,18 @@ pub enum DroneRequestType { pub struct Drone { mint_keypair: KeyPair, ip_cache: Vec, - airdrop_addr: SocketAddr, + _airdrop_addr: SocketAddr, transactions_addr: SocketAddr, requests_addr: SocketAddr, - time_slice: Duration, + pub time_slice: Duration, request_cap: u64, - request_current: u64, + pub request_current: u64, } impl Drone { pub fn new( mint_keypair: KeyPair, - airdrop_addr: SocketAddr, + _airdrop_addr: SocketAddr, transactions_addr: SocketAddr, requests_addr: SocketAddr, time_input: Option, @@ -43,16 +53,16 @@ impl Drone { ) -> Drone { let time_slice = match time_input { Some(time) => Duration::new(time, 0), - None => Duration::new(TIME_SLICE, 0) + None => Duration::new(TIME_SLICE, 0), }; let request_cap = match request_cap_input { Some(cap) => cap, - None => REQUEST_CAP + None => REQUEST_CAP, }; Drone { mint_keypair, ip_cache: Vec::new(), - airdrop_addr, + _airdrop_addr, transactions_addr, requests_addr, time_slice, @@ -61,6 +71,18 @@ impl Drone { } } + pub fn check_request_count(&mut self) -> bool { + if self.request_current <= self.request_cap { + true + } else { + false + } + } + + pub fn clear_request_count(&mut self) { + self.request_current = 0; + } + pub fn add_ip_to_cache(&mut self, ip: IpAddr) { self.ip_cache.push(ip); } @@ -75,42 +97,55 @@ impl Drone { if self.ip_cache.contains(&ip) { // Add proper error handling here Err(ip) - } - else { + } else { self.add_ip_to_cache(ip); Ok(ip) } - // Test for quota limit met } pub fn send_airdrop(&mut self, req: DroneRequest) -> Result { - let tx: Transaction; + self.request_current += 1; - let requests_socket = UdpSocket::bind("0.0.0.0:0").unwrap(); - let transactions_socket = UdpSocket::bind("0.0.0.0:0").unwrap(); + if self.check_request_count() { + let tx: Transaction; + let requests_socket = UdpSocket::bind("0.0.0.0:0").unwrap(); + let transactions_socket = UdpSocket::bind("0.0.0.0:0").unwrap(); - let mut client = ThinClient::new( - self.requests_addr, - requests_socket, - self.transactions_addr, - transactions_socket, - ); - let last_id = client.get_last_id(); + let mut client = ThinClient::new( + self.requests_addr, + requests_socket, + self.transactions_addr, + transactions_socket, + ); + let last_id = client.get_last_id(); - match req { - DroneRequest::GetAirdrop { request_type, client_public_key } => { - match request_type { + match req { + DroneRequest::GetAirdrop { + request_type, + client_public_key, + } => match request_type { DroneRequestType::SmallBatch => { - tx = Transaction::new(&self.mint_keypair, client_public_key, SMALL_BATCH, last_id); + tx = Transaction::new( + &self.mint_keypair, + client_public_key, + SMALL_BATCH, + last_id, + ); } DroneRequestType::TPSBatch => { - tx = Transaction::new(&self.mint_keypair, client_public_key, TPS_BATCH, last_id); + tx = Transaction::new( + &self.mint_keypair, + client_public_key, + TPS_BATCH, + last_id, + ); } - } + }, } + client.transfer_signed(tx) + } else { + Err(Error::new(ErrorKind::Other, "request limit reached")) } - - client.transfer_signed(tx) } } @@ -118,19 +153,58 @@ impl Drone { mod tests { use bank::Bank; use crdt::{get_ip_addr, TestNode}; - use drone::{Drone, DroneRequest, DroneRequestType, TIME_SLICE, REQUEST_CAP, SMALL_BATCH, TPS_BATCH}; + use drone::{Drone, DroneRequest, DroneRequestType, REQUEST_CAP, SMALL_BATCH, TIME_SLICE, + TPS_BATCH}; use logger; use mint::Mint; use server::Server; use signature::{KeyPair, KeyPairUtil}; use std::io::sink; use std::net::{SocketAddr, UdpSocket}; - use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; use std::thread::sleep; use std::time::Duration; use thin_client::ThinClient; + #[test] + fn test_check_request_count() { + let keypair = KeyPair::new(); + let mut addr: SocketAddr = "0.0.0.0:9900".parse().unwrap(); + addr.set_ip(get_ip_addr().unwrap()); + let mut transactions_addr = addr.clone(); + transactions_addr.set_port(8000); + let mut requests_addr = addr.clone(); + requests_addr.set_port(8003); + let mut drone = Drone::new( + keypair, + addr, + transactions_addr, + requests_addr, + None, + Some(3), + ); + assert!(drone.check_request_count()); + drone.request_current = 4; + assert!(!drone.check_request_count()); + } + + #[test] + fn test_clear_request_count() { + let keypair = KeyPair::new(); + let mut addr: SocketAddr = "0.0.0.0:9900".parse().unwrap(); + addr.set_ip(get_ip_addr().unwrap()); + let mut transactions_addr = addr.clone(); + transactions_addr.set_port(8000); + let mut requests_addr = addr.clone(); + requests_addr.set_port(8003); + let mut drone = Drone::new(keypair, addr, transactions_addr, requests_addr, None, None); + drone.request_current = drone.request_current + 256; + assert!(drone.request_current == 256); + drone.clear_request_count(); + assert!(drone.request_current == 0); + } + #[test] fn test_add_ip_to_cache() { let keypair = KeyPair::new(); @@ -178,7 +252,14 @@ mod tests { requests_addr.set_port(8003); let time_slice: Option = None; let request_cap: Option = None; - let drone = Drone::new(keypair, addr, transactions_addr, requests_addr, time_slice, request_cap); + let drone = Drone::new( + keypair, + addr, + transactions_addr, + requests_addr, + time_slice, + request_cap, + ); assert_eq!(drone.time_slice, Duration::new(TIME_SLICE, 0)); assert_eq!(drone.request_cap, REQUEST_CAP); } @@ -210,13 +291,26 @@ mod tests { let mut addr: SocketAddr = "0.0.0.0:9900".parse().unwrap(); addr.set_ip(get_ip_addr().unwrap()); - let mut drone = Drone::new(alice.keypair(), addr, leader.data.transactions_addr, leader.data.requests_addr, None, None); + let mut drone = Drone::new( + alice.keypair(), + addr, + leader.data.transactions_addr, + leader.data.requests_addr, + None, + None, + ); - let bob_req = DroneRequest::GetAirdrop{ request_type: DroneRequestType::SmallBatch, client_public_key: bob_pubkey }; + let bob_req = DroneRequest::GetAirdrop { + request_type: DroneRequestType::SmallBatch, + client_public_key: bob_pubkey, + }; let bob_result = drone.send_airdrop(bob_req).expect("send airdrop test"); assert!(bob_result > 0); - let carlos_req = DroneRequest::GetAirdrop{ request_type: DroneRequestType::TPSBatch, client_public_key: carlos_pubkey }; + let carlos_req = DroneRequest::GetAirdrop { + request_type: DroneRequestType::TPSBatch, + client_public_key: carlos_pubkey, + }; let carlos_result = drone.send_airdrop(carlos_req).expect("send airdrop test"); assert!(carlos_result > 0);