solana/src/drone.rs

333 lines
11 KiB
Rust
Raw Normal View History

2018-06-19 23:51:04 -07:00
//! 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.
2018-07-05 21:40:09 -07:00
use influx_db_client as influxdb;
use metrics;
2018-06-19 21:39:53 -07:00
use signature::{KeyPair, PublicKey};
use std::io;
2018-06-19 23:51:04 -07:00
use std::io::{Error, ErrorKind};
2018-06-19 21:39:53 -07:00
use std::net::{IpAddr, SocketAddr, UdpSocket};
use std::time::Duration;
use thin_client::ThinClient;
use transaction::Transaction;
2018-06-19 23:51:04 -07:00
pub const TIME_SLICE: u64 = 60;
2018-07-03 15:34:13 -07:00
pub const REQUEST_CAP: u64 = 1_000_000;
2018-06-19 21:39:53 -07:00
#[derive(Serialize, Deserialize, Debug, Clone, Copy)]
2018-06-19 21:39:53 -07:00
pub enum DroneRequest {
2018-06-19 23:51:04 -07:00
GetAirdrop {
2018-06-21 19:41:47 -07:00
airdrop_request_amount: u64,
2018-06-19 23:51:04 -07:00
client_public_key: PublicKey,
},
2018-06-19 21:39:53 -07:00
}
pub struct Drone {
mint_keypair: KeyPair,
ip_cache: Vec<IpAddr>,
2018-06-19 23:51:04 -07:00
_airdrop_addr: SocketAddr,
transactions_addr: SocketAddr,
requests_addr: SocketAddr,
2018-06-19 23:51:04 -07:00
pub time_slice: Duration,
2018-06-19 21:39:53 -07:00
request_cap: u64,
2018-06-19 23:51:04 -07:00
pub request_current: u64,
2018-06-19 21:39:53 -07:00
}
impl Drone {
pub fn new(
mint_keypair: KeyPair,
2018-06-19 23:51:04 -07:00
_airdrop_addr: SocketAddr,
transactions_addr: SocketAddr,
requests_addr: SocketAddr,
2018-06-19 21:39:53 -07:00
time_input: Option<u64>,
request_cap_input: Option<u64>,
) -> Drone {
let time_slice = match time_input {
Some(time) => Duration::new(time, 0),
2018-06-19 23:51:04 -07:00
None => Duration::new(TIME_SLICE, 0),
2018-06-19 21:39:53 -07:00
};
let request_cap = match request_cap_input {
Some(cap) => cap,
2018-06-19 23:51:04 -07:00
None => REQUEST_CAP,
2018-06-19 21:39:53 -07:00
};
Drone {
mint_keypair,
ip_cache: Vec::new(),
2018-06-19 23:51:04 -07:00
_airdrop_addr,
transactions_addr,
requests_addr,
2018-06-19 21:39:53 -07:00
time_slice,
request_cap,
request_current: 0,
}
}
2018-06-21 22:11:24 -07:00
pub fn check_request_limit(&mut self, request_amount: u64) -> bool {
(self.request_current + request_amount) <= self.request_cap
2018-06-19 23:51:04 -07:00
}
pub fn clear_request_count(&mut self) {
self.request_current = 0;
}
2018-06-19 21:39:53 -07:00
pub fn add_ip_to_cache(&mut self, ip: IpAddr) {
self.ip_cache.push(ip);
}
pub fn clear_ip_cache(&mut self) {
self.ip_cache.clear();
}
pub fn check_rate_limit(&mut self, ip: IpAddr) -> Result<IpAddr, IpAddr> {
// [WIP] This is placeholder code for a proper rate limiter.
// Right now it will only allow one total drone request per IP
if self.ip_cache.contains(&ip) {
// Add proper error handling here
Err(ip)
2018-06-19 23:51:04 -07:00
} else {
2018-06-19 21:39:53 -07:00
self.add_ip_to_cache(ip);
Ok(ip)
}
}
pub fn send_airdrop(&mut self, req: DroneRequest) -> Result<usize, io::Error> {
2018-06-21 22:11:24 -07:00
let tx: Transaction;
let request_amount: u64;
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,
2018-06-21 22:11:24 -07:00
requests_socket,
self.transactions_addr,
2018-06-21 22:11:24 -07:00
transactions_socket,
);
let last_id = client.get_last_id();
match req {
DroneRequest::GetAirdrop {
airdrop_request_amount,
client_public_key,
} => {
2018-07-11 13:40:46 -07:00
request_amount = airdrop_request_amount;
2018-06-21 22:11:24 -07:00
tx = Transaction::new(
&self.mint_keypair,
2018-06-19 23:51:04 -07:00
client_public_key,
2018-06-21 22:11:24 -07:00
airdrop_request_amount as i64,
last_id,
);
2018-06-19 21:39:53 -07:00
}
2018-06-21 22:11:24 -07:00
}
if self.check_request_limit(request_amount) {
self.request_current += request_amount;
2018-07-05 21:40:09 -07:00
metrics::submit(
influxdb::Point::new("drone")
.add_tag("op", influxdb::Value::String("airdrop".to_string()))
.add_field(
"request_amount",
influxdb::Value::Integer(request_amount as i64),
)
.add_field(
"request_current",
influxdb::Value::Integer(self.request_current as i64),
)
.to_owned(),
);
2018-07-11 13:40:46 -07:00
client.transfer_signed(&tx)
2018-06-19 23:51:04 -07:00
} else {
2018-06-21 19:41:47 -07:00
Err(Error::new(ErrorKind::Other, "token limit reached"))
2018-06-19 21:39:53 -07:00
}
}
}
2018-07-05 21:40:09 -07:00
impl Drop for Drone {
fn drop(&mut self) {
metrics::flush();
}
}
2018-06-19 21:39:53 -07:00
#[cfg(test)]
mod tests {
use bank::Bank;
use crdt::{get_ip_addr, TestNode};
2018-06-21 19:41:47 -07:00
use drone::{Drone, DroneRequest, REQUEST_CAP, TIME_SLICE};
2018-07-02 15:24:40 -07:00
use fullnode::FullNode;
2018-06-19 21:39:53 -07:00
use logger;
use mint::Mint;
use service::Service;
2018-06-19 21:39:53 -07:00
use signature::{KeyPair, KeyPairUtil};
use std::io::sink;
use std::net::{SocketAddr, UdpSocket};
2018-06-19 23:51:04 -07:00
use std::sync::atomic::{AtomicBool, Ordering};
2018-06-20 11:50:14 -07:00
use std::sync::Arc;
2018-06-19 21:39:53 -07:00
use std::thread::sleep;
use std::time::Duration;
use thin_client::ThinClient;
2018-06-19 23:51:04 -07:00
#[test]
2018-06-21 19:41:47 -07:00
fn test_check_request_limit() {
2018-06-19 23:51:04 -07:00
let keypair = KeyPair::new();
let mut addr: SocketAddr = "0.0.0.0:9900".parse().unwrap();
addr.set_ip(get_ip_addr().unwrap());
let transactions_addr = "0.0.0.0:0".parse().unwrap();
let requests_addr = "0.0.0.0:0".parse().unwrap();
2018-06-24 00:58:44 -07:00
let mut drone = Drone::new(
keypair,
addr,
transactions_addr,
requests_addr,
None,
Some(3),
);
2018-06-21 22:11:24 -07:00
assert!(drone.check_request_limit(1));
drone.request_current = 3;
assert!(!drone.check_request_limit(1));
2018-06-19 23:51:04 -07:00
}
#[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 transactions_addr = "0.0.0.0:0".parse().unwrap();
let requests_addr = "0.0.0.0:0".parse().unwrap();
let mut drone = Drone::new(keypair, addr, transactions_addr, requests_addr, None, None);
2018-06-19 23:51:04 -07:00
drone.request_current = drone.request_current + 256;
assert_eq!(drone.request_current, 256);
2018-06-19 23:51:04 -07:00
drone.clear_request_count();
assert_eq!(drone.request_current, 0);
2018-06-19 23:51:04 -07:00
}
2018-06-19 21:39:53 -07:00
#[test]
fn test_add_ip_to_cache() {
let keypair = KeyPair::new();
let mut addr: SocketAddr = "0.0.0.0:9900".parse().unwrap();
addr.set_ip(get_ip_addr().unwrap());
let transactions_addr = "0.0.0.0:0".parse().unwrap();
let requests_addr = "0.0.0.0:0".parse().unwrap();
let mut drone = Drone::new(keypair, addr, transactions_addr, requests_addr, None, None);
2018-06-19 21:39:53 -07:00
let ip = "127.0.0.1".parse().expect("create IpAddr from string");
assert_eq!(drone.ip_cache.len(), 0);
drone.add_ip_to_cache(ip);
assert_eq!(drone.ip_cache.len(), 1);
assert!(drone.ip_cache.contains(&ip));
}
#[test]
fn test_clear_ip_cache() {
let keypair = KeyPair::new();
let mut addr: SocketAddr = "0.0.0.0:9900".parse().unwrap();
addr.set_ip(get_ip_addr().unwrap());
let transactions_addr = "0.0.0.0:0".parse().unwrap();
let requests_addr = "0.0.0.0:0".parse().unwrap();
let mut drone = Drone::new(keypair, addr, transactions_addr, requests_addr, None, None);
2018-06-19 21:39:53 -07:00
let ip = "127.0.0.1".parse().expect("create IpAddr from string");
assert_eq!(drone.ip_cache.len(), 0);
drone.add_ip_to_cache(ip);
assert_eq!(drone.ip_cache.len(), 1);
drone.clear_ip_cache();
assert_eq!(drone.ip_cache.len(), 0);
assert!(drone.ip_cache.is_empty());
}
#[test]
fn test_drone_default_init() {
let keypair = KeyPair::new();
let mut addr: SocketAddr = "0.0.0.0:9900".parse().unwrap();
addr.set_ip(get_ip_addr().unwrap());
let transactions_addr = "0.0.0.0:0".parse().unwrap();
let requests_addr = "0.0.0.0:0".parse().unwrap();
2018-06-19 21:39:53 -07:00
let time_slice: Option<u64> = None;
let request_cap: Option<u64> = None;
2018-06-24 00:58:44 -07:00
let drone = Drone::new(
keypair,
addr,
transactions_addr,
requests_addr,
time_slice,
request_cap,
);
2018-06-19 21:39:53 -07:00
assert_eq!(drone.time_slice, Duration::new(TIME_SLICE, 0));
assert_eq!(drone.request_cap, REQUEST_CAP);
}
#[test]
2018-07-03 11:06:34 -07:00
#[ignore]
2018-06-19 21:39:53 -07:00
fn test_send_airdrop() {
2018-06-21 19:41:47 -07:00
const SMALL_BATCH: i64 = 50;
const TPS_BATCH: i64 = 5_000_000;
2018-06-19 21:39:53 -07:00
logger::setup();
let leader = TestNode::new();
let alice = Mint::new(10_000_000);
let bank = Bank::new(&alice);
let bob_pubkey = KeyPair::new().pubkey();
let carlos_pubkey = KeyPair::new().pubkey();
let exit = Arc::new(AtomicBool::new(false));
let leader_data = leader.data.clone();
2018-06-19 21:39:53 -07:00
2018-07-02 15:24:40 -07:00
let server = FullNode::new_leader(
2018-06-19 21:39:53 -07:00
bank,
2018-06-27 12:35:58 -07:00
0,
None,
2018-06-19 21:39:53 -07:00
Some(Duration::from_millis(30)),
leader,
2018-06-19 21:39:53 -07:00
exit.clone(),
sink(),
);
2018-07-03 11:06:34 -07:00
//TODO: this seems unstable
2018-06-19 21:39:53 -07:00
sleep(Duration::from_millis(900));
let mut addr: SocketAddr = "0.0.0.0:9900".parse().expect("bind to drone socket");
addr.set_ip(get_ip_addr().expect("drone get_ip_addr"));
2018-06-19 23:51:04 -07:00
let mut drone = Drone::new(
alice.keypair(),
addr,
2018-07-09 17:55:11 -07:00
leader_data.contact_info.tpu,
leader_data.contact_info.rpu,
2018-06-19 23:51:04 -07:00
None,
Some(150_000),
2018-06-19 23:51:04 -07:00
);
2018-06-19 21:39:53 -07:00
2018-06-19 23:51:04 -07:00
let bob_req = DroneRequest::GetAirdrop {
2018-06-21 19:41:47 -07:00
airdrop_request_amount: 50,
2018-06-19 23:51:04 -07:00
client_public_key: bob_pubkey,
};
2018-06-19 21:39:53 -07:00
let bob_result = drone.send_airdrop(bob_req).expect("send airdrop test");
assert!(bob_result > 0);
2018-06-19 23:51:04 -07:00
let carlos_req = DroneRequest::GetAirdrop {
2018-06-21 19:41:47 -07:00
airdrop_request_amount: 5_000_000,
2018-06-19 23:51:04 -07:00
client_public_key: carlos_pubkey,
};
2018-06-19 21:39:53 -07:00
let carlos_result = drone.send_airdrop(carlos_req).expect("send airdrop test");
assert!(carlos_result > 0);
let requests_socket = UdpSocket::bind("0.0.0.0:0").expect("drone bind to requests socket");
2018-06-22 13:41:43 -07:00
let transactions_socket =
UdpSocket::bind("0.0.0.0:0").expect("drone bind to transactions socket");
2018-06-24 00:58:44 -07:00
2018-06-19 21:39:53 -07:00
let mut client = ThinClient::new(
2018-07-09 17:55:11 -07:00
leader_data.contact_info.rpu,
2018-06-19 21:39:53 -07:00
requests_socket,
2018-07-09 17:55:11 -07:00
leader_data.contact_info.tpu,
2018-06-19 21:39:53 -07:00
transactions_socket,
);
let bob_balance = client.poll_get_balance(&bob_pubkey);
2018-06-21 19:41:47 -07:00
info!("Small request balance: {:?}", bob_balance);
2018-06-19 21:39:53 -07:00
assert_eq!(bob_balance.unwrap(), SMALL_BATCH);
let carlos_balance = client.poll_get_balance(&carlos_pubkey);
2018-06-21 19:41:47 -07:00
info!("TPS request balance: {:?}", carlos_balance);
2018-06-19 21:39:53 -07:00
assert_eq!(carlos_balance.unwrap(), TPS_BATCH);
exit.store(true, Ordering::Relaxed);
server.join().unwrap();
2018-06-19 21:39:53 -07:00
}
}