gaspersic/solana
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add request-count check and tests; fmt

Browse Source
This commit is contained in:
Tyera Eulberg
2018-06-20 00:51:04 -06:00
committed by Greg Fitzgerald
parent 64cf69045a
commit 73b0ff5b55
1 changed files with 130 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

166
src/drone.rs
View File

@@ -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<IpAddr>,
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<u64>,
@@ -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<usize, io::Error> {
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<u64> = None;
let request_cap: Option<u64> = 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);