YGGverse/aquatic
1
0
Fork 0
mirror of https://github.com/YGGverse/aquatic.git synced 2026-03-31 17:55:36 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Run rustfmt, clean up aquatic_http_protocol/Cargo.toml

Browse source
This commit is contained in:
Joakim Frostegård 2021-08-15 22:26:11 +02:00
parent 0cc312a78d
commit d0e716f80b
65 changed files with 1754 additions and 2590 deletions
Download patch file Download diff file Expand all files Collapse all files

65
aquatic/src/main.rs
View file

@ -1,43 +1,39 @@
use aquatic_cli_helpers::{Options, run_app_with_cli_and_config, print_help};
use aquatic_cli_helpers::{print_help, run_app_with_cli_and_config, Options};
use aquatic_http::config::Config as HttpConfig;
use aquatic_udp::config::Config as UdpConfig;
use aquatic_ws::config::Config as WsConfig;
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
const APP_NAME: &str = "aquatic: BitTorrent tracker";
fn main(){
::std::process::exit(match run(){
fn main() {
::std::process::exit(match run() {
Ok(()) => 0,
Err(None) => {
print_help(|| gen_info(), None);
0
},
Err(opt_err@Some(_)) => {
}
Err(opt_err @ Some(_)) => {
print_help(|| gen_info(), opt_err);
1
},
}
})
}
fn run() -> Result<(), Option<String>>{
fn run() -> Result<(), Option<String>> {
let mut arg_iter = ::std::env::args().skip(1);
let protocol = if let Some(protocol) = arg_iter.next(){
let protocol = if let Some(protocol) = arg_iter.next() {
protocol
} else {
return Err(None);
};
let options = match Options::parse_args(arg_iter){
let options = match Options::parse_args(arg_iter) {
Ok(options) => options,
Err(opt_err) => {
return Err(opt_err);
@ -45,44 +41,37 @@ fn run() -> Result<(), Option<String>>{
};
match protocol.as_str() {
"udp" => {
run_app_with_cli_and_config::<UdpConfig>(
aquatic_udp::APP_NAME,
aquatic_udp::run,
Some(options),
)
},
"http" => {
run_app_with_cli_and_config::<HttpConfig>(
aquatic_http::APP_NAME,
aquatic_http::run,
Some(options),
)
},
"ws" => {
run_app_with_cli_and_config::<WsConfig>(
aquatic_ws::APP_NAME,
aquatic_ws::run,
Some(options),
)
},
"udp" => run_app_with_cli_and_config::<UdpConfig>(
aquatic_udp::APP_NAME,
aquatic_udp::run,
Some(options),
),
"http" => run_app_with_cli_and_config::<HttpConfig>(
aquatic_http::APP_NAME,
aquatic_http::run,
Some(options),
),
"ws" => run_app_with_cli_and_config::<WsConfig>(
aquatic_ws::APP_NAME,
aquatic_ws::run,
Some(options),
),
arg => {
let opt_err = if arg == "-h" || arg == "--help" {
None
} else if arg.chars().next() == Some('-'){
} else if arg.chars().next() == Some('-') {
Some("First argument must be protocol".to_string())
} else {
Some("Invalid protocol".to_string())
};
return Err(opt_err)
},
return Err(opt_err);
}
}
Ok(())
}
fn gen_info() -> String {
let mut info = String::new();

101
aquatic_cli_helpers/src/lib.rs
View file

@ -2,10 +2,9 @@ use std::fs::File;
use std::io::Read;
use anyhow::Context;
use serde::{Serialize, Deserialize, de::DeserializeOwned};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use simplelog::{ConfigBuilder, LevelFilter, TermLogger, TerminalMode};
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum LogLevel {
@ -14,57 +13,50 @@ pub enum LogLevel {
Warn,
Info,
Debug,
Trace
Trace,
}
impl Default for LogLevel {
fn default() -> Self {
Self::Error
}
}
pub trait Config: Default + Serialize + DeserializeOwned {
fn get_log_level(&self) -> Option<LogLevel> {
None
}
}
#[derive(Debug, Default)]
pub struct Options {
config_file: Option<String>,
print_config: bool,
}
impl Options {
pub fn parse_args<I>(
mut arg_iter: I
) -> Result<Options, Option<String>>
where I: Iterator<Item = String>
pub fn parse_args<I>(mut arg_iter: I) -> Result<Options, Option<String>>
where
I: Iterator<Item = String>,
{
let mut options = Options::default();
loop {
if let Some(arg) = arg_iter.next(){
match arg.as_str(){
if let Some(arg) = arg_iter.next() {
match arg.as_str() {
"-c" | "--config-file" => {
if let Some(path) = arg_iter.next(){
if let Some(path) = arg_iter.next() {
options.config_file = Some(path);
} else {
return Err(
Some("No config file path given".to_string())
);
return Err(Some("No config file path given".to_string()));
}
},
}
"-p" | "--print-config" => {
options.print_config = true;
},
}
"-h" | "--help" => {
return Err(None);
},
}
_ => {
return Err(Some("Unrecognized argument".to_string()));
}
@ -78,31 +70,34 @@ impl Options {
}
}
pub fn run_app_with_cli_and_config<T>(
app_title: &str,
// Function that takes config file and runs application
app_fn: fn(T) -> anyhow::Result<()>,
opts: Option<Options>,
) where T: Config {
) where
T: Config,
{
::std::process::exit(match run_inner(app_title, app_fn, opts) {
Ok(()) => 0,
Err(err) => {
eprintln!("Error: {:#}", err);
1
},
}
})
}
fn run_inner<T>(
app_title: &str,
// Function that takes config file and runs application
app_fn: fn(T) -> anyhow::Result<()>,
// Possibly preparsed options
options: Option<Options>,
) -> anyhow::Result<()> where T: Config {
) -> anyhow::Result<()>
where
T: Config,
{
let options = if let Some(options) = options {
options
} else {
@ -110,18 +105,14 @@ fn run_inner<T>(
let app_path = arg_iter.next().unwrap();
match Options::parse_args(arg_iter){
match Options::parse_args(arg_iter) {
Ok(options) => options,
Err(opt_err) => {
let gen_info = || format!(
"{}\n\nUsage: {} [OPTIONS]",
app_title,
app_path
);
let gen_info = || format!("{}\n\nUsage: {} [OPTIONS]", app_title, app_path);
print_help(gen_info, opt_err);
return Ok(())
return Ok(());
}
}
};
@ -137,7 +128,7 @@ fn run_inner<T>(
T::default()
};
if let Some(log_level) = config.get_log_level(){
if let Some(log_level) = config.get_log_level() {
start_logger(log_level)?;
}
@ -145,11 +136,10 @@ fn run_inner<T>(
}
}
pub fn print_help<F>(
info_generator: F,
opt_error: Option<String>
) where F: FnOnce() -> String {
pub fn print_help<F>(info_generator: F, opt_error: Option<String>)
where
F: FnOnce() -> String,
{
println!("{}", info_generator());
println!("\nOptions:");
@ -162,36 +152,30 @@ pub fn print_help<F>(
}
}
fn config_from_toml_file<T>(path: String) -> anyhow::Result<T>
where T: DeserializeOwned
where
T: DeserializeOwned,
{
let mut file = File::open(path.clone()).with_context(||
format!("Couldn't open config file {}", path.clone())
)?;
let mut file = File::open(path.clone())
.with_context(|| format!("Couldn't open config file {}", path.clone()))?;
let mut data = String::new();
file.read_to_string(&mut data).with_context(||
format!("Couldn't read config file {}", path.clone())
)?;
file.read_to_string(&mut data)
.with_context(|| format!("Couldn't read config file {}", path.clone()))?;
toml::from_str(&data).with_context(||
format!("Couldn't parse config file {}", path.clone())
)
toml::from_str(&data).with_context(|| format!("Couldn't parse config file {}", path.clone()))
}
fn default_config_as_toml<T>() -> String
where T: Default + Serialize
where
T: Default + Serialize,
{
toml::to_string_pretty(&T::default())
.expect("Could not serialize default config to toml")
toml::to_string_pretty(&T::default()).expect("Could not serialize default config to toml")
}
fn start_logger(log_level: LogLevel) -> ::anyhow::Result<()> {
let level_filter = match log_level{
let level_filter = match log_level {
LogLevel::Off => LevelFilter::Off,
LogLevel::Error => LevelFilter::Error,
LogLevel::Warn => LevelFilter::Warn,
@ -206,11 +190,8 @@ fn start_logger(log_level: LogLevel) -> ::anyhow::Result<()> {
.set_location_level(LevelFilter::Off)
.build();
TermLogger::init(
level_filter,
simplelog_config,
TerminalMode::Stderr
).context("Couldn't initialize logger")?;
TermLogger::init(level_filter, simplelog_config, TerminalMode::Stderr)
.context("Couldn't initialize logger")?;
Ok(())
}

34
aquatic_common/src/lib.rs
View file

@ -1,10 +1,9 @@
use std::time::{Duration, Instant};
use std::net::IpAddr;
use std::time::{Duration, Instant};
use indexmap::IndexMap;
use rand::Rng;
/// Peer or connection valid until this instant
///
/// Used instead of "last seen" or similar to hopefully prevent arithmetic
@ -12,7 +11,6 @@ use rand::Rng;
#[derive(Debug, Clone, Copy)]
pub struct ValidUntil(pub Instant);
impl ValidUntil {
#[inline]
pub fn new(offset_seconds: u64) -> Self {
@ -20,7 +18,6 @@ impl ValidUntil {
}
}
/// Extract response peers
///
/// If there are more peers in map than `max_num_peers_to_take`, do a
@ -34,17 +31,18 @@ pub fn extract_response_peers<K, V, R, F>(
peer_map: &IndexMap<K, V>,
max_num_peers_to_take: usize,
sender_peer_map_key: K,
peer_conversion_function: F
peer_conversion_function: F,
) -> Vec<R>
where
K: Eq + ::std::hash::Hash,
F: Fn(&V) -> R
where
K: Eq + ::std::hash::Hash,
F: Fn(&V) -> R,
{
let peer_map_len = peer_map.len();
if peer_map_len <= max_num_peers_to_take + 1 {
peer_map.iter()
.filter_map(|(k, v)|{
peer_map
.iter()
.filter_map(|(k, v)| {
if *k == sender_peer_map_key {
None
} else {
@ -56,12 +54,9 @@ pub fn extract_response_peers<K, V, R, F>(
let half_num_to_take = max_num_peers_to_take / 2;
let half_peer_map_len = peer_map_len / 2;
let offset_first_half = rng.gen_range(
0..(half_peer_map_len + (peer_map_len % 2)) - half_num_to_take
);
let offset_second_half = rng.gen_range(
half_peer_map_len..peer_map_len - half_num_to_take
);
let offset_first_half =
rng.gen_range(0..(half_peer_map_len + (peer_map_len % 2)) - half_num_to_take);
let offset_second_half = rng.gen_range(half_peer_map_len..peer_map_len - half_num_to_take);
let end_first_half = offset_first_half + half_num_to_take;
let end_second_half = offset_second_half + half_num_to_take + (max_num_peers_to_take % 2);
@ -69,14 +64,14 @@ pub fn extract_response_peers<K, V, R, F>(
let mut peers: Vec<R> = Vec::with_capacity(max_num_peers_to_take);
for i in offset_first_half..end_first_half {
if let Some((k, peer)) = peer_map.get_index(i){
if let Some((k, peer)) = peer_map.get_index(i) {
if *k != sender_peer_map_key {
peers.push(peer_conversion_function(peer))
}
}
}
for i in offset_second_half..end_second_half {
if let Some((k, peer)) = peer_map.get_index(i){
if let Some((k, peer)) = peer_map.get_index(i) {
if *k != sender_peer_map_key {
peers.push(peer_conversion_function(peer))
}
@ -87,11 +82,10 @@ pub fn extract_response_peers<K, V, R, F>(
}
}
#[inline]
pub fn convert_ipv4_mapped_ipv6(ip_address: IpAddr) -> IpAddr {
if let IpAddr::V6(ip) = ip_address {
if let [0, 0, 0, 0, 0, 0xffff, ..] = ip.segments(){
if let [0, 0, 0, 0, 0, 0xffff, ..] = ip.segments() {
ip.to_ipv4().expect("convert ipv4-mapped ip").into()
} else {
ip_address

10
aquatic_http/src/bin/main.rs
View file

@ -1,15 +1,9 @@
use aquatic_cli_helpers::run_app_with_cli_and_config;
use aquatic_http::config::Config;
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
fn main(){
run_app_with_cli_and_config::<Config>(
aquatic_http::APP_NAME,
aquatic_http::run,
None
)
fn main() {
run_app_with_cli_and_config::<Config>(aquatic_http::APP_NAME, aquatic_http::run, None)
}

57
aquatic_http/src/lib/common.rs
View file

@ -1,32 +1,29 @@
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::Arc;
use crossbeam_channel::{Receiver, Sender};
use either::Either;
use crossbeam_channel::{Sender, Receiver};
use hashbrown::HashMap;
use indexmap::IndexMap;
use log::error;
use mio::Token;
use parking_lot::Mutex;
use smartstring::{SmartString, LazyCompact};
use smartstring::{LazyCompact, SmartString};
pub use aquatic_common::{ValidUntil, convert_ipv4_mapped_ipv6};
pub use aquatic_common::{convert_ipv4_mapped_ipv6, ValidUntil};
use aquatic_http_protocol::common::*;
use aquatic_http_protocol::request::Request;
use aquatic_http_protocol::response::{Response, ResponsePeer};
pub const LISTENER_TOKEN: Token = Token(0);
pub const CHANNEL_TOKEN: Token = Token(1);
pub trait Ip: ::std::fmt::Debug + Copy + Eq + ::std::hash::Hash {}
impl Ip for Ipv4Addr {}
impl Ip for Ipv6Addr {}
#[derive(Clone, Copy, Debug)]
pub struct ConnectionMeta {
/// Index of socket worker responsible for this connection. Required for
@ -36,7 +33,6 @@ pub struct ConnectionMeta {
pub poll_token: Token,
}
#[derive(Clone, Copy, Debug)]
pub struct PeerConnectionMeta<I: Ip> {
pub worker_index: usize,
@ -44,24 +40,19 @@ pub struct PeerConnectionMeta<I: Ip> {
pub peer_ip_address: I,
}
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum PeerStatus {
Seeding,
Leeching,
Stopped
Stopped,
}
impl PeerStatus {
/// Determine peer status from announce event and number of bytes left.
///
/// Likely, the last branch will be taken most of the time.
#[inline]
pub fn from_event_and_bytes_left(
event: AnnounceEvent,
opt_bytes_left: Option<usize>
) -> Self {
pub fn from_event_and_bytes_left(event: AnnounceEvent, opt_bytes_left: Option<usize>) -> Self {
if let AnnounceEvent::Stopped = event {
Self::Stopped
} else if let Some(0) = opt_bytes_left {
@ -72,7 +63,6 @@ impl PeerStatus {
}
}
#[derive(Debug, Clone, Copy)]
pub struct Peer<I: Ip> {
pub connection_meta: PeerConnectionMeta<I>,
@ -81,35 +71,30 @@ pub struct Peer<I: Ip> {
pub valid_until: ValidUntil,
}
impl <I: Ip>Peer<I> {
impl<I: Ip> Peer<I> {
pub fn to_response_peer(&self) -> ResponsePeer<I> {
ResponsePeer {
ip_address: self.connection_meta.peer_ip_address,
port: self.port
port: self.port,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct PeerMapKey<I: Ip> {
pub peer_id: PeerId,
pub ip_or_key: Either<I, SmartString<LazyCompact>>
pub ip_or_key: Either<I, SmartString<LazyCompact>>,
}
pub type PeerMap<I> = IndexMap<PeerMapKey<I>, Peer<I>>;
pub struct TorrentData<I: Ip> {
pub peers: PeerMap<I>,
pub num_seeders: usize,
pub num_leechers: usize,
}
impl <I: Ip> Default for TorrentData<I> {
impl<I: Ip> Default for TorrentData<I> {
#[inline]
fn default() -> Self {
Self {
@ -120,23 +105,19 @@ impl <I: Ip> Default for TorrentData<I> {
}
}
pub type TorrentMap<I> = HashMap<InfoHash, TorrentData<I>>;
#[derive(Default)]
pub struct TorrentMaps {
pub ipv4: TorrentMap<Ipv4Addr>,
pub ipv6: TorrentMap<Ipv6Addr>,
}
#[derive(Clone)]
pub struct State {
pub torrent_maps: Arc<Mutex<TorrentMaps>>,
}
impl Default for State {
fn default() -> Self {
Self {
@ -145,39 +126,27 @@ impl Default for State {
}
}
pub type RequestChannelSender = Sender<(ConnectionMeta, Request)>;
pub type RequestChannelReceiver = Receiver<(ConnectionMeta, Request)>;
pub type ResponseChannelReceiver = Receiver<(ConnectionMeta, Response)>;
#[derive(Clone)]
pub struct ResponseChannelSender {
senders: Vec<Sender<(ConnectionMeta, Response)>>,
}
impl ResponseChannelSender {
pub fn new(
senders: Vec<Sender<(ConnectionMeta, Response)>>,
) -> Self {
Self {
senders,
}
pub fn new(senders: Vec<Sender<(ConnectionMeta, Response)>>) -> Self {
Self { senders }
}
#[inline]
pub fn send(
&self,
meta: ConnectionMeta,
message: Response
){
if let Err(err) = self.senders[meta.worker_index].send((meta, message)){
pub fn send(&self, meta: ConnectionMeta, message: Response) {
if let Err(err) = self.senders[meta.worker_index].send((meta, message)) {
error!("ResponseChannelSender: couldn't send message: {:?}", err);
}
}
}
pub type SocketWorkerStatus = Option<Result<(), String>>;
pub type SocketWorkerStatuses = Arc<Mutex<Vec<SocketWorkerStatus>>>;

25
aquatic_http/src/lib/config.rs
View file

@ -1,10 +1,9 @@
use std::net::SocketAddr;
use serde::{Serialize, Deserialize};
use serde::{Deserialize, Serialize};
use aquatic_cli_helpers::LogLevel;
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct Config {
@ -24,14 +23,12 @@ pub struct Config {
pub privileges: PrivilegeConfig,
}
impl aquatic_cli_helpers::Config for Config {
fn get_log_level(&self) -> Option<LogLevel>{
fn get_log_level(&self) -> Option<LogLevel> {
Some(self.log_level)
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct TlsConfig {
@ -40,7 +37,6 @@ pub struct TlsConfig {
pub tls_pkcs12_password: String,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct NetworkConfig {
@ -54,7 +50,6 @@ pub struct NetworkConfig {
pub poll_timeout_microseconds: u64,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct ProtocolConfig {
@ -66,7 +61,6 @@ pub struct ProtocolConfig {
pub peer_announce_interval: usize,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct HandlerConfig {
@ -76,7 +70,6 @@ pub struct HandlerConfig {
pub channel_recv_timeout_microseconds: u64,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct CleaningConfig {
@ -88,7 +81,6 @@ pub struct CleaningConfig {
pub max_connection_age: u64,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct StatisticsConfig {
@ -96,7 +88,6 @@ pub struct StatisticsConfig {
pub interval: u64,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct PrivilegeConfig {
@ -108,7 +99,6 @@ pub struct PrivilegeConfig {
pub user: String,
}
impl Default for Config {
fn default() -> Self {
Self {
@ -125,7 +115,6 @@ impl Default for Config {
}
}
impl Default for NetworkConfig {
fn default() -> Self {
Self {
@ -139,7 +128,6 @@ impl Default for NetworkConfig {
}
}
impl Default for ProtocolConfig {
fn default() -> Self {
Self {
@ -150,7 +138,6 @@ impl Default for ProtocolConfig {
}
}
impl Default for HandlerConfig {
fn default() -> Self {
Self {
@ -160,7 +147,6 @@ impl Default for HandlerConfig {
}
}
impl Default for CleaningConfig {
fn default() -> Self {
Self {
@ -171,16 +157,12 @@ impl Default for CleaningConfig {
}
}
impl Default for StatisticsConfig {
fn default() -> Self {
Self {
interval: 0,
}
Self { interval: 0 }
}
}
impl Default for PrivilegeConfig {
fn default() -> Self {
Self {
@ -191,7 +173,6 @@ impl Default for PrivilegeConfig {
}
}
impl Default for TlsConfig {
fn default() -> Self {
Self {

123
aquatic_http/src/lib/handler.rs
View file

@ -1,13 +1,13 @@
use std::collections::BTreeMap;
use std::time::Duration;
use std::vec::Drain;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::sync::Arc;
use std::time::Duration;
use std::vec::Drain;
use either::Either;
use mio::Waker;
use parking_lot::MutexGuard;
use rand::{Rng, SeedableRng, rngs::SmallRng};
use rand::{rngs::SmallRng, Rng, SeedableRng};
use aquatic_common::extract_response_peers;
use aquatic_http_protocol::request::*;
@ -16,14 +16,13 @@ use aquatic_http_protocol::response::*;
use crate::common::*;
use crate::config::Config;
pub fn run_request_worker(
config: Config,
state: State,
request_channel_receiver: RequestChannelReceiver,
response_channel_sender: ResponseChannelSender,
wakers: Vec<Arc<Waker>>,
){
) {
let mut wake_socket_workers: Vec<bool> = (0..config.socket_workers).map(|_| false).collect();
let mut announce_requests = Vec::new();
@ -31,9 +30,7 @@ pub fn run_request_worker(
let mut rng = SmallRng::from_entropy();
let timeout = Duration::from_micros(
config.handlers.channel_recv_timeout_microseconds
);
let timeout = Duration::from_micros(config.handlers.channel_recv_timeout_microseconds);
loop {
let mut opt_torrent_map_guard: Option<MutexGuard<TorrentMaps>> = None;
@ -51,22 +48,22 @@ pub fn run_request_worker(
match opt_in_message {
Some((meta, Request::Announce(r))) => {
announce_requests.push((meta, r));
},
}
Some((meta, Request::Scrape(r))) => {
scrape_requests.push((meta, r));
},
}
None => {
if let Some(torrent_guard) = state.torrent_maps.try_lock(){
if let Some(torrent_guard) = state.torrent_maps.try_lock() {
opt_torrent_map_guard = Some(torrent_guard);
break
break;
}
}
}
}
let mut torrent_map_guard = opt_torrent_map_guard
.unwrap_or_else(|| state.torrent_maps.lock());
let mut torrent_map_guard =
opt_torrent_map_guard.unwrap_or_else(|| state.torrent_maps.lock());
handle_announce_requests(
&config,
@ -74,7 +71,7 @@ pub fn run_request_worker(
&mut torrent_map_guard,
&response_channel_sender,
&mut wake_socket_workers,
announce_requests.drain(..)
announce_requests.drain(..),
);
handle_scrape_requests(
@ -82,12 +79,12 @@ pub fn run_request_worker(
&mut torrent_map_guard,
&response_channel_sender,
&mut wake_socket_workers,
scrape_requests.drain(..)
scrape_requests.drain(..),
);
for (worker_index, wake) in wake_socket_workers.iter_mut().enumerate(){
for (worker_index, wake) in wake_socket_workers.iter_mut().enumerate() {
if *wake {
if let Err(err) = wakers[worker_index].wake(){
if let Err(err) = wakers[worker_index].wake() {
::log::error!("request handler couldn't wake poll: {:?}", err);
}
@ -97,7 +94,6 @@ pub fn run_request_worker(
}
}
pub fn handle_announce_requests(
config: &Config,
rng: &mut impl Rng,
@ -105,21 +101,18 @@ pub fn handle_announce_requests(
response_channel_sender: &ResponseChannelSender,
wake_socket_workers: &mut Vec<bool>,
requests: Drain<(ConnectionMeta, AnnounceRequest)>,
){
) {
let valid_until = ValidUntil::new(config.cleaning.max_peer_age);
for (meta, request) in requests {
let peer_ip = convert_ipv4_mapped_ipv6(
meta.peer_addr.ip()
);
let peer_ip = convert_ipv4_mapped_ipv6(meta.peer_addr.ip());
::log::debug!("peer ip: {:?}", peer_ip);
let response = match peer_ip {
IpAddr::V4(peer_ip_address) => {
let torrent_data: &mut TorrentData<Ipv4Addr> = torrent_maps.ipv4
.entry(request.info_hash)
.or_default();
let torrent_data: &mut TorrentData<Ipv4Addr> =
torrent_maps.ipv4.entry(request.info_hash).or_default();
let peer_connection_meta = PeerConnectionMeta {
worker_index: meta.worker_index,
@ -133,7 +126,7 @@ pub fn handle_announce_requests(
peer_connection_meta,
torrent_data,
request,
valid_until
valid_until,
);
let response = AnnounceResponse {
@ -145,16 +138,15 @@ pub fn handle_announce_requests(
};
Response::Announce(response)
},
}
IpAddr::V6(peer_ip_address) => {
let torrent_data: &mut TorrentData<Ipv6Addr> = torrent_maps.ipv6
.entry(request.info_hash)
.or_default();
let torrent_data: &mut TorrentData<Ipv6Addr> =
torrent_maps.ipv6.entry(request.info_hash).or_default();
let peer_connection_meta = PeerConnectionMeta {
worker_index: meta.worker_index,
poll_token: meta.poll_token,
peer_ip_address
peer_ip_address,
};
let (seeders, leechers, response_peers) = upsert_peer_and_get_response_peers(
@ -163,7 +155,7 @@ pub fn handle_announce_requests(
peer_connection_meta,
torrent_data,
request,
valid_until
valid_until,
);
let response = AnnounceResponse {
@ -175,15 +167,14 @@ pub fn handle_announce_requests(
};
Response::Announce(response)
},
}
};
response_channel_sender.send(meta, response);
wake_socket_workers[meta.worker_index] = true;
};
}
}
/// Insert/update peer. Return num_seeders, num_leechers and response peers
fn upsert_peer_and_get_response_peers<I: Ip>(
config: &Config,
@ -195,10 +186,8 @@ fn upsert_peer_and_get_response_peers<I: Ip>(
) -> (usize, usize, Vec<ResponsePeer<I>>) {
// Insert/update/remove peer who sent this request
let peer_status = PeerStatus::from_event_and_bytes_left(
request.event,
Some(request.bytes_left)
);
let peer_status =
PeerStatus::from_event_and_bytes_left(request.event, Some(request.bytes_left));
let peer = Peer {
connection_meta: request_sender_meta,
@ -209,11 +198,10 @@ fn upsert_peer_and_get_response_peers<I: Ip>(
::log::debug!("peer: {:?}", peer);
let ip_or_key = request.key
let ip_or_key = request
.key
.map(Either::Right)
.unwrap_or_else(||
Either::Left(request_sender_meta.peer_ip_address)
);
.unwrap_or_else(|| Either::Left(request_sender_meta.peer_ip_address));
let peer_map_key = PeerMapKey {
peer_id: request.peer_id,
@ -227,26 +215,24 @@ fn upsert_peer_and_get_response_peers<I: Ip>(
torrent_data.num_leechers += 1;
torrent_data.peers.insert(peer_map_key.clone(), peer)
},
}
PeerStatus::Seeding => {
torrent_data.num_seeders += 1;
torrent_data.peers.insert(peer_map_key.clone(), peer)
},
PeerStatus::Stopped => {
torrent_data.peers.remove(&peer_map_key)
}
PeerStatus::Stopped => torrent_data.peers.remove(&peer_map_key),
};
::log::debug!("opt_removed_peer: {:?}", opt_removed_peer);
match opt_removed_peer.map(|peer| peer.status){
match opt_removed_peer.map(|peer| peer.status) {
Some(PeerStatus::Leeching) => {
torrent_data.num_leechers -= 1;
},
}
Some(PeerStatus::Seeding) => {
torrent_data.num_seeders -= 1;
},
}
_ => {}
}
@ -262,38 +248,40 @@ fn upsert_peer_and_get_response_peers<I: Ip>(
&torrent_data.peers,
max_num_peers_to_take,
peer_map_key,
Peer::to_response_peer
Peer::to_response_peer,
);
(torrent_data.num_seeders, torrent_data.num_leechers, response_peers)
(
torrent_data.num_seeders,
torrent_data.num_leechers,
response_peers,
)
}
pub fn handle_scrape_requests(
config: &Config,
torrent_maps: &mut TorrentMaps,
response_channel_sender: &ResponseChannelSender,
wake_socket_workers: &mut Vec<bool>,
requests: Drain<(ConnectionMeta, ScrapeRequest)>,
){
) {
for (meta, request) in requests {
let num_to_take = request.info_hashes.len().min(
config.protocol.max_scrape_torrents
);
let num_to_take = request
.info_hashes
.len()
.min(config.protocol.max_scrape_torrents);
let mut response = ScrapeResponse {
files: BTreeMap::new(),
};
let peer_ip = convert_ipv4_mapped_ipv6(
meta.peer_addr.ip()
);
let peer_ip = convert_ipv4_mapped_ipv6(meta.peer_addr.ip());
// If request.info_hashes is empty, don't return scrape for all
// torrents, even though reference server does it. It is too expensive.
if peer_ip.is_ipv4(){
for info_hash in request.info_hashes.into_iter().take(num_to_take){
if let Some(torrent_data) = torrent_maps.ipv4.get(&info_hash){
if peer_ip.is_ipv4() {
for info_hash in request.info_hashes.into_iter().take(num_to_take) {
if let Some(torrent_data) = torrent_maps.ipv4.get(&info_hash) {
let stats = ScrapeStatistics {
complete: torrent_data.num_seeders,
downloaded: 0, // No implementation planned
@ -304,8 +292,8 @@ pub fn handle_scrape_requests(
}
}
} else {
for info_hash in request.info_hashes.into_iter().take(num_to_take){
if let Some(torrent_data) = torrent_maps.ipv6.get(&info_hash){
for info_hash in request.info_hashes.into_iter().take(num_to_take) {
if let Some(torrent_data) = torrent_maps.ipv6.get(&info_hash) {
let stats = ScrapeStatistics {
complete: torrent_data.num_seeders,
downloaded: 0, // No implementation planned
@ -317,8 +305,7 @@ pub fn handle_scrape_requests(
}
};
response_channel_sender.send(meta, Response::Scrape(response));
wake_socket_workers[meta.worker_index] = true;
};
}
}

73
aquatic_http/src/lib/lib.rs
View file

@ -1,6 +1,6 @@
use std::time::Duration;
use std::sync::Arc;
use std::thread::Builder;
use std::time::Duration;
use anyhow::Context;
use mio::{Poll, Waker};
@ -17,10 +17,8 @@ use common::*;
use config::Config;
use network::utils::create_tls_acceptor;
pub const APP_NAME: &str = "aquatic_http: HTTP/TLS BitTorrent tracker";
pub fn run(config: Config) -> anyhow::Result<()> {
let state = State::default();
@ -33,7 +31,6 @@ pub fn run(config: Config) -> anyhow::Result<()> {
}
}
pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
let opt_tls_acceptor = create_tls_acceptor(&config.network.tls)?;
@ -65,17 +62,19 @@ pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
out_message_senders.push(response_channel_sender);
wakers.push(waker);
Builder::new().name(format!("socket-{:02}", i + 1)).spawn(move || {
network::run_socket_worker(
config,
i,
socket_worker_statuses,
request_channel_sender,
response_channel_receiver,
opt_tls_acceptor,
poll
);
})?;
Builder::new()
.name(format!("socket-{:02}", i + 1))
.spawn(move || {
network::run_socket_worker(
config,
i,
socket_worker_statuses,
request_channel_sender,
response_channel_receiver,
opt_tls_acceptor,
poll,
);
})?;
}
// Wait for socket worker statuses. On error from any, quit program.
@ -84,14 +83,14 @@ pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
loop {
::std::thread::sleep(::std::time::Duration::from_millis(10));
if let Some(statuses) = socket_worker_statuses.try_lock(){
for opt_status in statuses.iter(){
if let Some(Err(err)) = opt_status {
if let Some(statuses) = socket_worker_statuses.try_lock() {
for opt_status in statuses.iter() {
if let Some(Err(err)) = opt_status {
return Err(::anyhow::anyhow!(err.to_owned()));
}
}
if statuses.iter().all(Option::is_some){
if statuses.iter().all(Option::is_some) {
if config.privileges.drop_privileges {
PrivDrop::default()
.chroot(config.privileges.chroot_path.clone())
@ -100,7 +99,7 @@ pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
.context("Couldn't drop root privileges")?;
}
break
break;
}
}
}
@ -114,32 +113,32 @@ pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
let response_channel_sender = response_channel_sender.clone();
let wakers = wakers.clone();
Builder::new().name(format!("request-{:02}", i + 1)).spawn(move || {
handler::run_request_worker(
config,
state,
request_channel_receiver,
response_channel_sender,
wakers,
);
})?;
Builder::new()
.name(format!("request-{:02}", i + 1))
.spawn(move || {
handler::run_request_worker(
config,
state,
request_channel_receiver,
response_channel_sender,
wakers,
);
})?;
}
if config.statistics.interval != 0 {
let state = state.clone();
let config = config.clone();
Builder::new().name("statistics".to_string()).spawn(move ||
loop {
::std::thread::sleep(Duration::from_secs(
config.statistics.interval
));
Builder::new()
.name("statistics".to_string())
.spawn(move || loop {
::std::thread::sleep(Duration::from_secs(config.statistics.interval));
tasks::print_statistics(&state);
}
).expect("spawn statistics thread");
})
.expect("spawn statistics thread");
}
Ok(())
}

134
aquatic_http/src/lib/network/connection.rs
View file

@ -1,12 +1,12 @@
use std::net::{SocketAddr};
use std::io::ErrorKind;
use std::io::{Read, Write};
use std::net::SocketAddr;
use std::sync::Arc;
use hashbrown::HashMap;
use mio::{Token, Poll};
use mio::net::TcpStream;
use native_tls::{TlsAcceptor, MidHandshakeTlsStream};
use mio::{Poll, Token};
use native_tls::{MidHandshakeTlsStream, TlsAcceptor};
use aquatic_http_protocol::request::{Request, RequestParseError};
@ -14,7 +14,6 @@ use crate::common::*;
use super::stream::Stream;
#[derive(Debug)]
pub enum RequestReadError {
NeedMoreData,
@ -23,7 +22,6 @@ pub enum RequestReadError {
Io(::std::io::Error),
}
pub struct EstablishedConnection {
stream: Stream,
pub peer_addr: SocketAddr,
@ -31,7 +29,6 @@ pub struct EstablishedConnection {
bytes_read: usize,
}
impl EstablishedConnection {
#[inline]
fn new(stream: Stream) -> Self {
@ -46,11 +43,11 @@ impl EstablishedConnection {
}
pub fn read_request(&mut self) -> Result<Request, RequestReadError> {
if (self.buf.len() - self.bytes_read < 512) & (self.buf.len() <= 3072){
if (self.buf.len() - self.bytes_read < 512) & (self.buf.len() <= 3072) {
self.buf.extend_from_slice(&[0; 1024]);
}
match self.stream.read(&mut self.buf[self.bytes_read..]){
match self.stream.read(&mut self.buf[self.bytes_read..]) {
Ok(0) => {
self.clear_buffer();
@ -60,10 +57,10 @@ impl EstablishedConnection {
self.bytes_read += bytes_read;
::log::debug!("read_request read {} bytes", bytes_read);
},
}
Err(err) if err.kind() == ErrorKind::WouldBlock => {
return Err(RequestReadError::NeedMoreData);
},
}
Err(err) => {
self.clear_buffer();
@ -71,20 +68,18 @@ impl EstablishedConnection {
}
}
match Request::from_bytes(&self.buf[..self.bytes_read]){
match Request::from_bytes(&self.buf[..self.bytes_read]) {
Ok(request) => {
self.clear_buffer();
Ok(request)
},
Err(RequestParseError::NeedMoreData) => {
Err(RequestReadError::NeedMoreData)
},
}
Err(RequestParseError::NeedMoreData) => Err(RequestReadError::NeedMoreData),
Err(RequestParseError::Invalid(err)) => {
self.clear_buffer();
Err(RequestReadError::Parse(err))
},
}
}
}
@ -92,9 +87,7 @@ impl EstablishedConnection {
let content_len = body.len() + 2; // 2 is for newlines at end
let content_len_num_digits = Self::num_digits_in_usize(content_len);
let mut response = Vec::with_capacity(
39 + content_len_num_digits + body.len()
);
let mut response = Vec::with_capacity(39 + content_len_num_digits + body.len());
response.extend_from_slice(b"HTTP/1.1 200 OK\r\nContent-Length: ");
::itoa::write(&mut response, content_len)?;
@ -130,40 +123,33 @@ impl EstablishedConnection {
}
#[inline]
pub fn clear_buffer(&mut self){
pub fn clear_buffer(&mut self) {
self.bytes_read = 0;
self.buf = Vec::new();
}
}
pub enum TlsHandshakeMachineError {
WouldBlock(TlsHandshakeMachine),
Failure(native_tls::Error)
Failure(native_tls::Error),
}
enum TlsHandshakeMachineInner {
TcpStream(TcpStream),
TlsMidHandshake(MidHandshakeTlsStream<TcpStream>),
}
pub struct TlsHandshakeMachine {
tls_acceptor: Arc<TlsAcceptor>,
inner: TlsHandshakeMachineInner,
}
impl <'a>TlsHandshakeMachine {
impl<'a> TlsHandshakeMachine {
#[inline]
fn new(
tls_acceptor: Arc<TlsAcceptor>,
tcp_stream: TcpStream
) -> Self {
fn new(tls_acceptor: Arc<TlsAcceptor>, tcp_stream: TcpStream) -> Self {
Self {
tls_acceptor,
inner: TlsHandshakeMachineInner::TcpStream(tcp_stream)
inner: TlsHandshakeMachineInner::TcpStream(tcp_stream),
}
}
@ -171,28 +157,20 @@ impl <'a>TlsHandshakeMachine {
/// the machine wrapped in an error for later attempts.
pub fn establish_tls(self) -> Result<EstablishedConnection, TlsHandshakeMachineError> {
let handshake_result = match self.inner {
TlsHandshakeMachineInner::TcpStream(stream) => {
self.tls_acceptor.accept(stream)
},
TlsHandshakeMachineInner::TlsMidHandshake(handshake) => {
handshake.handshake()
},
TlsHandshakeMachineInner::TcpStream(stream) => self.tls_acceptor.accept(stream),
TlsHandshakeMachineInner::TlsMidHandshake(handshake) => handshake.handshake(),
};
match handshake_result {
Ok(stream) => {
let established = EstablishedConnection::new(
Stream::TlsStream(stream)
);
let established = EstablishedConnection::new(Stream::TlsStream(stream));
::log::debug!("established tls connection");
Ok(established)
},
}
Err(native_tls::HandshakeError::WouldBlock(handshake)) => {
let inner = TlsHandshakeMachineInner::TlsMidHandshake(
handshake
);
let inner = TlsHandshakeMachineInner::TlsMidHandshake(handshake);
let machine = Self {
tls_acceptor: self.tls_acceptor,
@ -200,7 +178,7 @@ impl <'a>TlsHandshakeMachine {
};
Err(TlsHandshakeMachineError::WouldBlock(machine))
},
}
Err(native_tls::HandshakeError::Failure(err)) => {
Err(TlsHandshakeMachineError::Failure(err))
}
@ -208,19 +186,16 @@ impl <'a>TlsHandshakeMachine {
}
}
enum ConnectionInner {
Established(EstablishedConnection),
InProgress(TlsHandshakeMachine),
}
pub struct Connection {
pub valid_until: ValidUntil,
inner: ConnectionInner,
}
impl Connection {
#[inline]
pub fn new(
@ -230,42 +205,29 @@ impl Connection {
) -> Self {
// Setup handshake machine if TLS is requested
let inner = if let Some(tls_acceptor) = opt_tls_acceptor {
ConnectionInner::InProgress(
TlsHandshakeMachine::new(tls_acceptor.clone(), tcp_stream)
)
ConnectionInner::InProgress(TlsHandshakeMachine::new(tls_acceptor.clone(), tcp_stream))
} else {
::log::debug!("established tcp connection");
ConnectionInner::Established(
EstablishedConnection::new(Stream::TcpStream(tcp_stream))
)
ConnectionInner::Established(EstablishedConnection::new(Stream::TcpStream(tcp_stream)))
};
Self { valid_until, inner }
}
#[inline]
pub fn from_established(valid_until: ValidUntil, established: EstablishedConnection) -> Self {
Self {
valid_until,
inner,
inner: ConnectionInner::Established(established),
}
}
#[inline]
pub fn from_established(
valid_until: ValidUntil,
established: EstablishedConnection,
) -> Self {
pub fn from_in_progress(valid_until: ValidUntil, machine: TlsHandshakeMachine) -> Self {
Self {
valid_until,
inner: ConnectionInner::Established(established)
}
}
#[inline]
pub fn from_in_progress(
valid_until: ValidUntil,
machine: TlsHandshakeMachine,
) -> Self {
Self {
valid_until,
inner: ConnectionInner::InProgress(machine)
inner: ConnectionInner::InProgress(machine),
}
}
@ -290,40 +252,30 @@ impl Connection {
pub fn deregister(&mut self, poll: &mut Poll) -> ::std::io::Result<()> {
match &mut self.inner {
ConnectionInner::Established(established) => {
match &mut established.stream {
Stream::TcpStream(ref mut stream) => {
poll.registry().deregister(stream)
},
Stream::TlsStream(ref mut stream) => {
poll.registry().deregister(stream.get_mut())
},
}
ConnectionInner::Established(established) => match &mut established.stream {
Stream::TcpStream(ref mut stream) => poll.registry().deregister(stream),
Stream::TlsStream(ref mut stream) => poll.registry().deregister(stream.get_mut()),
},
ConnectionInner::InProgress(TlsHandshakeMachine { inner, ..}) => {
match inner {
TlsHandshakeMachineInner::TcpStream(ref mut stream) => {
poll.registry().deregister(stream)
},
TlsHandshakeMachineInner::TlsMidHandshake(ref mut mid_handshake) => {
poll.registry().deregister(mid_handshake.get_mut())
},
ConnectionInner::InProgress(TlsHandshakeMachine { inner, .. }) => match inner {
TlsHandshakeMachineInner::TcpStream(ref mut stream) => {
poll.registry().deregister(stream)
}
TlsHandshakeMachineInner::TlsMidHandshake(ref mut mid_handshake) => {
poll.registry().deregister(mid_handshake.get_mut())
}
},
}
}
}
pub type ConnectionMap = HashMap<Token, Connection>;
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_num_digits_in_usize(){
fn test_num_digits_in_usize() {
let f = EstablishedConnection::num_digits_in_usize;
assert_eq!(f(0), 1);

154
aquatic_http/src/lib/network/mod.rs
View file

@ -1,13 +1,13 @@
use std::time::{Duration, Instant};
use std::io::{ErrorKind, Cursor};
use std::io::{Cursor, ErrorKind};
use std::sync::Arc;
use std::time::{Duration, Instant};
use std::vec::Drain;
use hashbrown::HashMap;
use log::{info, debug, error};
use native_tls::TlsAcceptor;
use mio::{Events, Poll, Interest, Token};
use log::{debug, error, info};
use mio::net::TcpListener;
use mio::{Events, Interest, Poll, Token};
use native_tls::TlsAcceptor;
use aquatic_http_protocol::response::*;
@ -21,10 +21,8 @@ pub mod utils;
use connection::*;
use utils::*;
const CONNECTION_CLEAN_INTERVAL: usize = 2 ^ 22;
pub fn run_socket_worker(
config: Config,
socket_worker_index: usize,
@ -33,8 +31,8 @@ pub fn run_socket_worker(
response_channel_receiver: ResponseChannelReceiver,
opt_tls_acceptor: Option<TlsAcceptor>,
poll: Poll,
){
match create_listener(config.network.address, config.network.ipv6_only){
) {
match create_listener(config.network.address, config.network.ipv6_only) {
Ok(listener) => {
socket_worker_statuses.lock()[socket_worker_index] = Some(Ok(()));
@ -47,16 +45,14 @@ pub fn run_socket_worker(
opt_tls_acceptor,
poll,
);
},
}
Err(err) => {
socket_worker_statuses.lock()[socket_worker_index] = Some(
Err(format!("Couldn't open socket: {:#}", err))
);
socket_worker_statuses.lock()[socket_worker_index] =
Some(Err(format!("Couldn't open socket: {:#}", err)));
}
}
}
pub fn run_poll_loop(
config: Config,
socket_worker_index: usize,
@ -65,10 +61,8 @@ pub fn run_poll_loop(
listener: ::std::net::TcpListener,
opt_tls_acceptor: Option<TlsAcceptor>,
mut poll: Poll,
){
let poll_timeout = Duration::from_micros(
config.network.poll_timeout_microseconds
);
) {
let poll_timeout = Duration::from_micros(config.network.poll_timeout_microseconds);
let mut listener = TcpListener::from_std(listener);
let mut events = Events::with_capacity(config.network.poll_event_capacity);
@ -91,7 +85,7 @@ pub fn run_poll_loop(
poll.poll(&mut events, Some(poll_timeout))
.expect("failed polling");
for event in events.iter(){
for event in events.iter() {
let token = event.token();
if token == LISTENER_TOKEN {
@ -124,7 +118,7 @@ pub fn run_poll_loop(
&mut response_buffer,
local_responses.drain(..),
&response_channel_receiver,
&mut connections
&mut connections,
);
}
@ -137,7 +131,6 @@ pub fn run_poll_loop(
}
}
fn accept_new_streams(
config: &Config,
listener: &mut TcpListener,
@ -145,11 +138,11 @@ fn accept_new_streams(
connections: &mut ConnectionMap,
poll_token_counter: &mut Token,
opt_tls_acceptor: &Option<Arc<TlsAcceptor>>,
){
) {
let valid_until = ValidUntil::new(config.cleaning.max_connection_age);
loop {
match listener.accept(){
match listener.accept() {
Ok((mut stream, _)) => {
poll_token_counter.0 = poll_token_counter.0.wrapping_add(1);
@ -167,17 +160,13 @@ fn accept_new_streams(
.register(&mut stream, token, Interest::READABLE)
.unwrap();
let connection = Connection::new(
opt_tls_acceptor,
valid_until,
stream
);
let connection = Connection::new(opt_tls_acceptor, valid_until, stream);
connections.insert(token, connection);
},
}
Err(err) => {
if err.kind() == ErrorKind::WouldBlock {
break
break;
}
info!("error while accepting streams: {}", err);
@ -186,7 +175,6 @@ fn accept_new_streams(
}
}
/// On the stream given by poll_token, get TLS up and running if requested,
/// then read requests and pass on through channel.
pub fn handle_connection_read_event(
@ -197,119 +185,106 @@ pub fn handle_connection_read_event(
local_responses: &mut Vec<(ConnectionMeta, Response)>,
connections: &mut ConnectionMap,
poll_token: Token,
){
) {
let valid_until = ValidUntil::new(config.cleaning.max_connection_age);
loop {
// Get connection, updating valid_until
let connection = if let Some(c) = connections.get_mut(&poll_token){
let connection = if let Some(c) = connections.get_mut(&poll_token) {
c
} else {
// If there is no connection, there is no stream, so there
// shouldn't be any (relevant) poll events. In other words, it's
// safe to return here
return
return;
};
connection.valid_until = valid_until;
if let Some(established) = connection.get_established(){
match established.read_request(){
if let Some(established) = connection.get_established() {
match established.read_request() {
Ok(request) => {
let meta = ConnectionMeta {
worker_index: socket_worker_index,
poll_token,
peer_addr: established.peer_addr
peer_addr: established.peer_addr,
};
debug!("read request, sending to handler");
if let Err(err) = request_channel_sender
.send((meta, request))
{
error!(
"RequestChannelSender: couldn't send message: {:?}",
err
);
if let Err(err) = request_channel_sender.send((meta, request)) {
error!("RequestChannelSender: couldn't send message: {:?}", err);
}
break
},
break;
}
Err(RequestReadError::NeedMoreData) => {
info!("need more data");
// Stop reading data (defer to later events)
break;
},
}
Err(RequestReadError::Parse(err)) => {
info!("error reading request (invalid): {:#?}", err);
let meta = ConnectionMeta {
worker_index: socket_worker_index,
poll_token,
peer_addr: established.peer_addr
peer_addr: established.peer_addr,
};
let response = FailureResponse {
failure_reason: "invalid request".to_string()
failure_reason: "invalid request".to_string(),
};
local_responses.push(
(meta, Response::Failure(response))
);
local_responses.push((meta, Response::Failure(response)));
break;
},
}
Err(RequestReadError::StreamEnded) => {
::log::debug!("stream ended");
remove_connection(poll, connections, &poll_token);
break
},
break;
}
Err(RequestReadError::Io(err)) => {
::log::info!("error reading request (io): {}", err);
remove_connection(poll, connections, &poll_token);
break;
},
}
}
} else if let Some(handshake_machine) = connections.remove(&poll_token)
} else if let Some(handshake_machine) = connections
.remove(&poll_token)
.and_then(Connection::get_in_progress)
{
match handshake_machine.establish_tls(){
match handshake_machine.establish_tls() {
Ok(established) => {
let connection = Connection::from_established(
valid_until,
established
);
let connection = Connection::from_established(valid_until, established);
connections.insert(poll_token, connection);
},
}
Err(TlsHandshakeMachineError::WouldBlock(machine)) => {
let connection = Connection::from_in_progress(
valid_until,
machine
);
let connection = Connection::from_in_progress(valid_until, machine);
connections.insert(poll_token, connection);
// Break and wait for more data
break
},
break;
}
Err(TlsHandshakeMachineError::Failure(err)) => {
info!("tls handshake error: {}", err);
// TLS negotiation failed
break
break;
}
}
}
}
}
/// Read responses from channel, send to peers
pub fn send_responses(
config: &Config,
@ -318,13 +293,13 @@ pub fn send_responses(
local_responses: Drain<(ConnectionMeta, Response)>,
channel_responses: &ResponseChannelReceiver,
connections: &mut ConnectionMap,
){
) {
let channel_responses_len = channel_responses.len();
let channel_responses_drain = channel_responses.try_iter()
.take(channel_responses_len);
let channel_responses_drain = channel_responses.try_iter().take(channel_responses_len);
for (meta, response) in local_responses.chain(channel_responses_drain){
if let Some(established) = connections.get_mut(&meta.poll_token)
for (meta, response) in local_responses.chain(channel_responses_drain) {
if let Some(established) = connections
.get_mut(&meta.poll_token)
.and_then(Connection::get_established)
{
if established.peer_addr != meta.peer_addr {
@ -337,7 +312,7 @@ pub fn send_responses(
let bytes_written = response.write(buffer).unwrap();
match established.send_response(&buffer.get_mut()[..bytes_written]){
match established.send_response(&buffer.get_mut()[..bytes_written]) {
Ok(()) => {
::log::debug!(
"sent response: {:?} with response string {}",
@ -348,33 +323,29 @@ pub fn send_responses(
if !config.network.keep_alive {
remove_connection(poll, connections, &meta.poll_token);
}
},
}
Err(err) if err.kind() == ErrorKind::WouldBlock => {
debug!("send response: would block");
},
}
Err(err) => {
info!("error sending response: {}", err);
remove_connection(poll, connections, &meta.poll_token);
},
}
}
}
}
}
// Close and remove inactive connections
pub fn remove_inactive_connections(
poll: &mut Poll,
connections: &mut ConnectionMap,
){
pub fn remove_inactive_connections(poll: &mut Poll, connections: &mut ConnectionMap) {
let now = Instant::now();
connections.retain(|_, connection| {
let keep = connection.valid_until.0 >= now;
if !keep {
if let Err(err) = connection.deregister(poll){
if let Err(err) = connection.deregister(poll) {
::log::error!("deregister connection error: {}", err);
}
}
@ -385,14 +356,9 @@ pub fn remove_inactive_connections(
connections.shrink_to_fit();
}
fn remove_connection(
poll: &mut Poll,
connections: &mut ConnectionMap,
connection_token: &Token,
){
if let Some(mut connection) = connections.remove(connection_token){
if let Err(err) = connection.deregister(poll){
fn remove_connection(poll: &mut Poll, connections: &mut ConnectionMap, connection_token: &Token) {
if let Some(mut connection) = connections.remove(connection_token) {
if let Err(err) = connection.deregister(poll) {
::log::error!("deregister connection error: {}", err);
}
}

13
aquatic_http/src/lib/network/stream.rs
View file

@ -1,16 +1,14 @@
use std::net::{SocketAddr};
use std::io::{Read, Write};
use std::net::SocketAddr;
use mio::net::TcpStream;
use native_tls::TlsStream;
pub enum Stream {
TcpStream(TcpStream),
TlsStream(TlsStream<TcpStream>),
}
impl Stream {
#[inline]
pub fn get_peer_addr(&self) -> SocketAddr {
@ -21,7 +19,6 @@ impl Stream {
}
}
impl Read for Stream {
#[inline]
fn read(&mut self, buf: &mut [u8]) -> Result<usize, ::std::io::Error> {
@ -35,7 +32,7 @@ impl Read for Stream {
#[inline]
fn read_vectored(
&mut self,
bufs: &mut [::std::io::IoSliceMut<'_>]
bufs: &mut [::std::io::IoSliceMut<'_>],
) -> ::std::io::Result<usize> {
match self {
Self::TcpStream(stream) => stream.read_vectored(bufs),
@ -44,7 +41,6 @@ impl Read for Stream {
}
}
impl Write for Stream {
#[inline]
fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
@ -56,10 +52,7 @@ impl Write for Stream {
/// Not used but provided for completeness
#[inline]
fn write_vectored(
&mut self,
bufs: &[::std::io::IoSlice<'_>]
) -> ::std::io::Result<usize> {
fn write_vectored(&mut self, bufs: &[::std::io::IoSlice<'_>]) -> ::std::io::Result<usize> {
match self {
Self::TcpStream(stream) => stream.write_vectored(bufs),
Self::TlsStream(stream) => stream.write_vectored(bufs),

43
aquatic_http/src/lib/network/utils.rs
View file

@ -4,26 +4,21 @@ use std::net::SocketAddr;
use anyhow::Context;
use native_tls::{Identity, TlsAcceptor};
use socket2::{Socket, Domain, Type, Protocol};
use socket2::{Domain, Protocol, Socket, Type};
use crate::config::TlsConfig;
pub fn create_tls_acceptor(
config: &TlsConfig,
) -> anyhow::Result<Option<TlsAcceptor>> {
pub fn create_tls_acceptor(config: &TlsConfig) -> anyhow::Result<Option<TlsAcceptor>> {
if config.use_tls {
let mut identity_bytes = Vec::new();
let mut file = File::open(&config.tls_pkcs12_path)
.context("Couldn't open pkcs12 identity file")?;
let mut file =
File::open(&config.tls_pkcs12_path).context("Couldn't open pkcs12 identity file")?;
file.read_to_end(&mut identity_bytes)
.context("Couldn't read pkcs12 identity file")?;
let identity = Identity::from_pkcs12(
&identity_bytes[..],
&config.tls_pkcs12_password
).context("Couldn't parse pkcs12 identity file")?;
let identity = Identity::from_pkcs12(&identity_bytes[..], &config.tls_pkcs12_password)
.context("Couldn't parse pkcs12 identity file")?;
let acceptor = TlsAcceptor::new(identity)
.context("Couldn't create TlsAcceptor from pkcs12 identity")?;
@ -34,30 +29,34 @@ pub fn create_tls_acceptor(
}
}
pub fn create_listener(
address: SocketAddr,
ipv6_only: bool
ipv6_only: bool,
) -> ::anyhow::Result<::std::net::TcpListener> {
let builder = if address.is_ipv4(){
let builder = if address.is_ipv4() {
Socket::new(Domain::ipv4(), Type::stream(), Some(Protocol::tcp()))
} else {
Socket::new(Domain::ipv6(), Type::stream(), Some(Protocol::tcp()))
}.context("Couldn't create socket2::Socket")?;
}
.context("Couldn't create socket2::Socket")?;
if ipv6_only {
builder.set_only_v6(true)
builder
.set_only_v6(true)
.context("Couldn't put socket in ipv6 only mode")?
}
builder.set_nonblocking(true)
builder
.set_nonblocking(true)
.context("Couldn't put socket in non-blocking mode")?;
builder.set_reuse_port(true)
builder
.set_reuse_port(true)
.context("Couldn't put socket in reuse_port mode")?;
builder.bind(&address.into()).with_context(||
format!("Couldn't bind socket to address {}", address)
)?;
builder.listen(128)
builder
.bind(&address.into())
.with_context(|| format!("Couldn't bind socket to address {}", address))?;
builder
.listen(128)
.context("Couldn't listen for connections on socket")?;
Ok(builder.into_tcp_listener())

24
aquatic_http/src/lib/tasks.rs
View file

@ -4,19 +4,14 @@ use histogram::Histogram;
use crate::common::*;
pub fn clean_torrents(state: &State){
pub fn clean_torrents(state: &State) {
let mut torrent_maps = state.torrent_maps.lock();
clean_torrent_map(&mut torrent_maps.ipv4);
clean_torrent_map(&mut torrent_maps.ipv6);
}
fn clean_torrent_map<I: Ip>(
torrent_map: &mut TorrentMap<I>,
){
fn clean_torrent_map<I: Ip>(torrent_map: &mut TorrentMap<I>) {
let now = Instant::now();
torrent_map.retain(|_, torrent_data| {
@ -30,10 +25,10 @@ fn clean_torrent_map<I: Ip>(
match peer.status {
PeerStatus::Seeding => {
*num_seeders -= 1;
},
}
PeerStatus::Leeching => {
*num_leechers -= 1;
},
}
_ => (),
};
}
@ -47,24 +42,23 @@ fn clean_torrent_map<I: Ip>(
torrent_map.shrink_to_fit();
}
pub fn print_statistics(state: &State){
pub fn print_statistics(state: &State) {
let mut peers_per_torrent = Histogram::new();
{
let torrents = &mut state.torrent_maps.lock();
for torrent in torrents.ipv4.values(){
for torrent in torrents.ipv4.values() {
let num_peers = (torrent.num_seeders + torrent.num_leechers) as u64;
if let Err(err) = peers_per_torrent.increment(num_peers){
if let Err(err) = peers_per_torrent.increment(num_peers) {
eprintln!("error incrementing peers_per_torrent histogram: {}", err)
}
}
for torrent in torrents.ipv6.values(){
for torrent in torrents.ipv6.values() {
let num_peers = (torrent.num_seeders + torrent.num_leechers) as u64;
if let Err(err) = peers_per_torrent.increment(num_peers){
if let Err(err) = peers_per_torrent.increment(num_peers) {
eprintln!("error incrementing peers_per_torrent histogram: {}", err)
}
}

10
aquatic_http_load_test/src/common.rs
View file

@ -1,11 +1,10 @@
use std::sync::{Arc, atomic::AtomicUsize};
use std::sync::{atomic::AtomicUsize, Arc};
use rand_distr::Pareto;
pub use aquatic_http_protocol::common::*;
pub use aquatic_http_protocol::response::*;
pub use aquatic_http_protocol::request::*;
pub use aquatic_http_protocol::response::*;
#[derive(PartialEq, Eq, Clone)]
pub struct TorrentPeer {
@ -15,7 +14,6 @@ pub struct TorrentPeer {
pub port: u16,
}
#[derive(Default)]
pub struct Statistics {
pub requests: AtomicUsize,
@ -27,7 +25,6 @@ pub struct Statistics {
pub bytes_received: AtomicUsize,
}
#[derive(Clone)]
pub struct LoadTestState {
pub info_hashes: Arc<Vec<InfoHash>>,
@ -35,9 +32,8 @@ pub struct LoadTestState {
pub pareto: Arc<Pareto<f64>>,
}
#[derive(PartialEq, Eq, Clone, Copy)]
pub enum RequestType {
Announce,
Scrape
Scrape,
}

8
aquatic_http_load_test/src/config.rs
View file

@ -1,7 +1,6 @@
use std::net::SocketAddr;
use serde::{Serialize, Deserialize};
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
@ -14,10 +13,8 @@ pub struct Config {
pub torrents: TorrentConfig,
}
impl aquatic_cli_helpers::Config for Config {}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct NetworkConfig {
@ -26,7 +23,6 @@ pub struct NetworkConfig {
pub poll_event_capacity: usize,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct TorrentConfig {
@ -45,7 +41,6 @@ pub struct TorrentConfig {
pub weight_scrape: usize,
}
impl Default for Config {
fn default() -> Self {
Self {
@ -69,7 +64,6 @@ impl Default for NetworkConfig {
}
}
impl Default for TorrentConfig {
fn default() -> Self {
Self {

88
aquatic_http_load_test/src/main.rs
View file

@ -1,5 +1,5 @@
use std::sync::{atomic::Ordering, Arc};
use std::thread;
use std::sync::{Arc, atomic::Ordering};
use std::time::{Duration, Instant};
use rand::prelude::*;
@ -14,24 +14,20 @@ use common::*;
use config::*;
use network::*;
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
/// Multiply bytes during a second with this to get Mbit/s
const MBITS_FACTOR: f64 = 1.0 / ((1024.0 * 1024.0) / 8.0);
pub fn main(){
pub fn main() {
aquatic_cli_helpers::run_app_with_cli_and_config::<Config>(
"aquatic_http_load_test: BitTorrent load tester",
run,
None
None,
)
}
fn run(config: Config) -> ::anyhow::Result<()> {
if config.torrents.weight_announce + config.torrents.weight_scrape == 0 {
panic!("Error: at least one weight must be larger than zero.");
@ -47,10 +43,7 @@ fn run(config: Config) -> ::anyhow::Result<()> {
info_hashes.push(InfoHash(rng.gen()));
}
let pareto = Pareto::new(
1.0,
config.torrents.torrent_selection_pareto_shape
).unwrap();
let pareto = Pareto::new(1.0, config.torrents.torrent_selection_pareto_shape).unwrap();
let state = LoadTestState {
info_hashes: Arc::new(info_hashes),
@ -61,30 +54,18 @@ fn run(config: Config) -> ::anyhow::Result<()> {
// Start socket workers
for _ in 0..config.num_workers {
let config = config.clone();
let state = state.clone();
thread::spawn(move || run_socket_thread(
&config,
state,
1
));
thread::spawn(move || run_socket_thread(&config, state, 1));
}
monitor_statistics(
state,
&config
);
monitor_statistics(state, &config);
Ok(())
}
fn monitor_statistics(
state: LoadTestState,
config: &Config,
){
fn monitor_statistics(state: LoadTestState, config: &Config) {
let start_time = Instant::now();
let mut report_avg_response_vec: Vec<f64> = Vec::new();
@ -96,41 +77,52 @@ fn monitor_statistics(
let statistics = state.statistics.as_ref();
let responses_announce = statistics.responses_announce
.fetch_and(0, Ordering::SeqCst) as f64;
let responses_announce =
statistics.responses_announce.fetch_and(0, Ordering::SeqCst) as f64;
// let response_peers = statistics.response_peers
// .fetch_and(0, Ordering::SeqCst) as f64;
let requests_per_second = statistics.requests
.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_scrape_per_second = statistics.responses_scrape
.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_failure_per_second = statistics.responses_failure
.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let requests_per_second =
statistics.requests.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_scrape_per_second =
statistics.responses_scrape.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_failure_per_second =
statistics.responses_failure.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let bytes_sent_per_second = statistics.bytes_sent
.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let bytes_received_per_second = statistics.bytes_received
.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let bytes_sent_per_second =
statistics.bytes_sent.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let bytes_received_per_second =
statistics.bytes_received.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_announce_per_second = responses_announce / interval_f64;
let responses_announce_per_second = responses_announce / interval_f64;
let responses_per_second =
responses_announce_per_second +
responses_scrape_per_second +
responses_failure_per_second;
let responses_per_second = responses_announce_per_second
+ responses_scrape_per_second
+ responses_failure_per_second;
report_avg_response_vec.push(responses_per_second);
println!();
println!("Requests out: {:.2}/second", requests_per_second);
println!("Responses in: {:.2}/second", responses_per_second);
println!(" - Announce responses: {:.2}", responses_announce_per_second);
println!(
" - Announce responses: {:.2}",
responses_announce_per_second
);
println!(" - Scrape responses: {:.2}", responses_scrape_per_second);
println!(" - Failure responses: {:.2}", responses_failure_per_second);
println!(
" - Failure responses: {:.2}",
responses_failure_per_second
);
//println!("Peers per announce response: {:.2}", response_peers / responses_announce);
println!("Bandwidth out: {:.2}Mbit/s", bytes_sent_per_second * MBITS_FACTOR);
println!("Bandwidth in: {:.2}Mbit/s", bytes_received_per_second * MBITS_FACTOR);
println!(
"Bandwidth out: {:.2}Mbit/s",
bytes_sent_per_second * MBITS_FACTOR
);
println!(
"Bandwidth in: {:.2}Mbit/s",
bytes_received_per_second * MBITS_FACTOR
);
let time_elapsed = start_time.elapsed();
let duration = Duration::from_secs(config.duration as u64);
@ -151,7 +143,7 @@ fn monitor_statistics(
config
);
break
break;
}
}
}

126
aquatic_http_load_test/src/network.rs
View file

@ -1,16 +1,15 @@
use std::io::{Cursor, ErrorKind, Read, Write};
use std::sync::atomic::Ordering;
use std::time::Duration;
use std::io::{Read, Write, ErrorKind, Cursor};
use hashbrown::HashMap;
use mio::{net::TcpStream, Events, Poll, Interest, Token};
use rand::{rngs::SmallRng, prelude::*};
use mio::{net::TcpStream, Events, Interest, Poll, Token};
use rand::{prelude::*, rngs::SmallRng};
use crate::common::*;
use crate::config::*;
use crate::utils::create_random_request;
pub struct Connection {
stream: TcpStream,
read_buffer: [u8; 4096],
@ -18,7 +17,6 @@ pub struct Connection {
can_send: bool,
}
impl Connection {
pub fn create_and_register(
config: &Config,
@ -32,7 +30,7 @@ impl Connection {
.register(&mut stream, Token(*token_counter), Interest::READABLE)
.unwrap();
let connection = Connection {
let connection = Connection {
stream,
read_buffer: [0; 4096],
bytes_read: 0,
@ -46,14 +44,12 @@ impl Connection {
Ok(())
}
pub fn read_response(
&mut self,
state: &LoadTestState,
) -> bool { // bool = remove connection
pub fn read_response(&mut self, state: &LoadTestState) -> bool {
// bool = remove connection
loop {
match self.stream.read(&mut self.read_buffer[self.bytes_read..]){
match self.stream.read(&mut self.read_buffer[self.bytes_read..]) {
Ok(0) => {
if self.bytes_read == self.read_buffer.len(){
if self.bytes_read == self.read_buffer.len() {
eprintln!("read buffer is full");
}
@ -66,7 +62,7 @@ impl Connection {
let mut opt_body_start_index = None;
for (i, chunk) in interesting_bytes.windows(4).enumerate(){
for (i, chunk) in interesting_bytes.windows(4).enumerate() {
if chunk == b"\r\n\r\n" {
opt_body_start_index = Some(i + 4);
@ -77,30 +73,41 @@ impl Connection {
if let Some(body_start_index) = opt_body_start_index {
let interesting_bytes = &interesting_bytes[body_start_index..];
match Response::from_bytes(interesting_bytes){
match Response::from_bytes(interesting_bytes) {
Ok(response) => {
state.statistics.bytes_received
state
.statistics
.bytes_received
.fetch_add(self.bytes_read, Ordering::SeqCst);
match response {
Response::Announce(_) => {
state.statistics.responses_announce
state
.statistics
.responses_announce
.fetch_add(1, Ordering::SeqCst);
},
}
Response::Scrape(_) => {
state.statistics.responses_scrape
state
.statistics
.responses_scrape
.fetch_add(1, Ordering::SeqCst);
},
}
Response::Failure(response) => {
state.statistics.responses_failure
state
.statistics
.responses_failure
.fetch_add(1, Ordering::SeqCst);
println!("failure response: reason: {}", response.failure_reason);
},
println!(
"failure response: reason: {}",
response.failure_reason
);
}
}
self.bytes_read = 0;
self.can_send = true;
},
}
Err(err) => {
eprintln!(
"deserialize response error with {} bytes read: {:?}, text: {}",
@ -111,10 +118,10 @@ impl Connection {
}
}
}
},
}
Err(err) if err.kind() == ErrorKind::WouldBlock => {
break false;
},
}
Err(_) => {
self.bytes_read = 0;
@ -130,43 +137,40 @@ impl Connection {
state: &LoadTestState,
rng: &mut impl Rng,
request_buffer: &mut Cursor<&mut [u8]>,
) -> bool { // bool = remove connection
) -> bool {
// bool = remove connection
if !self.can_send {
return false;
}
let request = create_random_request(
&config,
&state,
rng
);
let request = create_random_request(&config, &state, rng);
request_buffer.set_position(0);
request.write(request_buffer).unwrap();
let position = request_buffer.position() as usize;
match self.send_request_inner(state, &request_buffer.get_mut()[..position]){
match self.send_request_inner(state, &request_buffer.get_mut()[..position]) {
Ok(()) => {
state.statistics.requests.fetch_add(1, Ordering::SeqCst);
self.can_send = false;
false
},
Err(_) => {
true
}
Err(_) => true,
}
}
fn send_request_inner(
&mut self,
state: &LoadTestState,
request: &[u8]
request: &[u8],
) -> ::std::io::Result<()> {
let bytes_sent = self.stream.write(request)?;
state.statistics.bytes_sent
state
.statistics
.bytes_sent
.fetch_add(bytes_sent, Ordering::SeqCst);
self.stream.flush()?;
@ -179,15 +183,9 @@ impl Connection {
}
}
pub type ConnectionMap = HashMap<usize, Connection>;
pub fn run_socket_thread(
config: &Config,
state: LoadTestState,
num_initial_requests: usize,
) {
pub fn run_socket_thread(config: &Config, state: LoadTestState, num_initial_requests: usize) {
let timeout = Duration::from_micros(config.network.poll_timeout_microseconds);
let create_conn_interval = 2 ^ config.network.connection_creation_interval;
@ -201,12 +199,8 @@ pub fn run_socket_thread(
let mut token_counter = 0usize;
for _ in 0..num_initial_requests {
Connection::create_and_register(
config,
&mut connections,
&mut poll,
&mut token_counter,
).unwrap();
Connection::create_and_register(config, &mut connections, &mut poll, &mut token_counter)
.unwrap();
}
let mut iter_counter = 0usize;
@ -218,14 +212,14 @@ pub fn run_socket_thread(
poll.poll(&mut events, Some(timeout))
.expect("failed polling");
for event in events.iter(){
if event.is_readable(){
for event in events.iter() {
if event.is_readable() {
let token = event.token();
if let Some(connection) = connections.get_mut(&token.0){
if let Some(connection) = connections.get_mut(&token.0) {
// Note that this does not indicate successfully reading
// response
if connection.read_response(&state){
if connection.read_response(&state) {
remove_connection(&mut poll, &mut connections, token.0);
num_to_create += 1;
@ -236,13 +230,9 @@ pub fn run_socket_thread(
}
}
for (k, connection) in connections.iter_mut(){
let remove_connection = connection.send_request(
config,
&state,
&mut rng,
&mut request_buffer
);
for (k, connection) in connections.iter_mut() {
let remove_connection =
connection.send_request(config, &state, &mut rng, &mut request_buffer);
if remove_connection {
drop_connections.push(*k);
@ -269,7 +259,8 @@ pub fn run_socket_thread(
&mut connections,
&mut poll,
&mut token_counter,
).is_ok();
)
.is_ok();
if ok {
num_to_create -= 1;
@ -280,14 +271,9 @@ pub fn run_socket_thread(
}
}
fn remove_connection(
poll: &mut Poll,
connections: &mut ConnectionMap,
connection_id: usize,
){
if let Some(mut connection) = connections.remove(&connection_id){
if let Err(err) = connection.deregister(poll){
fn remove_connection(poll: &mut Poll, connections: &mut ConnectionMap, connection_id: usize) {
if let Some(mut connection) = connections.remove(&connection_id) {
if let Err(err) = connection.deregister(poll) {
eprintln!("couldn't deregister connection: {}", err);
}
}

53
aquatic_http_load_test/src/utils.rs
View file

@ -1,13 +1,12 @@
use std::sync::Arc;
use rand::distributions::WeightedIndex;
use rand_distr::Pareto;
use rand::prelude::*;
use rand_distr::Pareto;
use crate::common::*;
use crate::config::*;
pub fn create_random_request(
config: &Config,
state: &LoadTestState,
@ -18,35 +17,18 @@ pub fn create_random_request(
config.torrents.weight_scrape as u32,
];
let items = [
RequestType::Announce,
RequestType::Scrape,
];
let items = [RequestType::Announce, RequestType::Scrape];
let dist = WeightedIndex::new(&weights)
.expect("random request weighted index");
let dist = WeightedIndex::new(&weights).expect("random request weighted index");
match items[dist.sample(rng)] {
RequestType::Announce => create_announce_request(
config,
state,
rng,
),
RequestType::Scrape => create_scrape_request(
config,
state,
rng,
)
RequestType::Announce => create_announce_request(config, state, rng),
RequestType::Scrape => create_scrape_request(config, state, rng),
}
}
#[inline]
fn create_announce_request(
config: &Config,
state: &LoadTestState,
rng: &mut impl Rng,
) -> Request {
fn create_announce_request(config: &Config, state: &LoadTestState, rng: &mut impl Rng) -> Request {
let (event, bytes_left) = {
if rng.gen_bool(config.torrents.peer_seeder_probability) {
(AnnounceEvent::Completed, 0)
@ -65,17 +47,12 @@ fn create_announce_request(
key: None,
numwant: None,
compact: true,
port: rng.gen()
port: rng.gen(),
})
}
#[inline]
fn create_scrape_request(
config: &Config,
state: &LoadTestState,
rng: &mut impl Rng,
) -> Request {
fn create_scrape_request(config: &Config, state: &LoadTestState, rng: &mut impl Rng) -> Request {
let mut scrape_hashes = Vec::with_capacity(5);
for _ in 0..5 {
@ -89,23 +66,13 @@ fn create_scrape_request(
})
}
#[inline]
fn select_info_hash_index(
config: &Config,
state: &LoadTestState,
rng: &mut impl Rng,
) -> usize {
fn select_info_hash_index(config: &Config, state: &LoadTestState, rng: &mut impl Rng) -> usize {
pareto_usize(rng, &state.pareto, config.torrents.number_of_torrents - 1)
}
#[inline]
fn pareto_usize(
rng: &mut impl Rng,
pareto: &Arc<Pareto<f64>>,
max: usize,
) -> usize {
fn pareto_usize(rng: &mut impl Rng, pareto: &Arc<Pareto<f64>>, max: usize) -> usize {
let p: f64 = pareto.sample(rng);
let p = (p.min(101.0f64) - 1.0) / 100.0;

5
aquatic_http_protocol/Cargo.toml
View file

@ -11,11 +11,6 @@ exclude = ["target"]
[lib]
name = "aquatic_http_protocol"
[[bench]]
name = "bench_request_from_path"
path = "benches/bench_request_from_path.rs"
harness = false
[[bench]]
name = "bench_request_from_bytes"
path = "benches/bench_request_from_bytes.rs"

15
aquatic_http_protocol/benches/bench_announce_response_to_bytes.rs
View file

@ -5,14 +5,13 @@ use criterion::{black_box, criterion_group, criterion_main, Criterion};
use aquatic_http_protocol::response::*;
pub fn bench(c: &mut Criterion) {
let mut peers = Vec::new();
for i in 0..100 {
peers.push(ResponsePeer {
ip_address: Ipv4Addr::new(127, 0, 0, i),
port: i as u16
port: i as u16,
})
}
@ -29,14 +28,16 @@ pub fn bench(c: &mut Criterion) {
let mut buffer = [0u8; 4096];
let mut buffer = ::std::io::Cursor::new(&mut buffer[..]);
c.bench_function("announce-response-to-bytes", |b| b.iter(|| {
buffer.set_position(0);
c.bench_function("announce-response-to-bytes", |b| {
b.iter(|| {
buffer.set_position(0);
Response::write(black_box(&response), black_box(&mut buffer)).unwrap();
}));
Response::write(black_box(&response), black_box(&mut buffer)).unwrap();
})
});
}
criterion_group!{
criterion_group! {
name = benches;
config = Criterion::default()
.sample_size(1000)

12
aquatic_http_protocol/benches/bench_request_from_bytes.rs
View file

@ -1,19 +1,17 @@
use std::time::Duration;
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use std::time::Duration;
use aquatic_http_protocol::request::Request;
static INPUT: &[u8] = b"GET /announce?info_hash=%04%0bkV%3f%5cr%14%a6%b7%98%adC%c3%c9.%40%24%00%b9&peer_id=-TR2940-5ert69muw5t8&port=11000&uploaded=0&downloaded=0&left=0&numwant=0&key=3ab4b977&compact=1&supportcrypto=1&event=stopped HTTP/1.1\r\n\r\n";
pub fn bench(c: &mut Criterion) {
c.bench_function("request-from-bytes", |b| b.iter(||
Request::from_bytes(black_box(INPUT))
));
c.bench_function("request-from-bytes", |b| {
b.iter(|| Request::from_bytes(black_box(INPUT)))
});
}
criterion_group!{
criterion_group! {
name = benches;
config = Criterion::default()
.sample_size(1000)

28
aquatic_http_protocol/src/common.rs
View file

@ -1,48 +1,43 @@
use std::str::FromStr;
use serde::{Serialize, Deserialize};
use serde::{Deserialize, Serialize};
use super::utils::*;
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct PeerId(
#[serde(
serialize_with = "serialize_20_bytes",
deserialize_with = "deserialize_20_bytes",
deserialize_with = "deserialize_20_bytes"
)]
pub [u8; 20]
pub [u8; 20],
);
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(transparent)]
pub struct InfoHash(
#[serde(
serialize_with = "serialize_20_bytes",
deserialize_with = "deserialize_20_bytes",
deserialize_with = "deserialize_20_bytes"
)]
pub [u8; 20]
pub [u8; 20],
);
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AnnounceEvent {
Started,
Stopped,
Completed,
Empty
Empty,
}
impl Default for AnnounceEvent {
fn default() -> Self {
Self::Empty
}
}
impl FromStr for AnnounceEvent {
type Err = String;
@ -52,18 +47,17 @@ impl FromStr for AnnounceEvent {
"stopped" => Ok(Self::Stopped),
"completed" => Ok(Self::Completed),
"empty" => Ok(Self::Empty),
value => Err(format!("Unknown value: {}", value))
value => Err(format!("Unknown value: {}", value)),
}
}
}
#[cfg(test)]
impl quickcheck::Arbitrary for InfoHash {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
let mut arr = [b'0'; 20];
for byte in arr.iter_mut(){
for byte in arr.iter_mut() {
*byte = u8::arbitrary(g);
}
@ -71,13 +65,12 @@ impl quickcheck::Arbitrary for InfoHash {
}
}
#[cfg(test)]
impl quickcheck::Arbitrary for PeerId {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
let mut arr = [b'0'; 20];
for byte in arr.iter_mut(){
for byte in arr.iter_mut() {
*byte = u8::arbitrary(g);
}
@ -85,11 +78,10 @@ impl quickcheck::Arbitrary for PeerId {
}
}
#[cfg(test)]
impl quickcheck::Arbitrary for AnnounceEvent {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
match (bool::arbitrary(g), bool::arbitrary(g)){
match (bool::arbitrary(g), bool::arbitrary(g)) {
(false, false) => Self::Started,
(true, false) => Self::Started,
(false, true) => Self::Completed,

119
aquatic_http_protocol/src/request.rs
View file

@ -1,12 +1,11 @@
use std::io::Write;
use anyhow::Context;
use smartstring::{SmartString, LazyCompact};
use smartstring::{LazyCompact, SmartString};
use super::common::*;
use super::utils::*;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AnnounceRequest {
pub info_hash: InfoHash,
@ -20,7 +19,6 @@ pub struct AnnounceRequest {
pub key: Option<SmartString<LazyCompact>>,
}
impl AnnounceRequest {
fn write<W: Write>(&self, output: &mut W) -> ::std::io::Result<()> {
output.write_all(b"GET /announce?info_hash=")?;
@ -61,13 +59,11 @@ impl AnnounceRequest {
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ScrapeRequest {
pub info_hashes: Vec<InfoHash>,
}
impl ScrapeRequest {
fn write<W: Write>(&self, output: &mut W) -> ::std::io::Result<()> {
output.write_all(b"GET /scrape?")?;
@ -91,49 +87,40 @@ impl ScrapeRequest {
}
}
#[derive(Debug)]
pub enum RequestParseError {
NeedMoreData,
Invalid(anyhow::Error),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Request {
Announce(AnnounceRequest),
Scrape(ScrapeRequest),
}
impl Request {
/// Parse Request from HTTP request bytes
pub fn from_bytes(bytes: &[u8]) -> Result<Self, RequestParseError> {
let mut headers = [httparse::EMPTY_HEADER; 16];
let mut http_request = httparse::Request::new(&mut headers);
let path = match http_request.parse(bytes){
let path = match http_request.parse(bytes) {
Ok(httparse::Status::Complete(_)) => {
if let Some(path) = http_request.path {
path
} else {
return Err(RequestParseError::Invalid(
anyhow::anyhow!("no http path")
))
return Err(RequestParseError::Invalid(anyhow::anyhow!("no http path")));
}
},
}
Ok(httparse::Status::Partial) => {
if let Some(path) = http_request.path {
path
} else {
return Err(RequestParseError::NeedMoreData)
return Err(RequestParseError::NeedMoreData);
}
}
Err(err) => {
return Err(RequestParseError::Invalid(
anyhow::Error::from(err)
))
},
Err(err) => return Err(RequestParseError::Invalid(anyhow::Error::from(err))),
};
Self::from_http_get_path(path).map_err(RequestParseError::Invalid)
@ -155,12 +142,10 @@ impl Request {
pub fn from_http_get_path(path: &str) -> anyhow::Result<Self> {
::log::debug!("request GET path: {}", path);
let mut split_parts= path.splitn(2, '?');
let mut split_parts = path.splitn(2, '?');
let location = split_parts.next()
.with_context(|| "no location")?;
let query_string = split_parts.next()
.with_context(|| "no query string")?;
let location = split_parts.next().with_context(|| "no location")?;
let query_string = split_parts.next().with_context(|| "no query string")?;
// -- Parse key-value pairs
@ -177,58 +162,61 @@ impl Request {
let mut ampersand_iter = ::memchr::memchr_iter(b'&', query_string_bytes);
let mut position = 0usize;
for equal_sign_index in ::memchr::memchr_iter(b'=', query_string_bytes){
let segment_end = ampersand_iter.next()
.unwrap_or_else(|| query_string.len());
for equal_sign_index in ::memchr::memchr_iter(b'=', query_string_bytes) {
let segment_end = ampersand_iter.next().unwrap_or_else(|| query_string.len());
let key = query_string.get(position..equal_sign_index)
let key = query_string
.get(position..equal_sign_index)
.with_context(|| format!("no key at {}..{}", position, equal_sign_index))?;
let value = query_string.get(equal_sign_index + 1..segment_end)
.with_context(|| format!("no value at {}..{}", equal_sign_index + 1, segment_end))?;
let value = query_string
.get(equal_sign_index + 1..segment_end)
.with_context(|| {
format!("no value at {}..{}", equal_sign_index + 1, segment_end)
})?;
match key {
"info_hash" => {
let value = urldecode_20_bytes(value)?;
info_hashes.push(InfoHash(value));
},
}
"peer_id" => {
let value = urldecode_20_bytes(value)?;
opt_peer_id = Some(PeerId(value));
},
}
"port" => {
opt_port = Some(value.parse::<u16>().with_context(|| "parse port")?);
},
"left" => {
}
"left" => {
opt_bytes_left = Some(value.parse::<usize>().with_context(|| "parse left")?);
},
"event" => {
event = value.parse::<AnnounceEvent>().map_err(|err|
anyhow::anyhow!("invalid event: {}", err)
)?;
},
"compact" => {
}
"event" => {
event = value
.parse::<AnnounceEvent>()
.map_err(|err| anyhow::anyhow!("invalid event: {}", err))?;
}
"compact" => {
if value != "1" {
return Err(anyhow::anyhow!("compact set, but not to 1"));
}
},
"numwant" => {
}
"numwant" => {
opt_numwant = Some(value.parse::<usize>().with_context(|| "parse numwant")?);
},
}
"key" => {
if value.len() > 100 {
return Err(anyhow::anyhow!("'key' is too long"))
return Err(anyhow::anyhow!("'key' is too long"));
}
opt_key = Some(::urlencoding::decode(value)?.into());
},
}
k => {
::log::debug!("ignored unrecognized key: {}", k)
}
}
if segment_end == query_string.len(){
break
if segment_end == query_string.len() {
break;
} else {
position = segment_end + 1;
}
@ -250,9 +238,7 @@ impl Request {
Ok(Request::Announce(request))
} else {
let request = ScrapeRequest {
info_hashes,
};
let request = ScrapeRequest { info_hashes };
Ok(Request::Scrape(request))
}
@ -266,17 +252,23 @@ impl Request {
}
}
#[cfg(test)]
mod tests {
use quickcheck::{Arbitrary, Gen, TestResult, quickcheck};
use quickcheck::{quickcheck, Arbitrary, Gen, TestResult};
use super::*;
static ANNOUNCE_REQUEST_PATH: &str = "/announce?info_hash=%04%0bkV%3f%5cr%14%a6%b7%98%adC%c3%c9.%40%24%00%b9&peer_id=-ABC940-5ert69muw5t8&port=12345&uploaded=0&downloaded=0&left=1&numwant=0&key=4ab4b877&compact=1&supportcrypto=1&event=started";
static SCRAPE_REQUEST_PATH: &str = "/scrape?info_hash=%04%0bkV%3f%5cr%14%a6%b7%98%adC%c3%c9.%40%24%00%b9";
static REFERENCE_INFO_HASH: [u8; 20] = [0x04, 0x0b, b'k', b'V', 0x3f, 0x5c, b'r', 0x14, 0xa6, 0xb7, 0x98, 0xad, b'C', 0xc3, 0xc9, b'.', 0x40, 0x24, 0x00, 0xb9];
static REFERENCE_PEER_ID: [u8; 20] = [b'-', b'A', b'B', b'C', b'9', b'4', b'0', b'-', b'5', b'e', b'r', b't', b'6', b'9', b'm', b'u', b'w', b'5', b't', b'8'];
static SCRAPE_REQUEST_PATH: &str =
"/scrape?info_hash=%04%0bkV%3f%5cr%14%a6%b7%98%adC%c3%c9.%40%24%00%b9";
static REFERENCE_INFO_HASH: [u8; 20] = [
0x04, 0x0b, b'k', b'V', 0x3f, 0x5c, b'r', 0x14, 0xa6, 0xb7, 0x98, 0xad, b'C', 0xc3, 0xc9,
b'.', 0x40, 0x24, 0x00, 0xb9,
];
static REFERENCE_PEER_ID: [u8; 20] = [
b'-', b'A', b'B', b'C', b'9', b'4', b'0', b'-', b'5', b'e', b'r', b't', b'6', b'9', b'm',
b'u', b'w', b'5', b't', b'8',
];
fn get_reference_announce_request() -> Request {
Request::Announce(AnnounceRequest {
@ -287,12 +279,12 @@ mod tests {
event: AnnounceEvent::Started,
compact: true,
numwant: Some(0),
key: Some("4ab4b877".into())
key: Some("4ab4b877".into()),
})
}
#[test]
fn test_announce_request_from_bytes(){
fn test_announce_request_from_bytes() {
let mut bytes = Vec::new();
bytes.extend_from_slice(b"GET ");
@ -306,7 +298,7 @@ mod tests {
}
#[test]
fn test_scrape_request_from_bytes(){
fn test_scrape_request_from_bytes() {
let mut bytes = Vec::new();
bytes.extend_from_slice(b"GET ");
@ -348,7 +340,7 @@ mod tests {
impl Arbitrary for Request {
fn arbitrary(g: &mut Gen) -> Self {
if Arbitrary::arbitrary(g){
if Arbitrary::arbitrary(g) {
Self::Announce(Arbitrary::arbitrary(g))
} else {
Self::Scrape(Arbitrary::arbitrary(g))
@ -357,9 +349,12 @@ mod tests {
}
#[test]
fn quickcheck_serde_identity_request(){
fn quickcheck_serde_identity_request() {
fn prop(request: Request) -> TestResult {
if let Request::Announce(AnnounceRequest { key: Some(ref key), ..}) = request {
if let Request::Announce(AnnounceRequest {
key: Some(ref key), ..
}) = request
{
if key.len() > 30 {
return TestResult::discard();
}

97
aquatic_http_protocol/src/response.rs
View file

@ -1,42 +1,38 @@
use std::net::{Ipv4Addr, Ipv6Addr};
use std::io::Write;
use std::net::{Ipv4Addr, Ipv6Addr};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use serde::{Serialize, Deserialize};
use super::common::*;
use super::utils::*;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResponsePeer<I: Eq>{
pub struct ResponsePeer<I: Eq> {
pub ip_address: I,
pub port: u16
pub port: u16,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[serde(transparent)]
pub struct ResponsePeerListV4(
#[serde(
serialize_with = "serialize_response_peers_ipv4",
deserialize_with = "deserialize_response_peers_ipv4",
deserialize_with = "deserialize_response_peers_ipv4"
)]
pub Vec<ResponsePeer<Ipv4Addr>>
pub Vec<ResponsePeer<Ipv4Addr>>,
);
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[serde(transparent)]
pub struct ResponsePeerListV6(
#[serde(
serialize_with = "serialize_response_peers_ipv6",
deserialize_with = "deserialize_response_peers_ipv6",
deserialize_with = "deserialize_response_peers_ipv6"
)]
pub Vec<ResponsePeer<Ipv6Addr>>
pub Vec<ResponsePeer<Ipv6Addr>>,
);
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScrapeStatistics {
pub complete: usize,
@ -44,7 +40,6 @@ pub struct ScrapeStatistics {
pub downloaded: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AnnounceResponse {
#[serde(rename = "interval")]
@ -57,47 +52,44 @@ pub struct AnnounceResponse {
pub peers6: ResponsePeerListV6,
}
impl AnnounceResponse {
fn write<W: Write>(&self, output: &mut W) -> ::std::io::Result<usize> {
let mut bytes_written = 0usize;
bytes_written += output.write(b"d8:completei")?;
bytes_written += output.write(
itoa::Buffer::new().format(self.complete).as_bytes()
)?;
bytes_written += output.write(itoa::Buffer::new().format(self.complete).as_bytes())?;
bytes_written += output.write(b"e10:incompletei")?;
bytes_written += output.write(
itoa::Buffer::new().format(self.incomplete).as_bytes()
)?;
bytes_written += output.write(itoa::Buffer::new().format(self.incomplete).as_bytes())?;
bytes_written += output.write(b"e8:intervali")?;
bytes_written += output.write(
itoa::Buffer::new().format(self.announce_interval).as_bytes()
itoa::Buffer::new()
.format(self.announce_interval)
.as_bytes(),
)?;
bytes_written += output.write(b"e5:peers")?;
bytes_written += output.write(
itoa::Buffer::new().format(self.peers.0.len() * 6).as_bytes()
itoa::Buffer::new()
.format(self.peers.0.len() * 6)
.as_bytes(),
)?;
bytes_written += output.write(b":")?;
for peer in self.peers.0.iter() {
bytes_written += output.write(
&u32::from(peer.ip_address).to_be_bytes()
)?;
bytes_written += output.write(&u32::from(peer.ip_address).to_be_bytes())?;
bytes_written += output.write(&peer.port.to_be_bytes())?;
}
bytes_written += output.write(b"6:peers6")?;
bytes_written += output.write(
itoa::Buffer::new().format(self.peers6.0.len() * 18).as_bytes()
itoa::Buffer::new()
.format(self.peers6.0.len() * 18)
.as_bytes(),
)?;
bytes_written += output.write(b":")?;
for peer in self.peers6.0.iter() {
bytes_written += output.write(
&u128::from(peer.ip_address).to_be_bytes()
)?;
bytes_written += output.write(&u128::from(peer.ip_address).to_be_bytes())?;
bytes_written += output.write(&peer.port.to_be_bytes())?;
}
bytes_written += output.write(b"e")?;
@ -106,31 +98,27 @@ impl AnnounceResponse {
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScrapeResponse {
/// BTreeMap instead of HashMap since keys need to be serialized in order
pub files: BTreeMap<InfoHash, ScrapeStatistics>,
}
impl ScrapeResponse {
fn write<W: Write>(&self, output: &mut W) -> ::std::io::Result<usize> {
let mut bytes_written = 0usize;
bytes_written += output.write(b"d5:filesd")?;
for (info_hash, statistics) in self.files.iter(){
for (info_hash, statistics) in self.files.iter() {
bytes_written += output.write(b"20:")?;
bytes_written += output.write(&info_hash.0)?;
bytes_written += output.write(b"d8:completei")?;
bytes_written += output.write(
itoa::Buffer::new().format(statistics.complete).as_bytes()
)?;
bytes_written +=
output.write(itoa::Buffer::new().format(statistics.complete).as_bytes())?;
bytes_written += output.write(b"e10:downloadedi0e10:incompletei")?;
bytes_written += output.write(
itoa::Buffer::new().format(statistics.incomplete).as_bytes()
)?;
bytes_written +=
output.write(itoa::Buffer::new().format(statistics.incomplete).as_bytes())?;
bytes_written += output.write(b"ee")?;
}
@ -140,14 +128,12 @@ impl ScrapeResponse {
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FailureResponse {
#[serde(rename = "failure reason")]
pub failure_reason: String,
}
impl FailureResponse {
fn write<W: Write>(&self, output: &mut W) -> ::std::io::Result<usize> {
let mut bytes_written = 0usize;
@ -155,9 +141,7 @@ impl FailureResponse {
let reason_bytes = self.failure_reason.as_bytes();
bytes_written += output.write(b"d14:failure reason")?;
bytes_written += output.write(
itoa::Buffer::new().format(reason_bytes.len()).as_bytes()
)?;
bytes_written += output.write(itoa::Buffer::new().format(reason_bytes.len()).as_bytes())?;
bytes_written += output.write(b":")?;
bytes_written += output.write(reason_bytes)?;
bytes_written += output.write(b"e")?;
@ -166,7 +150,6 @@ impl FailureResponse {
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum Response {
@ -175,7 +158,6 @@ pub enum Response {
Failure(FailureResponse),
}
impl Response {
pub fn write<W: Write>(&self, output: &mut W) -> ::std::io::Result<usize> {
match self {
@ -189,29 +171,26 @@ impl Response {
}
}
#[cfg(test)]
impl quickcheck::Arbitrary for ResponsePeer<Ipv4Addr> {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
Self {
ip_address: Ipv4Addr::arbitrary(g),
port: u16::arbitrary(g)
port: u16::arbitrary(g),
}
}
}
#[cfg(test)]
impl quickcheck::Arbitrary for ResponsePeer<Ipv6Addr> {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
Self {
ip_address: Ipv6Addr::arbitrary(g),
port: u16::arbitrary(g)
port: u16::arbitrary(g),
}
}
}
#[cfg(test)]
impl quickcheck::Arbitrary for ResponsePeerListV4 {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
@ -219,7 +198,6 @@ impl quickcheck::Arbitrary for ResponsePeerListV4 {
}
}
#[cfg(test)]
impl quickcheck::Arbitrary for ResponsePeerListV6 {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
@ -227,7 +205,6 @@ impl quickcheck::Arbitrary for ResponsePeerListV6 {
}
}
#[cfg(test)]
impl quickcheck::Arbitrary for ScrapeStatistics {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
@ -239,7 +216,6 @@ impl quickcheck::Arbitrary for ScrapeStatistics {
}
}
#[cfg(test)]
impl quickcheck::Arbitrary for AnnounceResponse {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
@ -253,7 +229,6 @@ impl quickcheck::Arbitrary for AnnounceResponse {
}
}
#[cfg(test)]
impl quickcheck::Arbitrary for ScrapeResponse {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
@ -263,7 +238,6 @@ impl quickcheck::Arbitrary for ScrapeResponse {
}
}
#[cfg(test)]
impl quickcheck::Arbitrary for FailureResponse {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
@ -273,7 +247,6 @@ impl quickcheck::Arbitrary for FailureResponse {
}
}
#[cfg(test)]
mod tests {
use quickcheck_macros::*;
@ -282,9 +255,7 @@ mod tests {
#[quickcheck]
fn test_announce_response_to_bytes(response: AnnounceResponse) -> bool {
let reference = bendy::serde::to_bytes(
&Response::Announce(response.clone())
).unwrap();
let reference = bendy::serde::to_bytes(&Response::Announce(response.clone())).unwrap();
let mut output = Vec::new();
@ -295,9 +266,7 @@ mod tests {
#[quickcheck]
fn test_scrape_response_to_bytes(response: ScrapeResponse) -> bool {
let reference = bendy::serde::to_bytes(
&Response::Scrape(response.clone())
).unwrap();
let reference = bendy::serde::to_bytes(&Response::Scrape(response.clone())).unwrap();
let mut hand_written = Vec::new();
@ -315,9 +284,7 @@ mod tests {
#[quickcheck]
fn test_failure_response_to_bytes(response: FailureResponse) -> bool {
let reference = bendy::serde::to_bytes(
&Response::Failure(response.clone())
).unwrap();
let reference = bendy::serde::to_bytes(&Response::Failure(response.clone())).unwrap();
let mut hand_written = Vec::new();

159
aquatic_http_protocol/src/utils.rs
View file

@ -1,23 +1,16 @@
use std::net::{Ipv4Addr, Ipv6Addr};
use std::io::Write;
use std::net::{Ipv4Addr, Ipv6Addr};
use anyhow::Context;
use serde::{Serializer, Deserializer, de::Visitor};
use serde::{de::Visitor, Deserializer, Serializer};
use super::response::ResponsePeer;
pub fn urlencode_20_bytes(
input: [u8; 20],
output: &mut impl Write
) -> ::std::io::Result<()> {
pub fn urlencode_20_bytes(input: [u8; 20], output: &mut impl Write) -> ::std::io::Result<()> {
let mut tmp = [b'%'; 60];
for i in 0..input.len() {
hex::encode_to_slice(
&input[i..i + 1],
&mut tmp[i * 3 + 1..i * 3 + 3]
).unwrap();
hex::encode_to_slice(&input[i..i + 1], &mut tmp[i * 3 + 1..i * 3 + 3]).unwrap();
}
output.write_all(&tmp)?;
@ -25,15 +18,13 @@ pub fn urlencode_20_bytes(
Ok(())
}
pub fn urldecode_20_bytes(value: &str) -> anyhow::Result<[u8; 20]> {
let mut out_arr = [0u8; 20];
let mut chars = value.chars();
for i in 0..20 {
let c = chars.next()
.with_context(|| "less than 20 chars")?;
let c = chars.next().with_context(|| "less than 20 chars")?;
if c as u32 > 255 {
return Err(anyhow::anyhow!(
@ -43,38 +34,37 @@ pub fn urldecode_20_bytes(value: &str) -> anyhow::Result<[u8; 20]> {
}
if c == '%' {
let first = chars.next()
let first = chars
.next()
.with_context(|| "missing first urldecode char in pair")?;
let second = chars.next()
let second = chars
.next()
.with_context(|| "missing second urldecode char in pair")?;
let hex = [first as u8, second as u8];
hex::decode_to_slice(&hex, &mut out_arr[i..i+1]).map_err(|err|
anyhow::anyhow!("hex decode error: {:?}", err)
)?;
hex::decode_to_slice(&hex, &mut out_arr[i..i + 1])
.map_err(|err| anyhow::anyhow!("hex decode error: {:?}", err))?;
} else {
out_arr[i] = c as u8;
}
}
if chars.next().is_some(){
if chars.next().is_some() {
return Err(anyhow::anyhow!("more than 20 chars"));
}
Ok(out_arr)
}
#[inline]
pub fn serialize_20_bytes<S>(
bytes: &[u8; 20],
serializer: S
) -> Result<S::Ok, S::Error> where S: Serializer {
pub fn serialize_20_bytes<S>(bytes: &[u8; 20], serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_bytes(bytes)
}
struct TwentyByteVisitor;
impl<'de> Visitor<'de> for TwentyByteVisitor {
@ -86,7 +76,8 @@ impl<'de> Visitor<'de> for TwentyByteVisitor {
#[inline]
fn visit_bytes<E>(self, value: &[u8]) -> Result<Self::Value, E>
where E: ::serde::de::Error,
where
E: ::serde::de::Error,
{
if value.len() != 20 {
return Err(::serde::de::Error::custom("not 20 bytes"));
@ -100,21 +91,21 @@ impl<'de> Visitor<'de> for TwentyByteVisitor {
}
}
#[inline]
pub fn deserialize_20_bytes<'de, D>(
deserializer: D
) -> Result<[u8; 20], D::Error>
where D: Deserializer<'de>
pub fn deserialize_20_bytes<'de, D>(deserializer: D) -> Result<[u8; 20], D::Error>
where
D: Deserializer<'de>,
{
deserializer.deserialize_any(TwentyByteVisitor)
}
pub fn serialize_response_peers_ipv4<S>(
response_peers: &[ResponsePeer<Ipv4Addr>],
serializer: S
) -> Result<S::Ok, S::Error> where S: Serializer {
serializer: S,
) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let mut bytes = Vec::with_capacity(response_peers.len() * 6);
for peer in response_peers {
@ -125,11 +116,13 @@ pub fn serialize_response_peers_ipv4<S>(
serializer.serialize_bytes(&bytes)
}
pub fn serialize_response_peers_ipv6<S>(
response_peers: &[ResponsePeer<Ipv6Addr>],
serializer: S
) -> Result<S::Ok, S::Error> where S: Serializer {
serializer: S,
) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let mut bytes = Vec::with_capacity(response_peers.len() * 6);
for peer in response_peers {
@ -140,10 +133,8 @@ pub fn serialize_response_peers_ipv6<S>(
serializer.serialize_bytes(&bytes)
}
struct ResponsePeersIpv4Visitor;
impl<'de> Visitor<'de> for ResponsePeersIpv4Visitor {
type Value = Vec<ResponsePeer<Ipv4Addr>>;
@ -153,45 +144,47 @@ impl<'de> Visitor<'de> for ResponsePeersIpv4Visitor {
#[inline]
fn visit_bytes<E>(self, value: &[u8]) -> Result<Self::Value, E>
where E: ::serde::de::Error,
where
E: ::serde::de::Error,
{
let chunks = value.chunks_exact(6);
if !chunks.remainder().is_empty(){
if !chunks.remainder().is_empty() {
return Err(::serde::de::Error::custom("trailing bytes"));
}
let mut ip_bytes = [0u8; 4];
let mut port_bytes = [0u8; 2];
let peers = chunks.into_iter().map(|chunk | {
ip_bytes.copy_from_slice(&chunk[0..4]);
port_bytes.copy_from_slice(&chunk[4..6]);
let peers = chunks
.into_iter()
.map(|chunk| {
ip_bytes.copy_from_slice(&chunk[0..4]);
port_bytes.copy_from_slice(&chunk[4..6]);
ResponsePeer {
ip_address: Ipv4Addr::from(u32::from_be_bytes(ip_bytes)),
port: u16::from_be_bytes(port_bytes),
}
}).collect();
ResponsePeer {
ip_address: Ipv4Addr::from(u32::from_be_bytes(ip_bytes)),
port: u16::from_be_bytes(port_bytes),
}
})
.collect();
Ok(peers)
}
}
#[inline]
pub fn deserialize_response_peers_ipv4<'de, D>(
deserializer: D
deserializer: D,
) -> Result<Vec<ResponsePeer<Ipv4Addr>>, D::Error>
where D: Deserializer<'de>
where
D: Deserializer<'de>,
{
deserializer.deserialize_any(ResponsePeersIpv4Visitor)
}
struct ResponsePeersIpv6Visitor;
impl<'de> Visitor<'de> for ResponsePeersIpv6Visitor {
type Value = Vec<ResponsePeer<Ipv6Addr>>;
@ -201,42 +194,45 @@ impl<'de> Visitor<'de> for ResponsePeersIpv6Visitor {
#[inline]
fn visit_bytes<E>(self, value: &[u8]) -> Result<Self::Value, E>
where E: ::serde::de::Error,
where
E: ::serde::de::Error,
{
let chunks = value.chunks_exact(18);
if !chunks.remainder().is_empty(){
if !chunks.remainder().is_empty() {
return Err(::serde::de::Error::custom("trailing bytes"));
}
let mut ip_bytes = [0u8; 16];
let mut port_bytes = [0u8; 2];
let peers = chunks.into_iter().map(|chunk| {
ip_bytes.copy_from_slice(&chunk[0..16]);
port_bytes.copy_from_slice(&chunk[16..18]);
let peers = chunks
.into_iter()
.map(|chunk| {
ip_bytes.copy_from_slice(&chunk[0..16]);
port_bytes.copy_from_slice(&chunk[16..18]);
ResponsePeer {
ip_address: Ipv6Addr::from(u128::from_be_bytes(ip_bytes)),
port: u16::from_be_bytes(port_bytes),
}
}).collect();
ResponsePeer {
ip_address: Ipv6Addr::from(u128::from_be_bytes(ip_bytes)),
port: u16::from_be_bytes(port_bytes),
}
})
.collect();
Ok(peers)
}
}
#[inline]
pub fn deserialize_response_peers_ipv6<'de, D>(
deserializer: D
deserializer: D,
) -> Result<Vec<ResponsePeer<Ipv6Addr>>, D::Error>
where D: Deserializer<'de>
where
D: Deserializer<'de>,
{
deserializer.deserialize_any(ResponsePeersIpv6Visitor)
}
#[cfg(test)]
mod tests {
use quickcheck_macros::*;
@ -246,10 +242,10 @@ mod tests {
use super::*;
#[test]
fn test_urlencode_20_bytes(){
fn test_urlencode_20_bytes() {
let mut input = [0u8; 20];
for (i, b) in input.iter_mut().enumerate(){
for (i, b) in input.iter_mut().enumerate() {
*b = i as u8 % 10;
}
@ -259,7 +255,7 @@ mod tests {
assert_eq!(output.len(), 60);
for (i, chunk) in output.chunks_exact(3).enumerate(){
for (i, chunk) in output.chunks_exact(3).enumerate() {
// Not perfect but should do the job
let reference = [b'%', b'0', input[i] + 48];
@ -284,9 +280,7 @@ mod tests {
g: u8,
h: u8,
) -> bool {
let input: [u8; 20] = [
a, b, c, d, e, f, g, h, b, c, d, a, e, f, g, h, a, b, d, c
];
let input: [u8; 20] = [a, b, c, d, e, f, g, h, b, c, d, a, e, f, g, h, a, b, d, c];
let mut output = Vec::new();
@ -302,9 +296,7 @@ mod tests {
}
#[quickcheck]
fn test_serde_response_peers_ipv4(
peers: Vec<ResponsePeer<Ipv4Addr>>,
) -> bool {
fn test_serde_response_peers_ipv4(peers: Vec<ResponsePeer<Ipv4Addr>>) -> bool {
let serialized = bendy::serde::to_bytes(&peers).unwrap();
let deserialized: Vec<ResponsePeer<Ipv4Addr>> =
::bendy::serde::from_bytes(&serialized).unwrap();
@ -313,9 +305,7 @@ mod tests {
}
#[quickcheck]
fn test_serde_response_peers_ipv6(
peers: Vec<ResponsePeer<Ipv6Addr>>,
) -> bool {
fn test_serde_response_peers_ipv6(peers: Vec<ResponsePeer<Ipv6Addr>>) -> bool {
let serialized = bendy::serde::to_bytes(&peers).unwrap();
let deserialized: Vec<ResponsePeer<Ipv6Addr>> =
::bendy::serde::from_bytes(&serialized).unwrap();
@ -324,12 +314,9 @@ mod tests {
}
#[quickcheck]
fn test_serde_info_hash(
info_hash: InfoHash,
) -> bool {
fn test_serde_info_hash(info_hash: InfoHash) -> bool {
let serialized = bendy::serde::to_bytes(&info_hash).unwrap();
let deserialized: InfoHash =
::bendy::serde::from_bytes(&serialized).unwrap();
let deserialized: InfoHash = ::bendy::serde::from_bytes(&serialized).unwrap();
info_hash == deserialized
}

5
aquatic_udp/src/bin/aquatic_udp.rs
View file

@ -1,11 +1,10 @@
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
fn main(){
fn main() {
aquatic_cli_helpers::run_app_with_cli_and_config::<aquatic_udp::config::Config>(
aquatic_udp::APP_NAME,
aquatic_udp::run,
None
None,
)
}

45
aquatic_udp/src/lib/common.rs
View file

@ -1,6 +1,6 @@
use std::net::{SocketAddr, IpAddr, Ipv4Addr, Ipv6Addr};
use std::sync::{Arc, atomic::AtomicUsize};
use std::hash::Hash;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::{atomic::AtomicUsize, Arc};
use hashbrown::HashMap;
use indexmap::IndexMap;
@ -9,56 +9,45 @@ use parking_lot::Mutex;
pub use aquatic_common::ValidUntil;
pub use aquatic_udp_protocol::*;
pub const MAX_PACKET_SIZE: usize = 4096;
pub trait Ip: Hash + PartialEq + Eq + Clone + Copy {
fn ip_addr(self) -> IpAddr;
}
impl Ip for Ipv4Addr {
fn ip_addr(self) -> IpAddr {
IpAddr::V4(self)
}
}
impl Ip for Ipv6Addr {
fn ip_addr(self) -> IpAddr {
IpAddr::V6(self)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ConnectionKey {
pub connection_id: ConnectionId,
pub socket_addr: SocketAddr
pub socket_addr: SocketAddr,
}
pub type ConnectionMap = HashMap<ConnectionKey, ValidUntil>;
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub enum PeerStatus {
Seeding,
Leeching,
Stopped
Stopped,
}
impl PeerStatus {
/// Determine peer status from announce event and number of bytes left.
///
/// Likely, the last branch will be taken most of the time.
#[inline]
pub fn from_event_and_bytes_left(
event: AnnounceEvent,
bytes_left: NumberOfBytes
) -> Self {
pub fn from_event_and_bytes_left(event: AnnounceEvent, bytes_left: NumberOfBytes) -> Self {
if event == AnnounceEvent::Stopped {
Self::Stopped
} else if bytes_left.0 == 0 {
@ -69,45 +58,39 @@ impl PeerStatus {
}
}
#[derive(Clone, Debug)]
pub struct Peer<I: Ip> {
pub ip_address: I,
pub port: Port,
pub status: PeerStatus,
pub valid_until: ValidUntil
pub valid_until: ValidUntil,
}
impl <I: Ip>Peer<I> {
impl<I: Ip> Peer<I> {
#[inline(always)]
pub fn to_response_peer(&self) -> ResponsePeer {
ResponsePeer {
ip_address: self.ip_address.ip_addr(),
port: self.port
port: self.port,
}
}
}
#[derive(PartialEq, Eq, Hash, Clone, Copy)]
pub struct PeerMapKey<I: Ip> {
pub ip: I,
pub peer_id: PeerId
pub peer_id: PeerId,
}
pub type PeerMap<I> = IndexMap<PeerMapKey<I>, Peer<I>>;
pub struct TorrentData<I: Ip> {
pub peers: PeerMap<I>,
pub num_seeders: usize,
pub num_leechers: usize,
}
impl <I: Ip>Default for TorrentData<I> {
impl<I: Ip> Default for TorrentData<I> {
fn default() -> Self {
Self {
peers: IndexMap::new(),
@ -117,17 +100,14 @@ impl <I: Ip>Default for TorrentData<I> {
}
}
pub type TorrentMap<I> = HashMap<InfoHash, TorrentData<I>>;
#[derive(Default)]
pub struct TorrentMaps {
pub ipv4: TorrentMap<Ipv4Addr>,
pub ipv6: TorrentMap<Ipv6Addr>,
}
#[derive(Default)]
pub struct Statistics {
pub requests_received: AtomicUsize,
@ -137,7 +117,6 @@ pub struct Statistics {
pub bytes_sent: AtomicUsize,
}
#[derive(Clone)]
pub struct State {
pub connections: Arc<Mutex<ConnectionMap>>,
@ -145,7 +124,6 @@ pub struct State {
pub statistics: Arc<Statistics>,
}
impl Default for State {
fn default() -> Self {
Self {
@ -156,11 +134,10 @@ impl Default for State {
}
}
#[cfg(test)]
mod tests {
#[test]
fn test_peer_status_from_event_and_bytes_left(){
fn test_peer_status_from_event_and_bytes_left() {
use crate::common::*;
use PeerStatus::*;

21
aquatic_udp/src/lib/config.rs
View file

@ -1,10 +1,9 @@
use std::net::SocketAddr;
use serde::{Serialize, Deserialize};
use serde::{Deserialize, Serialize};
use aquatic_cli_helpers::LogLevel;
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct Config {
@ -24,14 +23,12 @@ pub struct Config {
pub privileges: PrivilegeConfig,
}
impl aquatic_cli_helpers::Config for Config {
fn get_log_level(&self) -> Option<LogLevel> {
Some(self.log_level)
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct NetworkConfig {
@ -55,7 +52,6 @@ pub struct NetworkConfig {
pub poll_event_capacity: usize,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct ProtocolConfig {
@ -67,7 +63,6 @@ pub struct ProtocolConfig {
pub peer_announce_interval: i32,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct HandlerConfig {
@ -77,7 +72,6 @@ pub struct HandlerConfig {
pub channel_recv_timeout_microseconds: u64,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct StatisticsConfig {
@ -85,7 +79,6 @@ pub struct StatisticsConfig {
pub interval: u64,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct CleaningConfig {
@ -97,7 +90,6 @@ pub struct CleaningConfig {
pub max_connection_age: u64,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct PrivilegeConfig {
@ -109,7 +101,6 @@ pub struct PrivilegeConfig {
pub user: String,
}
impl Default for Config {
fn default() -> Self {
Self {
@ -126,7 +117,6 @@ impl Default for Config {
}
}
impl Default for NetworkConfig {
fn default() -> Self {
Self {
@ -137,7 +127,6 @@ impl Default for NetworkConfig {
}
}
impl Default for ProtocolConfig {
fn default() -> Self {
Self {
@ -148,7 +137,6 @@ impl Default for ProtocolConfig {
}
}
impl Default for HandlerConfig {
fn default() -> Self {
Self {
@ -158,16 +146,12 @@ impl Default for HandlerConfig {
}
}
impl Default for StatisticsConfig {
fn default() -> Self {
Self {
interval: 0,
}
Self { interval: 0 }
}
}
impl Default for CleaningConfig {
fn default() -> Self {
Self {
@ -178,7 +162,6 @@ impl Default for CleaningConfig {
}
}
impl Default for PrivilegeConfig {
fn default() -> Self {
Self {

206
aquatic_udp/src/lib/handlers.rs
View file

@ -1,10 +1,13 @@
use std::net::{SocketAddr, IpAddr};
use std::net::{IpAddr, SocketAddr};
use std::time::Duration;
use std::vec::Drain;
use crossbeam_channel::{Sender, Receiver};
use crossbeam_channel::{Receiver, Sender};
use parking_lot::MutexGuard;
use rand::{SeedableRng, Rng, rngs::{SmallRng, StdRng}};
use rand::{
rngs::{SmallRng, StdRng},
Rng, SeedableRng,
};
use aquatic_common::{convert_ipv4_mapped_ipv6, extract_response_peers};
use aquatic_udp_protocol::*;
@ -12,13 +15,12 @@ use aquatic_udp_protocol::*;
use crate::common::*;
use crate::config::Config;
pub fn run_request_worker(
state: State,
config: Config,
request_receiver: Receiver<(Request, SocketAddr)>,
response_sender: Sender<(Response, SocketAddr)>,
){
) {
let mut connect_requests: Vec<(ConnectRequest, SocketAddr)> = Vec::new();
let mut announce_requests: Vec<(AnnounceRequest, SocketAddr)> = Vec::new();
let mut scrape_requests: Vec<(ScrapeRequest, SocketAddr)> = Vec::new();
@ -28,9 +30,7 @@ pub fn run_request_worker(
let mut std_rng = StdRng::from_entropy();
let mut small_rng = SmallRng::from_rng(&mut std_rng).unwrap();
let timeout = Duration::from_micros(
config.handlers.channel_recv_timeout_microseconds
);
let timeout = Duration::from_micros(config.handlers.channel_recv_timeout_microseconds);
loop {
let mut opt_connections = None;
@ -42,48 +42,41 @@ pub fn run_request_worker(
// only if ConnectionMap mutex isn't locked.
for i in 0..config.handlers.max_requests_per_iter {
let (request, src): (Request, SocketAddr) = if i == 0 {
match request_receiver.recv(){
match request_receiver.recv() {
Ok(r) => r,
Err(_) => break, // Really shouldn't happen
}
} else {
match request_receiver.recv_timeout(timeout){
match request_receiver.recv_timeout(timeout) {
Ok(r) => r,
Err(_) => {
if let Some(guard) = state.connections.try_lock(){
if let Some(guard) = state.connections.try_lock() {
opt_connections = Some(guard);
break
break;
} else {
continue
continue;
}
},
}
}
};
match request {
Request::Connect(r) => {
connect_requests.push((r, src))
},
Request::Announce(r) => {
announce_requests.push((r, src))
},
Request::Scrape(r) => {
scrape_requests.push((r, src))
},
Request::Connect(r) => connect_requests.push((r, src)),
Request::Announce(r) => announce_requests.push((r, src)),
Request::Scrape(r) => scrape_requests.push((r, src)),
}
}
let mut connections: MutexGuard<ConnectionMap> = opt_connections.unwrap_or_else(||
state.connections.lock()
);
let mut connections: MutexGuard<ConnectionMap> =
opt_connections.unwrap_or_else(|| state.connections.lock());
handle_connect_requests(
&config,
&mut connections,
&mut std_rng,
connect_requests.drain(..),
&mut responses
&mut responses,
);
announce_requests.retain(|(request, src)| {
@ -92,12 +85,12 @@ pub fn run_request_worker(
socket_addr: *src,
};
if connections.contains_key(&connection_key){
if connections.contains_key(&connection_key) {
true
} else {
let response = ErrorResponse {
transaction_id: request.transaction_id,
message: "Connection invalid or expired".to_string()
message: "Connection invalid or expired".to_string(),
};
responses.push((response.into(), *src));
@ -112,12 +105,12 @@ pub fn run_request_worker(
socket_addr: *src,
};
if connections.contains_key(&connection_key){
if connections.contains_key(&connection_key) {
true
} else {
let response = ErrorResponse {
transaction_id: request.transaction_id,
message: "Connection invalid or expired".to_string()
message: "Connection invalid or expired".to_string(),
};
responses.push((response.into(), *src));
@ -128,32 +121,27 @@ pub fn run_request_worker(
::std::mem::drop(connections);
if !(announce_requests.is_empty() && scrape_requests.is_empty()){
let mut torrents= state.torrents.lock();
if !(announce_requests.is_empty() && scrape_requests.is_empty()) {
let mut torrents = state.torrents.lock();
handle_announce_requests(
&config,
&mut torrents,
&mut small_rng,
announce_requests.drain(..),
&mut responses
);
handle_scrape_requests(
&mut torrents,
scrape_requests.drain(..),
&mut responses
&mut responses,
);
handle_scrape_requests(&mut torrents, scrape_requests.drain(..), &mut responses);
}
for r in responses.drain(..){
if let Err(err) = response_sender.send(r){
for r in responses.drain(..) {
if let Err(err) = response_sender.send(r) {
::log::error!("error sending response to channel: {}", err);
}
}
}
}
#[inline]
pub fn handle_connect_requests(
config: &Config,
@ -161,7 +149,7 @@ pub fn handle_connect_requests(
rng: &mut StdRng,
requests: Drain<(ConnectRequest, SocketAddr)>,
responses: &mut Vec<(Response, SocketAddr)>,
){
) {
let valid_until = ValidUntil::new(config.cleaning.max_connection_age);
responses.extend(requests.map(|(request, src)| {
@ -174,18 +162,15 @@ pub fn handle_connect_requests(
connections.insert(key, valid_until);
let response = Response::Connect(
ConnectResponse {
connection_id,
transaction_id: request.transaction_id,
}
);
let response = Response::Connect(ConnectResponse {
connection_id,
transaction_id: request.transaction_id,
});
(response, src)
}));
}
#[inline]
pub fn handle_announce_requests(
config: &Config,
@ -193,40 +178,35 @@ pub fn handle_announce_requests(
rng: &mut SmallRng,
requests: Drain<(AnnounceRequest, SocketAddr)>,
responses: &mut Vec<(Response, SocketAddr)>,
){
) {
let peer_valid_until = ValidUntil::new(config.cleaning.max_peer_age);
responses.extend(requests.map(|(request, src)| {
let peer_ip = convert_ipv4_mapped_ipv6(src.ip());
let response = match peer_ip {
IpAddr::V4(ip) => {
handle_announce_request(
config,
rng,
&mut torrents.ipv4,
request,
ip,
peer_valid_until,
)
},
IpAddr::V6(ip) => {
handle_announce_request(
config,
rng,
&mut torrents.ipv6,
request,
ip,
peer_valid_until,
)
}
IpAddr::V4(ip) => handle_announce_request(
config,
rng,
&mut torrents.ipv4,
request,
ip,
peer_valid_until,
),
IpAddr::V6(ip) => handle_announce_request(
config,
rng,
&mut torrents.ipv6,
request,
ip,
peer_valid_until,
),
};
(response.into(), src)
}));
}
fn handle_announce_request<I: Ip>(
config: &Config,
rng: &mut SmallRng,
@ -240,10 +220,7 @@ fn handle_announce_request<I: Ip>(
peer_id: request.peer_id,
};
let peer_status = PeerStatus::from_event_and_bytes_left(
request.event,
request.bytes_left
);
let peer_status = PeerStatus::from_event_and_bytes_left(request.event, request.bytes_left);
let peer = Peer {
ip_address: peer_ip,
@ -252,47 +229,40 @@ fn handle_announce_request<I: Ip>(
valid_until: peer_valid_until,
};
let torrent_data = torrents
.entry(request.info_hash)
.or_default();
let torrent_data = torrents.entry(request.info_hash).or_default();
let opt_removed_peer = match peer_status {
PeerStatus::Leeching => {
torrent_data.num_leechers += 1;
torrent_data.peers.insert(peer_key, peer)
},
}
PeerStatus::Seeding => {
torrent_data.num_seeders += 1;
torrent_data.peers.insert(peer_key, peer)
},
PeerStatus::Stopped => {
torrent_data.peers.remove(&peer_key)
}
PeerStatus::Stopped => torrent_data.peers.remove(&peer_key),
};
match opt_removed_peer.map(|peer| peer.status){
match opt_removed_peer.map(|peer| peer.status) {
Some(PeerStatus::Leeching) => {
torrent_data.num_leechers -= 1;
},
}
Some(PeerStatus::Seeding) => {
torrent_data.num_seeders -= 1;
},
}
_ => {}
}
let max_num_peers_to_take = calc_max_num_peers_to_take(
config,
request.peers_wanted.0
);
let max_num_peers_to_take = calc_max_num_peers_to_take(config, request.peers_wanted.0);
let response_peers = extract_response_peers(
rng,
&torrent_data.peers,
max_num_peers_to_take,
peer_key,
Peer::to_response_peer
Peer::to_response_peer,
);
AnnounceResponse {
@ -300,29 +270,26 @@ fn handle_announce_request<I: Ip>(
announce_interval: AnnounceInterval(config.protocol.peer_announce_interval),
leechers: NumberOfPeers(torrent_data.num_leechers as i32),
seeders: NumberOfPeers(torrent_data.num_seeders as i32),
peers: response_peers
peers: response_peers,
}
}
#[inline]
pub fn handle_scrape_requests(
torrents: &mut MutexGuard<TorrentMaps>,
requests: Drain<(ScrapeRequest, SocketAddr)>,
responses: &mut Vec<(Response, SocketAddr)>,
){
) {
let empty_stats = create_torrent_scrape_statistics(0, 0);
responses.extend(requests.map(|(request, src)|{
let mut stats: Vec<TorrentScrapeStatistics> = Vec::with_capacity(
request.info_hashes.len()
);
responses.extend(requests.map(|(request, src)| {
let mut stats: Vec<TorrentScrapeStatistics> = Vec::with_capacity(request.info_hashes.len());
let peer_ip = convert_ipv4_mapped_ipv6(src.ip());
if peer_ip.is_ipv4(){
if peer_ip.is_ipv4() {
for info_hash in request.info_hashes.iter() {
if let Some(torrent_data) = torrents.ipv4.get(info_hash){
if let Some(torrent_data) = torrents.ipv4.get(info_hash) {
stats.push(create_torrent_scrape_statistics(
torrent_data.num_seeders as i32,
torrent_data.num_leechers as i32,
@ -333,7 +300,7 @@ pub fn handle_scrape_requests(
}
} else {
for info_hash in request.info_hashes.iter() {
if let Some(torrent_data) = torrents.ipv6.get(info_hash){
if let Some(torrent_data) = torrents.ipv6.get(info_hash) {
stats.push(create_torrent_scrape_statistics(
torrent_data.num_seeders as i32,
torrent_data.num_leechers as i32,
@ -353,44 +320,35 @@ pub fn handle_scrape_requests(
}));
}
#[inline]
fn calc_max_num_peers_to_take(
config: &Config,
peers_wanted: i32,
) -> usize {
fn calc_max_num_peers_to_take(config: &Config, peers_wanted: i32) -> usize {
if peers_wanted <= 0 {
config.protocol.max_response_peers as usize
} else {
::std::cmp::min(
config.protocol.max_response_peers as usize,
peers_wanted as usize
peers_wanted as usize,
)
}
}
#[inline(always)]
pub fn create_torrent_scrape_statistics(
seeders: i32,
leechers: i32
) -> TorrentScrapeStatistics {
pub fn create_torrent_scrape_statistics(seeders: i32, leechers: i32) -> TorrentScrapeStatistics {
TorrentScrapeStatistics {
seeders: NumberOfPeers(seeders),
completed: NumberOfDownloads(0), // No implementation planned
leechers: NumberOfPeers(leechers)
leechers: NumberOfPeers(leechers),
}
}
#[cfg(test)]
mod tests {
use std::net::Ipv4Addr;
use std::collections::HashSet;
use std::net::Ipv4Addr;
use indexmap::IndexMap;
use quickcheck::{quickcheck, TestResult};
use rand::thread_rng;
use quickcheck::{TestResult, quickcheck};
use super::*;
@ -413,14 +371,12 @@ mod tests {
}
#[test]
fn test_extract_response_peers(){
fn test_extract_response_peers() {
fn prop(data: (u16, u16)) -> TestResult {
let gen_num_peers = data.0 as u32;
let req_num_peers = data.1 as usize;
let mut peer_map: PeerMap<Ipv4Addr> = IndexMap::with_capacity(
gen_num_peers as usize
);
let mut peer_map: PeerMap<Ipv4Addr> = IndexMap::with_capacity(gen_num_peers as usize);
let mut opt_sender_key = None;
let mut opt_sender_peer = None;
@ -443,7 +399,7 @@ mod tests {
&peer_map,
req_num_peers,
opt_sender_key.unwrap_or_else(|| gen_peer_map_key_and_value(1).0),
Peer::to_response_peer
Peer::to_response_peer,
);
// Check that number of returned peers is correct
@ -451,8 +407,8 @@ mod tests {
let mut success = peers.len() <= req_num_peers;
if req_num_peers >= gen_num_peers as usize {
success &= peers.len() == gen_num_peers as usize ||
peers.len() + 1 == gen_num_peers as usize;
success &= peers.len() == gen_num_peers as usize
|| peers.len() + 1 == gen_num_peers as usize;
}
// Check that returned peers are unique (no overlap) and that sender
@ -461,7 +417,9 @@ mod tests {
let mut ip_addresses = HashSet::with_capacity(peers.len());
for peer in peers {
if peer == opt_sender_peer.clone().unwrap() || ip_addresses.contains(&peer.ip_address){
if peer == opt_sender_peer.clone().unwrap()
|| ip_addresses.contains(&peer.ip_address)
{
success = false;
break;

67
aquatic_udp/src/lib/lib.rs
View file

@ -1,6 +1,9 @@
use std::sync::{Arc, atomic::{AtomicUsize, Ordering}};
use std::time::Duration;
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
use std::thread::Builder;
use std::time::Duration;
use anyhow::Context;
use crossbeam_channel::unbounded;
@ -12,13 +15,11 @@ pub mod handlers;
pub mod network;
pub mod tasks;
use config::Config;
use common::State;
use config::Config;
pub const APP_NAME: &str = "aquatic_udp: UDP BitTorrent tracker";
pub fn run(config: Config) -> ::anyhow::Result<()> {
let state = State::default();
@ -48,11 +49,7 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
}
}
pub fn start_workers(
config: Config,
state: State
) -> ::anyhow::Result<Arc<AtomicUsize>> {
pub fn start_workers(config: Config, state: State) -> ::anyhow::Result<Arc<AtomicUsize>> {
let (request_sender, request_receiver) = unbounded();
let (response_sender, response_receiver) = unbounded();
@ -62,14 +59,12 @@ pub fn start_workers(
let request_receiver = request_receiver.clone();
let response_sender = response_sender.clone();
Builder::new().name(format!("request-{:02}", i + 1)).spawn(move ||
handlers::run_request_worker(
state,
config,
request_receiver,
response_sender
)
).with_context(|| "spawn request worker")?;
Builder::new()
.name(format!("request-{:02}", i + 1))
.spawn(move || {
handlers::run_request_worker(state, config, request_receiver, response_sender)
})
.with_context(|| "spawn request worker")?;
}
let num_bound_sockets = Arc::new(AtomicUsize::new(0));
@ -81,31 +76,33 @@ pub fn start_workers(
let response_receiver = response_receiver.clone();
let num_bound_sockets = num_bound_sockets.clone();
Builder::new().name(format!("socket-{:02}", i + 1)).spawn(move ||
network::run_socket_worker(
state,
config,
i,
request_sender,
response_receiver,
num_bound_sockets,
)
).with_context(|| "spawn socket worker")?;
Builder::new()
.name(format!("socket-{:02}", i + 1))
.spawn(move || {
network::run_socket_worker(
state,
config,
i,
request_sender,
response_receiver,
num_bound_sockets,
)
})
.with_context(|| "spawn socket worker")?;
}
if config.statistics.interval != 0 {
let state = state.clone();
let config = config.clone();
Builder::new().name("statistics-collector".to_string()).spawn(move ||
loop {
::std::thread::sleep(Duration::from_secs(
config.statistics.interval
));
Builder::new()
.name("statistics-collector".to_string())
.spawn(move || loop {
::std::thread::sleep(Duration::from_secs(config.statistics.interval));
tasks::gather_and_print_statistics(&state, &config);
}
).with_context(|| "spawn statistics worker")?;
})
.with_context(|| "spawn statistics worker")?;
}
Ok(num_bound_sockets)

102
aquatic_udp/src/lib/network.rs
View file

@ -1,20 +1,22 @@
use std::sync::{Arc, atomic::{AtomicUsize, Ordering}};
use std::io::{Cursor, ErrorKind};
use std::net::{SocketAddr, IpAddr};
use std::net::{IpAddr, SocketAddr};
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
use std::time::Duration;
use std::vec::Drain;
use crossbeam_channel::{Sender, Receiver};
use mio::{Events, Poll, Interest, Token};
use crossbeam_channel::{Receiver, Sender};
use mio::net::UdpSocket;
use socket2::{Socket, Domain, Type, Protocol};
use mio::{Events, Interest, Poll, Token};
use socket2::{Domain, Protocol, Socket, Type};
use aquatic_udp_protocol::{Request, Response, IpVersion};
use aquatic_udp_protocol::{IpVersion, Request, Response};
use crate::common::*;
use crate::config::Config;
pub fn run_socket_worker(
state: State,
config: Config,
@ -22,7 +24,7 @@ pub fn run_socket_worker(
request_sender: Sender<(Request, SocketAddr)>,
response_receiver: Receiver<(Response, SocketAddr)>,
num_bound_sockets: Arc<AtomicUsize>,
){
) {
let mut buffer = [0u8; MAX_PACKET_SIZE];
let mut socket = UdpSocket::from_std(create_socket(&config));
@ -47,10 +49,10 @@ pub fn run_socket_worker(
poll.poll(&mut events, Some(timeout))
.expect("failed polling");
for event in events.iter(){
for event in events.iter() {
let token = event.token();
if (token.0 == token_num) & event.is_readable(){
if (token.0 == token_num) & event.is_readable() {
read_requests(
&state,
&config,
@ -60,13 +62,16 @@ pub fn run_socket_worker(
&mut local_responses,
);
for r in requests.drain(..){
if let Err(err) = request_sender.send(r){
for r in requests.drain(..) {
if let Err(err) = request_sender.send(r) {
::log::error!("error sending to request_sender: {}", err);
}
}
state.statistics.readable_events.fetch_add(1, Ordering::SeqCst);
state
.statistics
.readable_events
.fetch_add(1, Ordering::SeqCst);
}
}
@ -76,33 +81,33 @@ pub fn run_socket_worker(
&mut socket,
&mut buffer,
&response_receiver,
local_responses.drain(..)
local_responses.drain(..),
);
}
}
fn create_socket(config: &Config) -> ::std::net::UdpSocket {
let socket = if config.network.address.is_ipv4(){
let socket = if config.network.address.is_ipv4() {
Socket::new(Domain::ipv4(), Type::dgram(), Some(Protocol::udp()))
} else {
Socket::new(Domain::ipv6(), Type::dgram(), Some(Protocol::udp()))
}.expect("create socket");
}
.expect("create socket");
socket.set_reuse_port(true)
.expect("socket: set reuse port");
socket.set_reuse_port(true).expect("socket: set reuse port");
socket.set_nonblocking(true)
socket
.set_nonblocking(true)
.expect("socket: set nonblocking");
socket.bind(&config.network.address.into()).unwrap_or_else(|err|
panic!("socket: bind to {}: {:?}", config.network.address, err)
);
socket
.bind(&config.network.address.into())
.unwrap_or_else(|err| panic!("socket: bind to {}: {:?}", config.network.address, err));
let recv_buffer_size = config.network.socket_recv_buffer_size;
if recv_buffer_size != 0 {
if let Err(err) = socket.set_recv_buffer_size(recv_buffer_size){
if let Err(err) = socket.set_recv_buffer_size(recv_buffer_size) {
::log::error!(
"socket: failed setting recv buffer to {}: {:?}",
recv_buffer_size,
@ -114,7 +119,6 @@ fn create_socket(config: &Config) -> ::std::net::UdpSocket {
socket.into_udp_socket()
}
#[inline]
fn read_requests(
state: &State,
@ -123,28 +127,26 @@ fn read_requests(
buffer: &mut [u8],
requests: &mut Vec<(Request, SocketAddr)>,
local_responses: &mut Vec<(Response, SocketAddr)>,
){
) {
let mut requests_received: usize = 0;
let mut bytes_received: usize = 0;
loop {
match socket.recv_from(&mut buffer[..]) {
Ok((amt, src)) => {
let request = Request::from_bytes(
&buffer[..amt],
config.protocol.max_scrape_torrents
);
let request =
Request::from_bytes(&buffer[..amt], config.protocol.max_scrape_torrents);
bytes_received += amt;
if request.is_ok(){
if request.is_ok() {
requests_received += 1;
}
match request {
Ok(request) => {
requests.push((request, src));
},
}
Err(err) => {
::log::debug!("request_from_bytes error: {:?}", err);
@ -166,9 +168,9 @@ fn read_requests(
local_responses.push((response.into(), src));
}
}
},
}
}
},
}
Err(err) => {
if err.kind() == ErrorKind::WouldBlock {
break;
@ -180,14 +182,17 @@ fn read_requests(
}
if config.statistics.interval != 0 {
state.statistics.requests_received
state
.statistics
.requests_received
.fetch_add(requests_received, Ordering::SeqCst);
state.statistics.bytes_received
state
.statistics
.bytes_received
.fetch_add(bytes_received, Ordering::SeqCst);
}
}
#[inline]
fn send_responses(
state: &State,
@ -196,15 +201,15 @@ fn send_responses(
buffer: &mut [u8],
response_receiver: &Receiver<(Response, SocketAddr)>,
local_responses: Drain<(Response, SocketAddr)>,
){
) {
let mut responses_sent: usize = 0;
let mut bytes_sent: usize = 0;
let mut cursor = Cursor::new(buffer);
let response_iterator = local_responses.into_iter().chain(
response_receiver.try_iter()
);
let response_iterator = local_responses
.into_iter()
.chain(response_receiver.try_iter());
for (response, src) in response_iterator {
cursor.set_position(0);
@ -215,11 +220,11 @@ fn send_responses(
let amt = cursor.position() as usize;
match socket.send_to(&cursor.get_ref()[..amt], src){
match socket.send_to(&cursor.get_ref()[..amt], src) {
Ok(amt) => {
responses_sent += 1;
bytes_sent += amt;
},
}
Err(err) => {
if err.kind() == ErrorKind::WouldBlock {
break;
@ -231,19 +236,22 @@ fn send_responses(
}
if config.statistics.interval != 0 {
state.statistics.responses_sent
state
.statistics
.responses_sent
.fetch_add(responses_sent, Ordering::SeqCst);
state.statistics.bytes_sent
state
.statistics
.bytes_sent
.fetch_add(bytes_sent, Ordering::SeqCst);
}
}
fn ip_version_from_ip(ip: IpAddr) -> IpVersion {
match ip {
IpAddr::V4(_) => IpVersion::IPv4,
IpAddr::V6(ip) => {
if let [0, 0, 0, 0, 0, 0xffff, ..] = ip.segments(){
if let [0, 0, 0, 0, 0, 0xffff, ..] = ip.segments() {
IpVersion::IPv4
} else {
IpVersion::IPv6

50
aquatic_udp/src/lib/tasks.rs
View file

@ -6,8 +6,7 @@ use histogram::Histogram;
use crate::common::*;
use crate::config::Config;
pub fn clean_connections_and_torrents(state: &State){
pub fn clean_connections_and_torrents(state: &State) {
let now = Instant::now();
{
@ -23,12 +22,8 @@ pub fn clean_connections_and_torrents(state: &State){
clean_torrent_map(&mut torrents.ipv6, now);
}
#[inline]
fn clean_torrent_map<I: Ip>(
torrents: &mut TorrentMap<I>,
now: Instant,
){
fn clean_torrent_map<I: Ip>(torrents: &mut TorrentMap<I>, now: Instant) {
torrents.retain(|_, torrent| {
let num_seeders = &mut torrent.num_seeders;
let num_leechers = &mut torrent.num_leechers;
@ -40,10 +35,10 @@ fn clean_torrent_map<I: Ip>(
match peer.status {
PeerStatus::Seeding => {
*num_seeders -= 1;
},
}
PeerStatus::Leeching => {
*num_leechers -= 1;
},
}
_ => (),
};
}
@ -57,28 +52,31 @@ fn clean_torrent_map<I: Ip>(
torrents.shrink_to_fit();
}
pub fn gather_and_print_statistics(
state: &State,
config: &Config,
){
pub fn gather_and_print_statistics(state: &State, config: &Config) {
let interval = config.statistics.interval;
let requests_received: f64 = state.statistics.requests_received
let requests_received: f64 = state
.statistics
.requests_received
.fetch_and(0, Ordering::SeqCst) as f64;
let responses_sent: f64 = state.statistics.responses_sent
let responses_sent: f64 = state
.statistics
.responses_sent
.fetch_and(0, Ordering::SeqCst) as f64;
let bytes_received: f64 = state.statistics.bytes_received
.fetch_and(0, Ordering::SeqCst) as f64;
let bytes_sent: f64 = state.statistics.bytes_sent
let bytes_received: f64 = state
.statistics
.bytes_received
.fetch_and(0, Ordering::SeqCst) as f64;
let bytes_sent: f64 = state.statistics.bytes_sent.fetch_and(0, Ordering::SeqCst) as f64;
let requests_per_second = requests_received / interval as f64;
let responses_per_second: f64 = responses_sent / interval as f64;
let bytes_received_per_second: f64 = bytes_received / interval as f64;
let bytes_sent_per_second: f64 = bytes_sent / interval as f64;
let readable_events: f64 = state.statistics.readable_events
let readable_events: f64 = state
.statistics
.readable_events
.fetch_and(0, Ordering::SeqCst) as f64;
let requests_per_readable_event = if readable_events == 0.0 {
0.0
@ -88,9 +86,7 @@ pub fn gather_and_print_statistics(
println!(
"stats: {:.2} requests/second, {:.2} responses/second, {:.2} requests/readable event",
requests_per_second,
responses_per_second,
requests_per_readable_event
requests_per_second, responses_per_second, requests_per_readable_event
);
println!(
@ -104,17 +100,17 @@ pub fn gather_and_print_statistics(
{
let torrents = &mut state.torrents.lock();
for torrent in torrents.ipv4.values(){
for torrent in torrents.ipv4.values() {
let num_peers = (torrent.num_seeders + torrent.num_leechers) as u64;
if let Err(err) = peers_per_torrent.increment(num_peers){
if let Err(err) = peers_per_torrent.increment(num_peers) {
::log::error!("error incrementing peers_per_torrent histogram: {}", err)
}
}
for torrent in torrents.ipv6.values(){
for torrent in torrents.ipv6.values() {
let num_peers = (torrent.num_seeders + torrent.num_leechers) as u64;
if let Err(err) = peers_per_torrent.increment(num_peers){
if let Err(err) = peers_per_torrent.increment(num_peers) {
::log::error!("error incrementing peers_per_torrent histogram: {}", err)
}
}

28
aquatic_udp_bench/src/announce.rs
View file

@ -1,7 +1,7 @@
use std::net::SocketAddr;
use std::time::{Duration, Instant};
use crossbeam_channel::{Sender, Receiver};
use crossbeam_channel::{Receiver, Sender};
use indicatif::ProgressIterator;
use rand::Rng;
use rand_distr::Pareto;
@ -12,7 +12,6 @@ use aquatic_udp::config::Config;
use crate::common::*;
use crate::config::BenchConfig;
pub fn bench_announce_handler(
state: &State,
bench_config: &BenchConfig,
@ -22,12 +21,7 @@ pub fn bench_announce_handler(
rng: &mut impl Rng,
info_hashes: &[InfoHash],
) -> (usize, Duration) {
let requests = create_requests(
state,
rng,
info_hashes,
bench_config.num_announce_requests
);
let requests = create_requests(state, rng, info_hashes, bench_config.num_announce_requests);
let p = aquatic_config.handlers.max_requests_per_iter * bench_config.num_threads;
let mut num_responses = 0usize;
@ -40,8 +34,8 @@ pub fn bench_announce_handler(
let before = Instant::now();
for round in (0..bench_config.num_rounds).progress_with(pb){
for request_chunk in requests.chunks(p){
for round in (0..bench_config.num_rounds).progress_with(pb) {
for request_chunk in requests.chunks(p) {
for (request, src) in request_chunk {
request_sender.send((request.clone().into(), *src)).unwrap();
}
@ -49,7 +43,7 @@ pub fn bench_announce_handler(
while let Ok((Response::Announce(r), _)) = response_receiver.try_recv() {
num_responses += 1;
if let Some(last_peer) = r.peers.last(){
if let Some(last_peer) = r.peers.last() {
dummy ^= last_peer.port.0;
}
}
@ -61,7 +55,7 @@ pub fn bench_announce_handler(
if let Ok((Response::Announce(r), _)) = response_receiver.recv() {
num_responses += 1;
if let Some(last_peer) = r.peers.last(){
if let Some(last_peer) = r.peers.last() {
dummy ^= last_peer.port.0;
}
}
@ -77,7 +71,6 @@ pub fn bench_announce_handler(
(num_responses, elapsed)
}
pub fn create_requests(
state: &State,
rng: &mut impl Rng,
@ -92,12 +85,9 @@ pub fn create_requests(
let connections = state.connections.lock();
let connection_keys: Vec<ConnectionKey> = connections.keys()
.take(number)
.cloned()
.collect();
let connection_keys: Vec<ConnectionKey> = connections.keys().take(number).cloned().collect();
for connection_key in connection_keys.into_iter(){
for connection_key in connection_keys.into_iter() {
let info_hash_index = pareto_usize(rng, pareto, max_index);
let request = AnnounceRequest {
@ -112,7 +102,7 @@ pub fn create_requests(
ip_address: None,
key: PeerKey(rng.gen()),
peers_wanted: NumberOfPeers(rng.gen()),
port: Port(rng.gen())
port: Port(rng.gen()),
};
requests.push((request, connection_key.socket_addr));

9
aquatic_udp_bench/src/common.rs
View file

@ -2,11 +2,9 @@ use indicatif::{ProgressBar, ProgressStyle};
use rand::Rng;
use rand_distr::Pareto;
pub const PARETO_SHAPE: f64 = 0.1;
pub const NUM_INFO_HASHES: usize = 10_000;
pub fn create_progress_bar(name: &str, iterations: u64) -> ProgressBar {
let t = format!("{:<8} {}", name, "{wide_bar} {pos:>2}/{len:>2}");
let style = ProgressStyle::default_bar().template(&t);
@ -14,12 +12,7 @@ pub fn create_progress_bar(name: &str, iterations: u64) -> ProgressBar {
ProgressBar::new(iterations).with_style(style)
}
pub fn pareto_usize(
rng: &mut impl Rng,
pareto: Pareto<f64>,
max: usize,
) -> usize {
pub fn pareto_usize(rng: &mut impl Rng, pareto: Pareto<f64>, max: usize) -> usize {
let p: f64 = rng.sample(pareto);
let p = (p.min(101.0f64) - 1.0) / 100.0;

5
aquatic_udp_bench/src/config.rs
View file

@ -1,5 +1,4 @@
use serde::{Serialize, Deserialize};
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct BenchConfig {
@ -11,7 +10,6 @@ pub struct BenchConfig {
pub num_hashes_per_scrape_request: usize,
}
impl Default for BenchConfig {
fn default() -> Self {
Self {
@ -25,5 +23,4 @@ impl Default for BenchConfig {
}
}
impl aquatic_cli_helpers::Config for BenchConfig {}

16
aquatic_udp_bench/src/connect.rs
View file

@ -1,8 +1,8 @@
use std::time::{Duration, Instant};
use crossbeam_channel::{Sender, Receiver};
use crossbeam_channel::{Receiver, Sender};
use indicatif::ProgressIterator;
use rand::{Rng, SeedableRng, thread_rng, rngs::SmallRng};
use rand::{rngs::SmallRng, thread_rng, Rng, SeedableRng};
use std::net::SocketAddr;
use aquatic_udp::common::*;
@ -11,16 +11,13 @@ use aquatic_udp::config::Config;
use crate::common::*;
use crate::config::BenchConfig;
pub fn bench_connect_handler(
bench_config: &BenchConfig,
aquatic_config: &Config,
request_sender: &Sender<(Request, SocketAddr)>,
response_receiver: &Receiver<(Response, SocketAddr)>,
) -> (usize, Duration) {
let requests = create_requests(
bench_config.num_connect_requests
);
let requests = create_requests(bench_config.num_connect_requests);
let p = aquatic_config.handlers.max_requests_per_iter * bench_config.num_threads;
let mut num_responses = 0usize;
@ -33,8 +30,8 @@ pub fn bench_connect_handler(
let before = Instant::now();
for round in (0..bench_config.num_rounds).progress_with(pb){
for request_chunk in requests.chunks(p){
for round in (0..bench_config.num_rounds).progress_with(pb) {
for request_chunk in requests.chunks(p) {
for (request, src) in request_chunk {
request_sender.send((request.clone().into(), *src)).unwrap();
}
@ -48,7 +45,7 @@ pub fn bench_connect_handler(
let total = bench_config.num_connect_requests * (round + 1);
while num_responses < total {
if let Ok((Response::Connect(r), _)) = response_receiver.recv(){
if let Ok((Response::Connect(r), _)) = response_receiver.recv() {
num_responses += 1;
dummy ^= r.connection_id.0;
}
@ -64,7 +61,6 @@ pub fn bench_connect_handler(
(num_responses, elapsed)
}
pub fn create_requests(number: usize) -> Vec<(ConnectRequest, SocketAddr)> {
let mut rng = SmallRng::from_rng(thread_rng()).unwrap();

45
aquatic_udp_bench/src/main.rs
View file

@ -15,14 +15,14 @@
//! ```
use crossbeam_channel::unbounded;
use std::time::Duration;
use num_format::{Locale, ToFormattedString};
use rand::{Rng, thread_rng, rngs::SmallRng, SeedableRng};
use rand::{rngs::SmallRng, thread_rng, Rng, SeedableRng};
use std::time::Duration;
use aquatic_cli_helpers::run_app_with_cli_and_config;
use aquatic_udp::common::*;
use aquatic_udp::config::Config;
use aquatic_udp::handlers;
use aquatic_cli_helpers::run_app_with_cli_and_config;
use config::BenchConfig;
@ -32,20 +32,17 @@ mod config;
mod connect;
mod scrape;
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
fn main(){
fn main() {
run_app_with_cli_and_config::<BenchConfig>(
"aquatic_udp_bench: Run aquatic_udp benchmarks",
run,
None
None,
)
}
pub fn run(bench_config: BenchConfig) -> ::anyhow::Result<()> {
// Setup common state, spawn request handlers
@ -62,12 +59,7 @@ pub fn run(bench_config: BenchConfig) -> ::anyhow::Result<()> {
let response_sender = response_sender.clone();
::std::thread::spawn(move || {
handlers::run_request_worker(
state,
config,
request_receiver,
response_sender
)
handlers::run_request_worker(state, config, request_receiver, response_sender)
});
}
@ -90,7 +82,7 @@ pub fn run(bench_config: BenchConfig) -> ::anyhow::Result<()> {
&request_sender,
&response_receiver,
&mut rng,
&info_hashes
&info_hashes,
);
let s = scrape::bench_scrape_handler(
@ -100,13 +92,12 @@ pub fn run(bench_config: BenchConfig) -> ::anyhow::Result<()> {
&request_sender,
&response_receiver,
&mut rng,
&info_hashes
&info_hashes,
);
println!(
"\n# Results over {} rounds with {} threads",
bench_config.num_rounds,
bench_config.num_threads,
bench_config.num_rounds, bench_config.num_threads,
);
print_results("Connect: ", c.0, c.1);
@ -116,28 +107,18 @@ pub fn run(bench_config: BenchConfig) -> ::anyhow::Result<()> {
Ok(())
}
pub fn print_results(
request_type: &str,
num_responses: usize,
duration: Duration,
) {
let per_second = (
(num_responses as f64 / (duration.as_micros() as f64 / 1000000.0)
) as usize).to_formatted_string(&Locale::se);
pub fn print_results(request_type: &str, num_responses: usize, duration: Duration) {
let per_second = ((num_responses as f64 / (duration.as_micros() as f64 / 1000000.0)) as usize)
.to_formatted_string(&Locale::se);
let time_per_request = duration.as_nanos() as f64 / (num_responses as f64);
println!(
"{} {:>10} requests/second, {:>8.2} ns/request",
request_type,
per_second,
time_per_request,
request_type, per_second, time_per_request,
);
}
fn create_info_hashes(rng: &mut impl Rng) -> Vec<InfoHash> {
let mut info_hashes = Vec::new();

22
aquatic_udp_bench/src/scrape.rs
View file

@ -1,7 +1,7 @@
use std::net::SocketAddr;
use std::time::{Duration, Instant};
use crossbeam_channel::{Sender, Receiver};
use crossbeam_channel::{Receiver, Sender};
use indicatif::ProgressIterator;
use rand::Rng;
use rand_distr::Pareto;
@ -12,7 +12,6 @@ use aquatic_udp::config::Config;
use crate::common::*;
use crate::config::BenchConfig;
pub fn bench_scrape_handler(
state: &State,
bench_config: &BenchConfig,
@ -41,8 +40,8 @@ pub fn bench_scrape_handler(
let before = Instant::now();
for round in (0..bench_config.num_rounds).progress_with(pb){
for request_chunk in requests.chunks(p){
for round in (0..bench_config.num_rounds).progress_with(pb) {
for request_chunk in requests.chunks(p) {
for (request, src) in request_chunk {
request_sender.send((request.clone().into(), *src)).unwrap();
}
@ -50,7 +49,7 @@ pub fn bench_scrape_handler(
while let Ok((Response::Scrape(r), _)) = response_receiver.try_recv() {
num_responses += 1;
if let Some(stat) = r.torrent_stats.last(){
if let Some(stat) = r.torrent_stats.last() {
dummy ^= stat.leechers.0;
}
}
@ -59,10 +58,10 @@ pub fn bench_scrape_handler(
let total = bench_config.num_scrape_requests * (round + 1);
while num_responses < total {
if let Ok((Response::Scrape(r), _)) = response_receiver.recv(){
if let Ok((Response::Scrape(r), _)) = response_receiver.recv() {
num_responses += 1;
if let Some(stat) = r.torrent_stats.last(){
if let Some(stat) = r.torrent_stats.last() {
dummy ^= stat.leechers.0;
}
}
@ -78,8 +77,6 @@ pub fn bench_scrape_handler(
(num_responses, elapsed)
}
pub fn create_requests(
state: &State,
rng: &mut impl Rng,
@ -93,14 +90,11 @@ pub fn create_requests(
let connections = state.connections.lock();
let connection_keys: Vec<ConnectionKey> = connections.keys()
.take(number)
.cloned()
.collect();
let connection_keys: Vec<ConnectionKey> = connections.keys().take(number).cloned().collect();
let mut requests = Vec::new();
for connection_key in connection_keys.into_iter(){
for connection_key in connection_keys.into_iter() {
let mut request_info_hashes = Vec::new();
for _ in 0..hashes_per_request {

18
aquatic_udp_load_test/src/common.rs
View file

@ -1,17 +1,15 @@
use std::net::SocketAddr;
use std::sync::{Arc, atomic::AtomicUsize};
use std::sync::{atomic::AtomicUsize, Arc};
use hashbrown::HashMap;
use parking_lot::Mutex;
use serde::{Serialize, Deserialize};
use serde::{Deserialize, Serialize};
use aquatic_udp_protocol::*;
#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
pub struct ThreadId(pub u8);
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct Config {
@ -31,7 +29,6 @@ pub struct Config {
pub handler: HandlerConfig,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct NetworkConfig {
@ -67,7 +64,6 @@ pub struct NetworkConfig {
pub recv_buffer: usize,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct HandlerConfig {
@ -100,7 +96,6 @@ pub struct HandlerConfig {
pub additional_request_factor: f64,
}
impl Default for Config {
fn default() -> Self {
Self {
@ -127,7 +122,6 @@ impl Default for NetworkConfig {
}
}
impl Default for HandlerConfig {
fn default() -> Self {
Self {
@ -145,7 +139,6 @@ impl Default for HandlerConfig {
}
}
#[derive(PartialEq, Eq, Clone)]
pub struct TorrentPeer {
pub info_hash: InfoHash,
@ -155,10 +148,8 @@ pub struct TorrentPeer {
pub port: Port,
}
pub type TorrentPeerMap = HashMap<TransactionId, TorrentPeer>;
#[derive(Default)]
pub struct Statistics {
pub requests: AtomicUsize,
@ -169,7 +160,6 @@ pub struct Statistics {
pub responses_error: AtomicUsize,
}
#[derive(Clone)]
pub struct LoadTestState {
pub torrent_peers: Arc<Mutex<TorrentPeerMap>>,
@ -177,15 +167,13 @@ pub struct LoadTestState {
pub statistics: Arc<Statistics>,
}
#[derive(PartialEq, Eq, Clone, Copy)]
pub enum RequestType {
Announce,
Connect,
Scrape
Scrape,
}
#[derive(Default)]
pub struct SocketWorkerLocalStatistics {
pub requests: usize,

179
aquatic_udp_load_test/src/handler.rs
View file

@ -13,20 +13,17 @@ use aquatic_udp_protocol::*;
use crate::common::*;
use crate::utils::*;
pub fn run_handler_thread(
config: &Config,
state: LoadTestState,
pareto: Pareto<f64>,
request_senders: Vec<Sender<Request>>,
response_receiver: Receiver<(ThreadId, Response)>,
){
) {
let state = &state;
let mut rng1 = SmallRng::from_rng(thread_rng())
.expect("create SmallRng from thread_rng()");
let mut rng2 = SmallRng::from_rng(thread_rng())
.expect("create SmallRng from thread_rng()");
let mut rng1 = SmallRng::from_rng(thread_rng()).expect("create SmallRng from thread_rng()");
let mut rng2 = SmallRng::from_rng(thread_rng()).expect("create SmallRng from thread_rng()");
let timeout = Duration::from_micros(config.handler.channel_timeout);
@ -40,30 +37,30 @@ pub fn run_handler_thread(
// only if ConnectionMap mutex isn't locked.
for i in 0..config.handler.max_responses_per_iter {
let response = if i == 0 {
match response_receiver.recv(){
match response_receiver.recv() {
Ok(r) => r,
Err(_) => break, // Really shouldn't happen
}
} else {
match response_receiver.recv_timeout(timeout){
match response_receiver.recv_timeout(timeout) {
Ok(r) => r,
Err(_) => {
if let Some(guard) = state.torrent_peers.try_lock(){
if let Some(guard) = state.torrent_peers.try_lock() {
opt_torrent_peers = Some(guard);
break
break;
} else {
continue
continue;
}
},
}
}
};
responses.push(response);
}
let mut torrent_peers: MutexGuard<TorrentPeerMap> = opt_torrent_peers
.unwrap_or_else(|| state.torrent_peers.lock());
let mut torrent_peers: MutexGuard<TorrentPeerMap> =
opt_torrent_peers.unwrap_or_else(|| state.torrent_peers.lock());
let requests = process_responses(
&mut rng1,
@ -71,69 +68,60 @@ pub fn run_handler_thread(
&state.info_hashes,
config,
&mut torrent_peers,
responses.drain(..)
responses.drain(..),
);
// Somewhat dubious heuristic for deciding how fast to create
// and send additional requests (requests not having anything
// to do with previously sent requests)
let num_additional_to_send = {
let num_additional_requests = requests.iter()
.map(|v| v.len())
.sum::<usize>() as f64;
let num_additional_requests = requests.iter().map(|v| v.len()).sum::<usize>() as f64;
let num_new_requests_per_socket = num_additional_requests /
config.num_socket_workers as f64;
let num_new_requests_per_socket =
num_additional_requests / config.num_socket_workers as f64;
((num_new_requests_per_socket / 1.2) * config.handler.additional_request_factor) as usize + 10
((num_new_requests_per_socket / 1.2) * config.handler.additional_request_factor)
as usize
+ 10
};
for (channel_index, new_requests) in requests.into_iter().enumerate(){
for (channel_index, new_requests) in requests.into_iter().enumerate() {
let channel = &request_senders[channel_index];
for _ in 0..num_additional_to_send {
let request = create_connect_request(
generate_transaction_id(&mut rng2)
);
let request = create_connect_request(generate_transaction_id(&mut rng2));
channel.send(request)
channel
.send(request)
.expect("send request to channel in handler worker");
}
for request in new_requests.into_iter(){
channel.send(request)
for request in new_requests.into_iter() {
channel
.send(request)
.expect("send request to channel in handler worker");
}
}
}
}
fn process_responses(
rng: &mut impl Rng,
pareto: Pareto<f64>,
info_hashes: &Arc<Vec<InfoHash>>,
config: &Config,
torrent_peers: &mut TorrentPeerMap,
responses: Drain<(ThreadId, Response)>
responses: Drain<(ThreadId, Response)>,
) -> Vec<Vec<Request>> {
let mut new_requests = Vec::with_capacity(
config.num_socket_workers as usize
);
let mut new_requests = Vec::with_capacity(config.num_socket_workers as usize);
for _ in 0..config.num_socket_workers {
new_requests.push(Vec::new());
}
for (socket_thread_id, response) in responses.into_iter() {
let opt_request = process_response(
rng,
pareto,
info_hashes,
&config,
torrent_peers,
response
);
let opt_request =
process_response(rng, pareto, info_hashes, &config, torrent_peers, response);
if let Some(new_request) = opt_request {
new_requests[socket_thread_id.0 as usize].push(new_request);
@ -143,77 +131,63 @@ fn process_responses(
new_requests
}
fn process_response(
rng: &mut impl Rng,
pareto: Pareto<f64>,
info_hashes: &Arc<Vec<InfoHash>>,
config: &Config,
torrent_peers: &mut TorrentPeerMap,
response: Response
response: Response,
) -> Option<Request> {
match response {
Response::Connect(r) => {
// Fetch the torrent peer or create it if is doesn't exists. Update
// the connection id if fetched. Create a request and move the
// torrent peer appropriately.
let torrent_peer = torrent_peers.remove(&r.transaction_id)
let torrent_peer = torrent_peers
.remove(&r.transaction_id)
.map(|mut torrent_peer| {
torrent_peer.connection_id = r.connection_id;
torrent_peer
})
.unwrap_or_else(|| {
create_torrent_peer(
config,
rng,
pareto,
info_hashes,
r.connection_id
)
create_torrent_peer(config, rng, pareto, info_hashes, r.connection_id)
});
let new_transaction_id = generate_transaction_id(rng);
let request = create_random_request(
config,
rng,
info_hashes,
new_transaction_id,
&torrent_peer
);
let request =
create_random_request(config, rng, info_hashes, new_transaction_id, &torrent_peer);
torrent_peers.insert(new_transaction_id, torrent_peer);
Some(request)
},
Response::Announce(r) => {
if_torrent_peer_move_and_create_random_request(
config,
rng,
info_hashes,
torrent_peers,
r.transaction_id
)
},
Response::Scrape(r) => {
if_torrent_peer_move_and_create_random_request(
config,
rng,
info_hashes,
torrent_peers,
r.transaction_id
)
},
}
Response::Announce(r) => if_torrent_peer_move_and_create_random_request(
config,
rng,
info_hashes,
torrent_peers,
r.transaction_id,
),
Response::Scrape(r) => if_torrent_peer_move_and_create_random_request(
config,
rng,
info_hashes,
torrent_peers,
r.transaction_id,
),
Response::Error(r) => {
if !r.message.to_lowercase().contains("connection"){
eprintln!("Received error response which didn't contain the word 'connection': {}", r.message);
if !r.message.to_lowercase().contains("connection") {
eprintln!(
"Received error response which didn't contain the word 'connection': {}",
r.message
);
}
if let Some(torrent_peer) = torrent_peers.remove(&r.transaction_id){
if let Some(torrent_peer) = torrent_peers.remove(&r.transaction_id) {
let new_transaction_id = generate_transaction_id(rng);
torrent_peers.insert(new_transaction_id, torrent_peer);
@ -226,7 +200,6 @@ fn process_response(
}
}
fn if_torrent_peer_move_and_create_random_request(
config: &Config,
rng: &mut impl Rng,
@ -234,16 +207,11 @@ fn if_torrent_peer_move_and_create_random_request(
torrent_peers: &mut TorrentPeerMap,
transaction_id: TransactionId,
) -> Option<Request> {
if let Some(torrent_peer) = torrent_peers.remove(&transaction_id){
if let Some(torrent_peer) = torrent_peers.remove(&transaction_id) {
let new_transaction_id = generate_transaction_id(rng);
let request = create_random_request(
config,
rng,
info_hashes,
new_transaction_id,
&torrent_peer
);
let request =
create_random_request(config, rng, info_hashes, new_transaction_id, &torrent_peer);
torrent_peers.insert(new_transaction_id, torrent_peer);
@ -253,13 +221,12 @@ fn if_torrent_peer_move_and_create_random_request(
}
}
fn create_random_request(
config: &Config,
rng: &mut impl Rng,
info_hashes: &Arc<Vec<InfoHash>>,
transaction_id: TransactionId,
torrent_peer: &TorrentPeer
torrent_peer: &TorrentPeer,
) -> Request {
let weights = vec![
config.handler.weight_announce as u32,
@ -273,26 +240,15 @@ fn create_random_request(
RequestType::Scrape,
];
let dist = WeightedIndex::new(&weights)
.expect("random request weighted index");
let dist = WeightedIndex::new(&weights).expect("random request weighted index");
match items[dist.sample(rng)] {
RequestType::Announce => create_announce_request(
config,
rng,
torrent_peer,
transaction_id
),
RequestType::Announce => create_announce_request(config, rng, torrent_peer, transaction_id),
RequestType::Connect => create_connect_request(transaction_id),
RequestType::Scrape => create_scrape_request(
&info_hashes,
torrent_peer,
transaction_id
)
RequestType::Scrape => create_scrape_request(&info_hashes, torrent_peer, transaction_id),
}
}
fn create_announce_request(
config: &Config,
rng: &mut impl Rng,
@ -319,11 +275,11 @@ fn create_announce_request(
ip_address: None,
key: PeerKey(12345),
peers_wanted: NumberOfPeers(100),
port: torrent_peer.port
}).into()
port: torrent_peer.port,
})
.into()
}
fn create_scrape_request(
info_hashes: &Arc<Vec<InfoHash>>,
torrent_peer: &TorrentPeer,
@ -341,5 +297,6 @@ fn create_scrape_request(
connection_id: torrent_peer.connection_id,
transaction_id,
info_hashes: scape_hashes,
}).into()
})
.into()
}

120
aquatic_udp_load_test/src/main.rs
View file

@ -1,6 +1,6 @@
use std::net::{SocketAddr, Ipv4Addr, Ipv6Addr};
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::{atomic::Ordering, Arc};
use std::thread;
use std::sync::{Arc, atomic::Ordering};
use std::time::{Duration, Instant};
use crossbeam_channel::unbounded;
@ -15,16 +15,14 @@ mod network;
mod utils;
use common::*;
use utils::*;
use network::*;
use handler::run_handler_thread;
use network::*;
use utils::*;
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
pub fn main(){
pub fn main() {
aquatic_cli_helpers::run_app_with_cli_and_config::<Config>(
"aquatic_udp_load_test: BitTorrent load tester",
run,
@ -32,12 +30,12 @@ pub fn main(){
)
}
impl aquatic_cli_helpers::Config for Config {}
fn run(config: Config) -> ::anyhow::Result<()> {
if config.handler.weight_announce + config.handler.weight_connect + config.handler.weight_scrape == 0 {
if config.handler.weight_announce + config.handler.weight_connect + config.handler.weight_scrape
== 0
{
panic!("Error: at least one weight must be larger than zero.");
}
@ -55,10 +53,7 @@ fn run(config: Config) -> ::anyhow::Result<()> {
statistics: Arc::new(Statistics::default()),
};
let pareto = Pareto::new(
1.0,
config.handler.torrent_selection_pareto_shape
).unwrap();
let pareto = Pareto::new(1.0, config.handler.torrent_selection_pareto_shape).unwrap();
// Start socket workers
@ -72,7 +67,8 @@ fn run(config: Config) -> ::anyhow::Result<()> {
let port = config.network.first_port + (i as u16);
let addr = if config.network.multiple_client_ips {
let ip = if config.network.ipv6_client { // FIXME: test ipv6
let ip = if config.network.ipv6_client {
// FIXME: test ipv6
Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1 + i as u16).into()
} else {
Ipv4Addr::new(127, 0, 0, 1 + i).into()
@ -95,54 +91,37 @@ fn run(config: Config) -> ::anyhow::Result<()> {
let response_sender = response_sender.clone();
let state = state.clone();
thread::spawn(move || run_socket_thread(
state,
response_sender,
receiver,
&config,
addr,
thread_id
));
thread::spawn(move || {
run_socket_thread(state, response_sender, receiver, &config, addr, thread_id)
});
}
for _ in 0..config.num_request_workers {
let config = config.clone();
let state= state.clone();
let state = state.clone();
let request_senders = request_senders.clone();
let response_receiver = response_receiver.clone();
thread::spawn(move || run_handler_thread(
&config,
state,
pareto,
request_senders,
response_receiver,
));
thread::spawn(move || {
run_handler_thread(&config, state, pareto, request_senders, response_receiver)
});
}
// Bootstrap request cycle by adding a request to each request channel
for sender in request_senders.iter(){
let request = create_connect_request(
generate_transaction_id(&mut thread_rng())
);
for sender in request_senders.iter() {
let request = create_connect_request(generate_transaction_id(&mut thread_rng()));
sender.send(request)
sender
.send(request)
.expect("bootstrap: add initial request to request queue");
}
monitor_statistics(
state,
&config
);
monitor_statistics(state, &config);
Ok(())
}
fn monitor_statistics(
state: LoadTestState,
config: &Config,
){
fn monitor_statistics(state: LoadTestState, config: &Config) {
let start_time = Instant::now();
let mut report_avg_response_vec: Vec<f64> = Vec::new();
@ -154,38 +133,45 @@ fn monitor_statistics(
let statistics = state.statistics.as_ref();
let responses_announce = statistics.responses_announce
.fetch_and(0, Ordering::SeqCst) as f64;
let response_peers = statistics.response_peers
.fetch_and(0, Ordering::SeqCst) as f64;
let responses_announce =
statistics.responses_announce.fetch_and(0, Ordering::SeqCst) as f64;
let response_peers = statistics.response_peers.fetch_and(0, Ordering::SeqCst) as f64;
let requests_per_second = statistics.requests
.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_connect_per_second = statistics.responses_connect
.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_scrape_per_second = statistics.responses_scrape
.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_error_per_second = statistics.responses_error
.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let requests_per_second =
statistics.requests.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_connect_per_second =
statistics.responses_connect.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_scrape_per_second =
statistics.responses_scrape.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_error_per_second =
statistics.responses_error.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_announce_per_second = responses_announce / interval_f64;
let responses_announce_per_second = responses_announce / interval_f64;
let responses_per_second =
responses_connect_per_second +
responses_announce_per_second +
responses_scrape_per_second +
responses_error_per_second;
let responses_per_second = responses_connect_per_second
+ responses_announce_per_second
+ responses_scrape_per_second
+ responses_error_per_second;
report_avg_response_vec.push(responses_per_second);
println!();
println!("Requests out: {:.2}/second", requests_per_second);
println!("Responses in: {:.2}/second", responses_per_second);
println!(" - Connect responses: {:.2}", responses_connect_per_second);
println!(" - Announce responses: {:.2}", responses_announce_per_second);
println!(
" - Connect responses: {:.2}",
responses_connect_per_second
);
println!(
" - Announce responses: {:.2}",
responses_announce_per_second
);
println!(" - Scrape responses: {:.2}", responses_scrape_per_second);
println!(" - Error responses: {:.2}", responses_error_per_second);
println!("Peers per announce response: {:.2}", response_peers / responses_announce);
println!(
"Peers per announce response: {:.2}",
response_peers / responses_announce
);
let time_elapsed = start_time.elapsed();
let duration = Duration::from_secs(config.duration as u64);
@ -206,7 +192,7 @@ fn monitor_statistics(
config
);
break
break;
}
}
}

101
aquatic_udp_load_test/src/network.rs
View file

@ -4,57 +4,54 @@ use std::sync::atomic::Ordering;
use std::time::Duration;
use crossbeam_channel::{Receiver, Sender};
use mio::{net::UdpSocket, Events, Poll, Interest, Token};
use socket2::{Socket, Domain, Type, Protocol};
use mio::{net::UdpSocket, Events, Interest, Poll, Token};
use socket2::{Domain, Protocol, Socket, Type};
use aquatic_udp_protocol::*;
use crate::common::*;
const MAX_PACKET_SIZE: usize = 4096;
pub fn create_socket(
config: &Config,
addr: SocketAddr
) -> ::std::net::UdpSocket {
let socket = if addr.is_ipv4(){
pub fn create_socket(config: &Config, addr: SocketAddr) -> ::std::net::UdpSocket {
let socket = if addr.is_ipv4() {
Socket::new(Domain::ipv4(), Type::dgram(), Some(Protocol::udp()))
} else {
Socket::new(Domain::ipv6(), Type::dgram(), Some(Protocol::udp()))
}.expect("create socket");
}
.expect("create socket");
socket.set_nonblocking(true)
socket
.set_nonblocking(true)
.expect("socket: set nonblocking");
if config.network.recv_buffer != 0 {
if let Err(err) = socket.set_recv_buffer_size(config.network.recv_buffer){
if let Err(err) = socket.set_recv_buffer_size(config.network.recv_buffer) {
eprintln!(
"socket: failed setting recv buffer to {}: {:?}",
config.network.recv_buffer,
err
config.network.recv_buffer, err
);
}
}
socket.bind(&addr.into())
socket
.bind(&addr.into())
.unwrap_or_else(|err| panic!("socket: bind to {}: {:?}", addr, err));
socket.connect(&config.server_address.into())
socket
.connect(&config.server_address.into())
.expect("socket: connect to server");
socket.into_udp_socket()
}
pub fn run_socket_thread(
state: LoadTestState,
response_channel_sender: Sender<(ThreadId, Response)>,
request_receiver: Receiver<Request>,
config: &Config,
addr: SocketAddr,
thread_id: ThreadId
thread_id: ThreadId,
) {
let mut socket = UdpSocket::from_std(create_socket(config, addr));
let mut buffer = [0u8; MAX_PACKET_SIZE];
@ -78,23 +75,23 @@ pub fn run_socket_thread(
poll.poll(&mut events, Some(timeout))
.expect("failed polling");
for event in events.iter(){
if (event.token() == token) & event.is_readable(){
for event in events.iter() {
if (event.token() == token) & event.is_readable() {
read_responses(
thread_id,
&socket,
&mut buffer,
&mut local_state,
&mut responses
&mut responses,
);
for r in responses.drain(..){
response_channel_sender.send(r)
.unwrap_or_else(|err| panic!(
for r in responses.drain(..) {
response_channel_sender.send(r).unwrap_or_else(|err| {
panic!(
"add response to channel in socket worker {}: {:?}",
thread_id.0,
err
));
thread_id.0, err
)
});
}
poll.registry()
@ -107,7 +104,7 @@ pub fn run_socket_thread(
&mut socket,
&mut buffer,
&request_receiver,
&mut local_state
&mut local_state,
);
}
@ -116,40 +113,39 @@ pub fn run_socket_thread(
&mut socket,
&mut buffer,
&request_receiver,
&mut local_state
&mut local_state,
);
}
}
fn read_responses(
thread_id: ThreadId,
socket: &UdpSocket,
buffer: &mut [u8],
ls: &mut SocketWorkerLocalStatistics,
responses: &mut Vec<(ThreadId, Response)>,
){
) {
while let Ok(amt) = socket.recv(buffer) {
match Response::from_bytes(&buffer[0..amt]){
match Response::from_bytes(&buffer[0..amt]) {
Ok(response) => {
match response {
Response::Announce(ref r) => {
ls.responses_announce += 1;
ls.response_peers += r.peers.len();
},
}
Response::Scrape(_) => {
ls.responses_scrape += 1;
},
}
Response::Connect(_) => {
ls.responses_connect += 1;
},
}
Response::Error(_) => {
ls.responses_error += 1;
},
}
}
responses.push((thread_id, response))
},
}
Err(err) => {
eprintln!("Received invalid response: {:#?}", err);
}
@ -157,20 +153,19 @@ fn read_responses(
}
}
fn send_requests(
state: &LoadTestState,
socket: &mut UdpSocket,
buffer: &mut [u8],
receiver: &Receiver<Request>,
statistics: &mut SocketWorkerLocalStatistics,
){
) {
let mut cursor = Cursor::new(buffer);
while let Ok(request) = receiver.try_recv() {
cursor.set_position(0);
if let Err(err) = request.write(&mut cursor){
if let Err(err) = request.write(&mut cursor) {
eprintln!("request_to_bytes err: {}", err);
}
@ -180,24 +175,36 @@ fn send_requests(
match socket.send(&inner[..position]) {
Ok(_) => {
statistics.requests += 1;
},
}
Err(err) => {
eprintln!("Couldn't send packet: {:?}", err);
}
}
}
state.statistics.requests
state
.statistics
.requests
.fetch_add(statistics.requests, Ordering::SeqCst);
state.statistics.responses_connect
state
.statistics
.responses_connect
.fetch_add(statistics.responses_connect, Ordering::SeqCst);
state.statistics.responses_announce
state
.statistics
.responses_announce
.fetch_add(statistics.responses_announce, Ordering::SeqCst);
state.statistics.responses_scrape
state
.statistics
.responses_scrape
.fetch_add(statistics.responses_scrape, Ordering::SeqCst);
state.statistics.responses_error
state
.statistics
.responses_error
.fetch_add(statistics.responses_error, Ordering::SeqCst);
state.statistics.response_peers
state
.statistics
.response_peers
.fetch_add(statistics.response_peers, Ordering::SeqCst);
*statistics = SocketWorkerLocalStatistics::default();

34
aquatic_udp_load_test/src/utils.rs
View file

@ -1,30 +1,25 @@
use std::sync::Arc;
use rand_distr::Pareto;
use rand::prelude::*;
use rand_distr::Pareto;
use aquatic_udp_protocol::*;
use crate::common::*;
pub fn create_torrent_peer(
config: &Config,
rng: &mut impl Rng,
pareto: Pareto<f64>,
info_hashes: &Arc<Vec<InfoHash>>,
connection_id: ConnectionId
connection_id: ConnectionId,
) -> TorrentPeer {
let num_scape_hashes = rng.gen_range(
1..config.handler.scrape_max_torrents
);
let num_scape_hashes = rng.gen_range(1..config.handler.scrape_max_torrents);
let mut scrape_hash_indeces = Vec::new();
for _ in 0..num_scape_hashes {
scrape_hash_indeces.push(
select_info_hash_index(config, rng, pareto)
)
scrape_hash_indeces.push(select_info_hash_index(config, rng, pareto))
}
let info_hash_index = select_info_hash_index(config, rng, pareto);
@ -34,52 +29,37 @@ pub fn create_torrent_peer(
scrape_hash_indeces,
connection_id,
peer_id: generate_peer_id(),
port: Port(rand::random())
port: Port(rand::random()),
}
}
fn select_info_hash_index(
config: &Config,
rng: &mut impl Rng,
pareto: Pareto<f64>,
) -> usize {
fn select_info_hash_index(config: &Config, rng: &mut impl Rng, pareto: Pareto<f64>) -> usize {
pareto_usize(rng, pareto, config.handler.number_of_torrents - 1)
}
pub fn pareto_usize(
rng: &mut impl Rng,
pareto: Pareto<f64>,
max: usize,
) -> usize {
pub fn pareto_usize(rng: &mut impl Rng, pareto: Pareto<f64>, max: usize) -> usize {
let p: f64 = rng.sample(pareto);
let p = (p.min(101.0f64) - 1.0) / 100.0;
(p * max as f64) as usize
}
pub fn generate_peer_id() -> PeerId {
PeerId(random_20_bytes())
}
pub fn generate_info_hash() -> InfoHash {
InfoHash(random_20_bytes())
}
pub fn generate_transaction_id(rng: &mut impl Rng) -> TransactionId {
TransactionId(rng.gen())
}
pub fn create_connect_request(transaction_id: TransactionId) -> Request {
(ConnectRequest { transaction_id }).into()
}
// Don't use SmallRng here for now
fn random_20_bytes() -> [u8; 20] {
let mut bytes = [0; 20];

33
aquatic_udp_protocol/src/common.rs
View file

@ -1,46 +1,40 @@
use std::net::IpAddr;
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub enum IpVersion {
IPv4,
IPv6
IPv6,
}
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub struct AnnounceInterval (pub i32);
pub struct AnnounceInterval(pub i32);
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub struct InfoHash (pub [u8; 20]);
pub struct InfoHash(pub [u8; 20]);
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub struct ConnectionId (pub i64);
pub struct ConnectionId(pub i64);
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub struct TransactionId (pub i32);
pub struct TransactionId(pub i32);
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub struct NumberOfBytes (pub i64);
pub struct NumberOfBytes(pub i64);
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub struct NumberOfPeers (pub i32);
pub struct NumberOfPeers(pub i32);
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub struct NumberOfDownloads (pub i32);
pub struct NumberOfDownloads(pub i32);
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub struct Port (pub u16);
pub struct Port(pub u16);
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug, PartialOrd, Ord)]
pub struct PeerId (pub [u8; 20]);
pub struct PeerId(pub [u8; 20]);
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub struct PeerKey (pub u32);
pub struct PeerKey(pub u32);
#[derive(Hash, PartialEq, Eq, Clone, Debug)]
pub struct ResponsePeer {
@ -48,8 +42,6 @@ pub struct ResponsePeer {
pub port: Port,
}
#[cfg(test)]
impl quickcheck::Arbitrary for IpVersion {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
@ -61,7 +53,6 @@ impl quickcheck::Arbitrary for IpVersion {
}
}
#[cfg(test)]
impl quickcheck::Arbitrary for InfoHash {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
@ -75,7 +66,6 @@ impl quickcheck::Arbitrary for InfoHash {
}
}
#[cfg(test)]
impl quickcheck::Arbitrary for PeerId {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
@ -89,7 +79,6 @@ impl quickcheck::Arbitrary for PeerId {
}
}
#[cfg(test)]
impl quickcheck::Arbitrary for ResponsePeer {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {

133
aquatic_udp_protocol/src/request.rs
View file

@ -2,23 +2,20 @@ use std::convert::TryInto;
use std::io::{self, Cursor, Read, Write};
use std::net::Ipv4Addr;
use byteorder::{ReadBytesExt, WriteBytesExt, NetworkEndian};
use byteorder::{NetworkEndian, ReadBytesExt, WriteBytesExt};
use super::common::*;
const PROTOCOL_IDENTIFIER: i64 = 4_497_486_125_440;
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub enum AnnounceEvent {
Started,
Stopped,
Completed,
None
None,
}
impl AnnounceEvent {
#[inline]
pub fn from_i32(i: i32) -> Self {
@ -26,7 +23,7 @@ impl AnnounceEvent {
1 => Self::Completed,
2 => Self::Started,
3 => Self::Stopped,
_ => Self::None
_ => Self::None,
}
}
@ -36,18 +33,16 @@ impl AnnounceEvent {
AnnounceEvent::None => 0,
AnnounceEvent::Completed => 1,
AnnounceEvent::Started => 2,
AnnounceEvent::Stopped => 3
AnnounceEvent::Stopped => 3,
}
}
}
#[derive(PartialEq, Eq, Clone, Debug)]
pub struct ConnectRequest {
pub transaction_id: TransactionId
pub transaction_id: TransactionId,
}
#[derive(PartialEq, Eq, Clone, Debug)]
pub struct AnnounceRequest {
pub connection_id: ConnectionId,
@ -61,18 +56,16 @@ pub struct AnnounceRequest {
pub ip_address: Option<Ipv4Addr>,
pub key: PeerKey,
pub peers_wanted: NumberOfPeers,
pub port: Port
pub port: Port,
}
#[derive(PartialEq, Eq, Clone, Debug)]
pub struct ScrapeRequest {
pub connection_id: ConnectionId,
pub transaction_id: TransactionId,
pub info_hashes: Vec<InfoHash>
pub info_hashes: Vec<InfoHash>,
}
#[derive(Debug)]
pub struct RequestParseError {
pub transaction_id: Option<TransactionId>,
@ -80,20 +73,19 @@ pub struct RequestParseError {
pub error: Option<io::Error>,
}
impl RequestParseError {
pub fn new(err: io::Error, transaction_id: i32) -> Self {
Self {
transaction_id: Some(TransactionId(transaction_id)),
message: None,
error: Some(err)
error: Some(err),
}
}
pub fn io(err: io::Error) -> Self {
Self {
transaction_id: None,
message: None,
error: Some(err)
error: Some(err),
}
}
pub fn text(transaction_id: i32, message: &str) -> Self {
@ -105,7 +97,6 @@ impl RequestParseError {
}
}
#[derive(PartialEq, Eq, Clone, Debug)]
pub enum Request {
Connect(ConnectRequest),
@ -113,28 +104,24 @@ pub enum Request {
Scrape(ScrapeRequest),
}
impl From<ConnectRequest> for Request {
fn from(r: ConnectRequest) -> Self {
Self::Connect(r)
}
}
impl From<AnnounceRequest> for Request {
fn from(r: AnnounceRequest) -> Self {
Self::Announce(r)
}
}
impl From<ScrapeRequest> for Request {
fn from(r: ScrapeRequest) -> Self {
Self::Scrape(r)
}
}
impl Request {
pub fn write(self, bytes: &mut impl Write) -> Result<(), io::Error> {
match self {
@ -142,7 +129,7 @@ impl Request {
bytes.write_i64::<NetworkEndian>(PROTOCOL_IDENTIFIER)?;
bytes.write_i32::<NetworkEndian>(0)?;
bytes.write_i32::<NetworkEndian>(r.transaction_id.0)?;
},
}
Request::Announce(r) => {
bytes.write_i64::<NetworkEndian>(r.connection_id.0)?;
@ -158,15 +145,12 @@ impl Request {
bytes.write_i32::<NetworkEndian>(r.event.to_i32())?;
bytes.write_all(&r.ip_address.map_or(
[0; 4],
|ip| ip.octets()
))?;
bytes.write_all(&r.ip_address.map_or([0; 4], |ip| ip.octets()))?;
bytes.write_u32::<NetworkEndian>(r.key.0)?;
bytes.write_i32::<NetworkEndian>(r.peers_wanted.0)?;
bytes.write_u16::<NetworkEndian>(r.port.0)?;
},
}
Request::Scrape(r) => {
bytes.write_i64::<NetworkEndian>(r.connection_id.0)?;
@ -182,17 +166,17 @@ impl Request {
Ok(())
}
pub fn from_bytes(
bytes: &[u8],
max_scrape_torrents: u8,
) -> Result<Self, RequestParseError> {
pub fn from_bytes(bytes: &[u8], max_scrape_torrents: u8) -> Result<Self, RequestParseError> {
let mut cursor = Cursor::new(bytes);
let connection_id = cursor.read_i64::<NetworkEndian>()
let connection_id = cursor
.read_i64::<NetworkEndian>()
.map_err(RequestParseError::io)?;
let action = cursor.read_i32::<NetworkEndian>()
let action = cursor
.read_i32::<NetworkEndian>()
.map_err(RequestParseError::io)?;
let transaction_id = cursor.read_i32::<NetworkEndian>()
let transaction_id = cursor
.read_i32::<NetworkEndian>()
.map_err(RequestParseError::io)?;
match action {
@ -200,15 +184,16 @@ impl Request {
0 => {
if connection_id == PROTOCOL_IDENTIFIER {
Ok((ConnectRequest {
transaction_id: TransactionId(transaction_id)
}).into())
transaction_id: TransactionId(transaction_id),
})
.into())
} else {
Err(RequestParseError::text(
transaction_id,
"Protocol identifier missing"
"Protocol identifier missing",
))
}
},
}
// Announce
1 => {
@ -216,28 +201,38 @@ impl Request {
let mut peer_id = [0; 20];
let mut ip = [0; 4];
cursor.read_exact(&mut info_hash)
cursor
.read_exact(&mut info_hash)
.map_err(|err| RequestParseError::new(err, transaction_id))?;
cursor.read_exact(&mut peer_id)
cursor
.read_exact(&mut peer_id)
.map_err(|err| RequestParseError::new(err, transaction_id))?;
let bytes_downloaded = cursor.read_i64::<NetworkEndian>()
let bytes_downloaded = cursor
.read_i64::<NetworkEndian>()
.map_err(|err| RequestParseError::new(err, transaction_id))?;
let bytes_left = cursor.read_i64::<NetworkEndian>()
let bytes_left = cursor
.read_i64::<NetworkEndian>()
.map_err(|err| RequestParseError::new(err, transaction_id))?;
let bytes_uploaded = cursor.read_i64::<NetworkEndian>()
let bytes_uploaded = cursor
.read_i64::<NetworkEndian>()
.map_err(|err| RequestParseError::new(err, transaction_id))?;
let event = cursor.read_i32::<NetworkEndian>()
let event = cursor
.read_i32::<NetworkEndian>()
.map_err(|err| RequestParseError::new(err, transaction_id))?;
cursor.read_exact(&mut ip)
cursor
.read_exact(&mut ip)
.map_err(|err| RequestParseError::new(err, transaction_id))?;
let key = cursor.read_u32::<NetworkEndian>()
let key = cursor
.read_u32::<NetworkEndian>()
.map_err(|err| RequestParseError::new(err, transaction_id))?;
let peers_wanted = cursor.read_i32::<NetworkEndian>()
let peers_wanted = cursor
.read_i32::<NetworkEndian>()
.map_err(|err| RequestParseError::new(err, transaction_id))?;
let port = cursor.read_u16::<NetworkEndian>()
let port = cursor
.read_u16::<NetworkEndian>()
.map_err(|err| RequestParseError::new(err, transaction_id))?;
let opt_ip = if ip == [0; 4] {
@ -258,16 +253,18 @@ impl Request {
ip_address: opt_ip,
key: PeerKey(key),
peers_wanted: NumberOfPeers(peers_wanted),
port: Port(port)
}).into())
},
port: Port(port),
})
.into())
}
// Scrape
2 => {
let position = cursor.position() as usize;
let inner = cursor.into_inner();
let info_hashes = (&inner[position..]).chunks_exact(20)
let info_hashes = (&inner[position..])
.chunks_exact(20)
.take(max_scrape_torrents as usize)
.map(|chunk| InfoHash(chunk.try_into().unwrap()))
.collect();
@ -275,16 +272,16 @@ impl Request {
Ok((ScrapeRequest {
connection_id: ConnectionId(connection_id),
transaction_id: TransactionId(transaction_id),
info_hashes
}).into())
info_hashes,
})
.into())
}
_ => Err(RequestParseError::text(transaction_id, "Invalid action"))
_ => Err(RequestParseError::text(transaction_id, "Invalid action")),
}
}
}
#[cfg(test)]
mod tests {
use quickcheck_macros::quickcheck;
@ -293,7 +290,7 @@ mod tests {
impl quickcheck::Arbitrary for AnnounceEvent {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
match (bool::arbitrary(g), bool::arbitrary(g)){
match (bool::arbitrary(g), bool::arbitrary(g)) {
(false, false) => Self::Started,
(true, false) => Self::Started,
(false, true) => Self::Completed,
@ -324,16 +321,16 @@ mod tests {
ip_address: None,
key: PeerKey(u32::arbitrary(g)),
peers_wanted: NumberOfPeers(i32::arbitrary(g)),
port: Port(u16::arbitrary(g))
port: Port(u16::arbitrary(g)),
}
}
}
impl quickcheck::Arbitrary for ScrapeRequest {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
let info_hashes = (0..u8::arbitrary(g)).map(|_| {
InfoHash::arbitrary(g)
}).collect();
let info_hashes = (0..u8::arbitrary(g))
.map(|_| InfoHash::arbitrary(g))
.collect();
Self {
connection_id: ConnectionId(i64::arbitrary(g)),
@ -359,23 +356,17 @@ mod tests {
}
#[quickcheck]
fn test_connect_request_convert_identity(
request: ConnectRequest
) -> bool {
fn test_connect_request_convert_identity(request: ConnectRequest) -> bool {
same_after_conversion(request.into())
}
#[quickcheck]
fn test_announce_request_convert_identity(
request: AnnounceRequest
) -> bool {
fn test_announce_request_convert_identity(request: AnnounceRequest) -> bool {
same_after_conversion(request.into())
}
#[quickcheck]
fn test_scrape_request_convert_identity(
request: ScrapeRequest
) -> bool {
fn test_scrape_request_convert_identity(request: ScrapeRequest) -> bool {
same_after_conversion(request.into())
}
}

180
aquatic_udp_protocol/src/response.rs
View file

@ -1,51 +1,45 @@
use std::convert::TryInto;
use std::io::{self, Cursor, Write};
use std::net::{IpAddr, Ipv6Addr, Ipv4Addr};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use byteorder::{ReadBytesExt, WriteBytesExt, NetworkEndian};
use byteorder::{NetworkEndian, ReadBytesExt, WriteBytesExt};
use super::common::*;
#[derive(PartialEq, Eq, Debug, Copy, Clone)]
pub struct TorrentScrapeStatistics {
pub seeders: NumberOfPeers,
pub completed: NumberOfDownloads,
pub leechers: NumberOfPeers
pub leechers: NumberOfPeers,
}
#[derive(PartialEq, Eq, Clone, Debug)]
pub struct ConnectResponse {
pub connection_id: ConnectionId,
pub transaction_id: TransactionId
pub transaction_id: TransactionId,
}
#[derive(PartialEq, Eq, Clone, Debug)]
pub struct AnnounceResponse {
pub transaction_id: TransactionId,
pub announce_interval: AnnounceInterval,
pub leechers: NumberOfPeers,
pub seeders: NumberOfPeers,
pub peers: Vec<ResponsePeer>
pub peers: Vec<ResponsePeer>,
}
#[derive(PartialEq, Eq, Clone, Debug)]
pub struct ScrapeResponse {
pub transaction_id: TransactionId,
pub torrent_stats: Vec<TorrentScrapeStatistics>
pub torrent_stats: Vec<TorrentScrapeStatistics>,
}
#[derive(PartialEq, Eq, Clone, Debug)]
pub struct ErrorResponse {
pub transaction_id: TransactionId,
pub message: String
pub message: String,
}
#[derive(PartialEq, Eq, Clone, Debug)]
pub enum Response {
Connect(ConnectResponse),
@ -54,35 +48,30 @@ pub enum Response {
Error(ErrorResponse),
}
impl From<ConnectResponse> for Response {
fn from(r: ConnectResponse) -> Self {
Self::Connect(r)
}
}
impl From<AnnounceResponse> for Response {
fn from(r: AnnounceResponse) -> Self {
Self::Announce(r)
}
}
impl From<ScrapeResponse> for Response {
fn from(r: ScrapeResponse) -> Self {
Self::Scrape(r)
}
}
impl From<ErrorResponse> for Response {
fn from(r: ErrorResponse) -> Self {
Self::Error(r)
}
}
impl Response {
/// Returning IPv6 peers doesn't really work with UDP. It is not supported
/// by https://libtorrent.org/udp_tracker_protocol.html. There is a
@ -91,17 +80,13 @@ impl Response {
/// addresses. Clients seem not to support it very well, but due to a lack
/// of alternative solutions, it is implemented here.
#[inline]
pub fn write(
self,
bytes: &mut impl Write,
ip_version: IpVersion
) -> Result<(), io::Error> {
pub fn write(self, bytes: &mut impl Write, ip_version: IpVersion) -> Result<(), io::Error> {
match self {
Response::Connect(r) => {
bytes.write_i32::<NetworkEndian>(0)?;
bytes.write_i32::<NetworkEndian>(r.transaction_id.0)?;
bytes.write_i64::<NetworkEndian>(r.connection_id.0)?;
},
}
Response::Announce(r) => {
if ip_version == IpVersion::IPv4 {
bytes.write_i32::<NetworkEndian>(1)?;
@ -132,7 +117,7 @@ impl Response {
}
}
}
},
}
Response::Scrape(r) => {
bytes.write_i32::<NetworkEndian>(2)?;
bytes.write_i32::<NetworkEndian>(r.transaction_id.0)?;
@ -142,13 +127,13 @@ impl Response {
bytes.write_i32::<NetworkEndian>(torrent_stat.completed.0)?;
bytes.write_i32::<NetworkEndian>(torrent_stat.leechers.0)?;
}
},
}
Response::Error(r) => {
bytes.write_i32::<NetworkEndian>(3)?;
bytes.write_i32::<NetworkEndian>(r.transaction_id.0)?;
bytes.write_all(r.message.as_bytes())?;
},
}
}
Ok(())
@ -168,9 +153,10 @@ impl Response {
Ok((ConnectResponse {
connection_id: ConnectionId(connection_id),
transaction_id: TransactionId(transaction_id)
}).into())
},
transaction_id: TransactionId(transaction_id),
})
.into())
}
// Announce
1 => {
let announce_interval = cursor.read_i32::<NetworkEndian>()?;
@ -180,49 +166,57 @@ impl Response {
let position = cursor.position() as usize;
let inner = cursor.into_inner();
let peers = inner[position..].chunks_exact(6).map(|chunk| {
let ip_bytes: [u8; 4] = (&chunk[..4]).try_into().unwrap();
let ip_address = IpAddr::V4(Ipv4Addr::from(ip_bytes));
let port = (&chunk[4..]).read_u16::<NetworkEndian>().unwrap();
let peers = inner[position..]
.chunks_exact(6)
.map(|chunk| {
let ip_bytes: [u8; 4] = (&chunk[..4]).try_into().unwrap();
let ip_address = IpAddr::V4(Ipv4Addr::from(ip_bytes));
let port = (&chunk[4..]).read_u16::<NetworkEndian>().unwrap();
ResponsePeer {
ip_address,
port: Port(port),
}
}).collect();
ResponsePeer {
ip_address,
port: Port(port),
}
})
.collect();
Ok((AnnounceResponse {
transaction_id: TransactionId(transaction_id),
announce_interval: AnnounceInterval(announce_interval),
leechers: NumberOfPeers(leechers),
seeders: NumberOfPeers(seeders),
peers
}).into())
},
peers,
})
.into())
}
// Scrape
2 => {
let position = cursor.position() as usize;
let inner = cursor.into_inner();
let stats = inner[position..].chunks_exact(12).map(|chunk| {
let mut cursor: Cursor<&[u8]> = Cursor::new(&chunk[..]);
let stats = inner[position..]
.chunks_exact(12)
.map(|chunk| {
let mut cursor: Cursor<&[u8]> = Cursor::new(&chunk[..]);
let seeders = cursor.read_i32::<NetworkEndian>().unwrap();
let downloads = cursor.read_i32::<NetworkEndian>().unwrap();
let leechers = cursor.read_i32::<NetworkEndian>().unwrap();
let seeders = cursor.read_i32::<NetworkEndian>().unwrap();
let downloads = cursor.read_i32::<NetworkEndian>().unwrap();
let leechers = cursor.read_i32::<NetworkEndian>().unwrap();
TorrentScrapeStatistics {
seeders: NumberOfPeers(seeders),
completed: NumberOfDownloads(downloads),
leechers:NumberOfPeers(leechers)
}
}).collect();
TorrentScrapeStatistics {
seeders: NumberOfPeers(seeders),
completed: NumberOfDownloads(downloads),
leechers: NumberOfPeers(leechers),
}
})
.collect();
Ok((ScrapeResponse {
transaction_id: TransactionId(transaction_id),
torrent_stats: stats
}).into())
},
torrent_stats: stats,
})
.into())
}
// Error
3 => {
let position = cursor.position() as usize;
@ -230,9 +224,10 @@ impl Response {
Ok((ErrorResponse {
transaction_id: TransactionId(transaction_id),
message: String::from_utf8_lossy(&inner[position..]).into()
}).into())
},
message: String::from_utf8_lossy(&inner[position..]).into(),
})
.into())
}
// IPv6 announce
4 => {
let announce_interval = cursor.read_i32::<NetworkEndian>()?;
@ -242,36 +237,38 @@ impl Response {
let position = cursor.position() as usize;
let inner = cursor.into_inner();
let peers = inner[position..].chunks_exact(18).map(|chunk| {
let ip_bytes: [u8; 16] = (&chunk[..16]).try_into().unwrap();
let ip_address = IpAddr::V6(Ipv6Addr::from(ip_bytes));
let port = (&chunk[16..]).read_u16::<NetworkEndian>().unwrap();
let peers = inner[position..]
.chunks_exact(18)
.map(|chunk| {
let ip_bytes: [u8; 16] = (&chunk[..16]).try_into().unwrap();
let ip_address = IpAddr::V6(Ipv6Addr::from(ip_bytes));
let port = (&chunk[16..]).read_u16::<NetworkEndian>().unwrap();
ResponsePeer {
ip_address,
port: Port(port),
}
}).collect();
ResponsePeer {
ip_address,
port: Port(port),
}
})
.collect();
Ok((AnnounceResponse {
transaction_id: TransactionId(transaction_id),
announce_interval: AnnounceInterval(announce_interval),
leechers: NumberOfPeers(leechers),
seeders: NumberOfPeers(seeders),
peers
}).into())
},
_ => {
Ok((ErrorResponse {
transaction_id: TransactionId(transaction_id),
message: "Invalid action".to_string()
}).into())
peers,
})
.into())
}
_ => Ok((ErrorResponse {
transaction_id: TransactionId(transaction_id),
message: "Invalid action".to_string(),
})
.into()),
}
}
}
#[cfg(test)]
mod tests {
use quickcheck_macros::quickcheck;
@ -299,9 +296,9 @@ mod tests {
impl quickcheck::Arbitrary for AnnounceResponse {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
let peers = (0..u8::arbitrary(g)).map(|_| {
ResponsePeer::arbitrary(g)
}).collect();
let peers = (0..u8::arbitrary(g))
.map(|_| ResponsePeer::arbitrary(g))
.collect();
Self {
transaction_id: TransactionId(i32::arbitrary(g)),
@ -315,9 +312,9 @@ mod tests {
impl quickcheck::Arbitrary for ScrapeResponse {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
let torrent_stats = (0..u8::arbitrary(g)).map(|_| {
TorrentScrapeStatistics::arbitrary(g)
}).collect();
let torrent_stats = (0..u8::arbitrary(g))
.map(|_| TorrentScrapeStatistics::arbitrary(g))
.collect();
Self {
transaction_id: TransactionId(i32::arbitrary(g)),
@ -326,10 +323,7 @@ mod tests {
}
}
fn same_after_conversion(
response: Response,
ip_version: IpVersion
) -> bool {
fn same_after_conversion(response: Response, ip_version: IpVersion) -> bool {
let mut buf = Vec::new();
response.clone().write(&mut buf, ip_version).unwrap();
@ -345,16 +339,12 @@ mod tests {
}
#[quickcheck]
fn test_connect_response_convert_identity(
response: ConnectResponse
) -> bool {
fn test_connect_response_convert_identity(response: ConnectResponse) -> bool {
same_after_conversion(response.into(), IpVersion::IPv4)
}
#[quickcheck]
fn test_announce_response_convert_identity(
data: (AnnounceResponse, IpVersion)
) -> bool {
fn test_announce_response_convert_identity(data: (AnnounceResponse, IpVersion)) -> bool {
let mut r = data.0;
if data.1 == IpVersion::IPv4 {
@ -367,9 +357,7 @@ mod tests {
}
#[quickcheck]
fn test_scrape_response_convert_identity(
response: ScrapeResponse
) -> bool {
fn test_scrape_response_convert_identity(response: ScrapeResponse) -> bool {
same_after_conversion(response.into(), IpVersion::IPv4)
}
}

10
aquatic_ws/src/bin/main.rs
View file

@ -1,15 +1,9 @@
use aquatic_cli_helpers::run_app_with_cli_and_config;
use aquatic_ws::config::Config;
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
fn main(){
run_app_with_cli_and_config::<Config>(
aquatic_ws::APP_NAME,
aquatic_ws::run,
None
)
fn main() {
run_app_with_cli_and_config::<Config>(aquatic_ws::APP_NAME, aquatic_ws::run, None)
}

35
aquatic_ws/src/lib/common.rs
View file

@ -1,7 +1,7 @@
use std::net::{SocketAddr, IpAddr};
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use crossbeam_channel::{Sender, Receiver};
use crossbeam_channel::{Receiver, Sender};
use hashbrown::HashMap;
use indexmap::IndexMap;
use log::error;
@ -12,11 +12,9 @@ pub use aquatic_common::ValidUntil;
use aquatic_ws_protocol::*;
pub const LISTENER_TOKEN: Token = Token(0);
pub const CHANNEL_TOKEN: Token = Token(1);
#[derive(Clone, Copy, Debug)]
pub struct ConnectionMeta {
/// Index of socket worker responsible for this connection. Required for
@ -29,24 +27,19 @@ pub struct ConnectionMeta {
pub poll_token: Token,
}
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum PeerStatus {
Seeding,
Leeching,
Stopped
Stopped,
}
impl PeerStatus {
/// Determine peer status from announce event and number of bytes left.
///
/// Likely, the last branch will be taken most of the time.
#[inline]
pub fn from_event_and_bytes_left(
event: AnnounceEvent,
opt_bytes_left: Option<usize>
) -> Self {
pub fn from_event_and_bytes_left(event: AnnounceEvent, opt_bytes_left: Option<usize>) -> Self {
if let AnnounceEvent::Stopped = event {
Self::Stopped
} else if let Some(0) = opt_bytes_left {
@ -57,7 +50,6 @@ impl PeerStatus {
}
}
#[derive(Clone, Copy)]
pub struct Peer {
pub connection_meta: ConnectionMeta,
@ -65,17 +57,14 @@ pub struct Peer {
pub valid_until: ValidUntil,
}
pub type PeerMap = IndexMap<PeerId, Peer>;
pub struct TorrentData {
pub peers: PeerMap,
pub num_seeders: usize,
pub num_leechers: usize,
}
impl Default for TorrentData {
#[inline]
fn default() -> Self {
@ -87,23 +76,19 @@ impl Default for TorrentData {
}
}
pub type TorrentMap = HashMap<InfoHash, TorrentData>;
#[derive(Default)]
pub struct TorrentMaps {
pub ipv4: TorrentMap,
pub ipv6: TorrentMap,
}
#[derive(Clone)]
pub struct State {
pub torrent_maps: Arc<Mutex<TorrentMaps>>,
}
impl Default for State {
fn default() -> Self {
Self {
@ -112,33 +97,25 @@ impl Default for State {
}
}
pub type InMessageSender = Sender<(ConnectionMeta, InMessage)>;
pub type InMessageReceiver = Receiver<(ConnectionMeta, InMessage)>;
pub type OutMessageReceiver = Receiver<(ConnectionMeta, OutMessage)>;
#[derive(Clone)]
pub struct OutMessageSender(Vec<Sender<(ConnectionMeta, OutMessage)>>);
impl OutMessageSender {
pub fn new(senders: Vec<Sender<(ConnectionMeta, OutMessage)>>) -> Self {
Self(senders)
}
#[inline]
pub fn send(
&self,
meta: ConnectionMeta,
message: OutMessage
){
if let Err(err) = self.0[meta.worker_index].send((meta, message)){
pub fn send(&self, meta: ConnectionMeta, message: OutMessage) {
if let Err(err) = self.0[meta.worker_index].send((meta, message)) {
error!("OutMessageSender: couldn't send message: {:?}", err);
}
}
}
pub type SocketWorkerStatus = Option<Result<(), String>>;
pub type SocketWorkerStatuses = Arc<Mutex<Vec<SocketWorkerStatus>>>;

21
aquatic_ws/src/lib/config.rs
View file

@ -1,10 +1,9 @@
use std::net::SocketAddr;
use serde::{Serialize, Deserialize};
use serde::{Deserialize, Serialize};
use aquatic_cli_helpers::LogLevel;
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct Config {
@ -24,14 +23,12 @@ pub struct Config {
pub privileges: PrivilegeConfig,
}
impl aquatic_cli_helpers::Config for Config {
fn get_log_level(&self) -> Option<LogLevel> {
Some(self.log_level)
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct NetworkConfig {
@ -47,7 +44,6 @@ pub struct NetworkConfig {
pub websocket_max_frame_size: usize,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct HandlerConfig {
@ -57,7 +53,6 @@ pub struct HandlerConfig {
pub channel_recv_timeout_microseconds: u64,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct ProtocolConfig {
@ -69,7 +64,6 @@ pub struct ProtocolConfig {
pub peer_announce_interval: usize,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct CleaningConfig {
@ -81,7 +75,6 @@ pub struct CleaningConfig {
pub max_connection_age: u64,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct StatisticsConfig {
@ -89,7 +82,6 @@ pub struct StatisticsConfig {
pub interval: u64,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct PrivilegeConfig {
@ -101,7 +93,6 @@ pub struct PrivilegeConfig {
pub user: String,
}
impl Default for Config {
fn default() -> Self {
Self {
@ -118,7 +109,6 @@ impl Default for Config {
}
}
impl Default for NetworkConfig {
fn default() -> Self {
Self {
@ -135,7 +125,6 @@ impl Default for NetworkConfig {
}
}
impl Default for ProtocolConfig {
fn default() -> Self {
Self {
@ -146,7 +135,6 @@ impl Default for ProtocolConfig {
}
}
impl Default for HandlerConfig {
fn default() -> Self {
Self {
@ -156,7 +144,6 @@ impl Default for HandlerConfig {
}
}
impl Default for CleaningConfig {
fn default() -> Self {
Self {
@ -167,16 +154,12 @@ impl Default for CleaningConfig {
}
}
impl Default for StatisticsConfig {
fn default() -> Self {
Self {
interval: 0,
}
Self { interval: 0 }
}
}
impl Default for PrivilegeConfig {
fn default() -> Self {
Self {

91
aquatic_ws/src/lib/handler.rs
View file

@ -1,11 +1,11 @@
use std::sync::Arc;
use std::time::Duration;
use std::vec::Drain;
use std::sync::Arc;
use hashbrown::HashMap;
use mio::Waker;
use parking_lot::MutexGuard;
use rand::{Rng, SeedableRng, rngs::SmallRng};
use rand::{rngs::SmallRng, Rng, SeedableRng};
use aquatic_common::extract_response_peers;
use aquatic_ws_protocol::*;
@ -13,26 +13,21 @@ use aquatic_ws_protocol::*;
use crate::common::*;
use crate::config::Config;
pub fn run_request_worker(
config: Config,
state: State,
in_message_receiver: InMessageReceiver,
out_message_sender: OutMessageSender,
wakers: Vec<Arc<Waker>>,
){
let mut wake_socket_workers: Vec<bool> = (0..config.socket_workers)
.map(|_| false)
.collect();
) {
let mut wake_socket_workers: Vec<bool> = (0..config.socket_workers).map(|_| false).collect();
let mut announce_requests = Vec::new();
let mut scrape_requests = Vec::new();
let mut rng = SmallRng::from_entropy();
let timeout = Duration::from_micros(
config.handlers.channel_recv_timeout_microseconds
);
let timeout = Duration::from_micros(config.handlers.channel_recv_timeout_microseconds);
loop {
let mut opt_torrent_map_guard: Option<MutexGuard<TorrentMaps>> = None;
@ -47,22 +42,22 @@ pub fn run_request_worker(
match opt_in_message {
Some((meta, InMessage::AnnounceRequest(r))) => {
announce_requests.push((meta, r));
},
}
Some((meta, InMessage::ScrapeRequest(r))) => {
scrape_requests.push((meta, r));
},
}
None => {
if let Some(torrent_guard) = state.torrent_maps.try_lock(){
if let Some(torrent_guard) = state.torrent_maps.try_lock() {
opt_torrent_map_guard = Some(torrent_guard);
break
break;
}
}
}
}
let mut torrent_map_guard = opt_torrent_map_guard
.unwrap_or_else(|| state.torrent_maps.lock());
let mut torrent_map_guard =
opt_torrent_map_guard.unwrap_or_else(|| state.torrent_maps.lock());
handle_announce_requests(
&config,
@ -70,7 +65,7 @@ pub fn run_request_worker(
&mut torrent_map_guard,
&out_message_sender,
&mut wake_socket_workers,
announce_requests.drain(..)
announce_requests.drain(..),
);
handle_scrape_requests(
@ -78,12 +73,12 @@ pub fn run_request_worker(
&mut torrent_map_guard,
&out_message_sender,
&mut wake_socket_workers,
scrape_requests.drain(..)
scrape_requests.drain(..),
);
for (worker_index, wake) in wake_socket_workers.iter_mut().enumerate(){
for (worker_index, wake) in wake_socket_workers.iter_mut().enumerate() {
if *wake {
if let Err(err) = wakers[worker_index].wake(){
if let Err(err) = wakers[worker_index].wake() {
::log::error!("request handler couldn't wake poll: {:?}", err);
}
@ -93,7 +88,6 @@ pub fn run_request_worker(
}
}
pub fn handle_announce_requests(
config: &Config,
rng: &mut impl Rng,
@ -101,11 +95,11 @@ pub fn handle_announce_requests(
out_message_sender: &OutMessageSender,
wake_socket_workers: &mut Vec<bool>,
requests: Drain<(ConnectionMeta, AnnounceRequest)>,
){
) {
let valid_until = ValidUntil::new(config.cleaning.max_peer_age);
for (request_sender_meta, request) in requests {
let torrent_data: &mut TorrentData = if request_sender_meta.converted_peer_ip.is_ipv4(){
let torrent_data: &mut TorrentData = if request_sender_meta.converted_peer_ip.is_ipv4() {
torrent_maps.ipv4.entry(request.info_hash).or_default()
} else {
torrent_maps.ipv6.entry(request.info_hash).or_default()
@ -117,8 +111,9 @@ pub fn handle_announce_requests(
// requests using them, causing all sorts of issues. Checking naive
// (non-converted) socket addresses is enough, since state is split
// on converted peer ip.
if let Some(previous_peer) = torrent_data.peers.get(&request.peer_id){
if request_sender_meta.naive_peer_addr != previous_peer.connection_meta.naive_peer_addr {
if let Some(previous_peer) = torrent_data.peers.get(&request.peer_id) {
if request_sender_meta.naive_peer_addr != previous_peer.connection_meta.naive_peer_addr
{
continue;
}
}
@ -127,7 +122,7 @@ pub fn handle_announce_requests(
{
let peer_status = PeerStatus::from_event_and_bytes_left(
request.event.unwrap_or_default(),
request.bytes_left
request.bytes_left,
);
let peer = Peer {
@ -141,24 +136,22 @@ pub fn handle_announce_requests(
torrent_data.num_leechers += 1;
torrent_data.peers.insert(request.peer_id, peer)
},
}
PeerStatus::Seeding => {
torrent_data.num_seeders += 1;
torrent_data.peers.insert(request.peer_id, peer)
},
PeerStatus::Stopped => {
torrent_data.peers.remove(&request.peer_id)
}
PeerStatus::Stopped => torrent_data.peers.remove(&request.peer_id),
};
match opt_removed_peer.map(|peer| peer.status){
match opt_removed_peer.map(|peer| peer.status) {
Some(PeerStatus::Leeching) => {
torrent_data.num_leechers -= 1;
},
}
Some(PeerStatus::Seeding) => {
torrent_data.num_seeders -= 1;
},
}
_ => {}
}
}
@ -166,8 +159,7 @@ pub fn handle_announce_requests(
// If peer sent offers, send them on to random peers
if let Some(offers) = request.offers {
// FIXME: config: also maybe check this when parsing request
let max_num_peers_to_take = offers.len()
.min(config.protocol.max_offers);
let max_num_peers_to_take = offers.len().min(config.protocol.max_offers);
#[inline]
fn f(peer: &Peer) -> Peer {
@ -179,12 +171,10 @@ pub fn handle_announce_requests(
&torrent_data.peers,
max_num_peers_to_take,
request.peer_id,
f
f,
);
for (offer, offer_receiver) in offers.into_iter()
.zip(offer_receivers)
{
for (offer, offer_receiver) in offers.into_iter().zip(offer_receivers) {
let middleman_offer = MiddlemanOfferToPeer {
action: AnnounceAction,
info_hash: request.info_hash,
@ -195,7 +185,7 @@ pub fn handle_announce_requests(
out_message_sender.send(
offer_receiver.connection_meta,
OutMessage::Offer(middleman_offer)
OutMessage::Offer(middleman_offer),
);
wake_socket_workers[offer_receiver.connection_meta.worker_index] = true;
}
@ -205,9 +195,7 @@ pub fn handle_announce_requests(
if let (Some(answer), Some(answer_receiver_id), Some(offer_id)) =
(request.answer, request.to_peer_id, request.offer_id)
{
if let Some(answer_receiver) = torrent_data.peers
.get(&answer_receiver_id)
{
if let Some(answer_receiver) = torrent_data.peers.get(&answer_receiver_id) {
let middleman_answer = MiddlemanAnswerToPeer {
action: AnnounceAction,
peer_id: request.peer_id,
@ -218,7 +206,7 @@ pub fn handle_announce_requests(
out_message_sender.send(
answer_receiver.connection_meta,
OutMessage::Answer(middleman_answer)
OutMessage::Answer(middleman_answer),
);
wake_socket_workers[answer_receiver.connection_meta.worker_index] = true;
}
@ -237,31 +225,28 @@ pub fn handle_announce_requests(
}
}
pub fn handle_scrape_requests(
config: &Config,
torrent_maps: &mut TorrentMaps,
out_message_sender: &OutMessageSender,
wake_socket_workers: &mut Vec<bool>,
requests: Drain<(ConnectionMeta, ScrapeRequest)>,
){
) {
for (meta, request) in requests {
let info_hashes = if let Some(info_hashes) = request.info_hashes {
info_hashes.as_vec()
} else {
continue
continue;
};
let num_to_take = info_hashes.len().min(
config.protocol.max_scrape_torrents
);
let num_to_take = info_hashes.len().min(config.protocol.max_scrape_torrents);
let mut response = ScrapeResponse {
action: ScrapeAction,
files: HashMap::with_capacity(num_to_take),
};
let torrent_map: &mut TorrentMap = if meta.converted_peer_ip.is_ipv4(){
let torrent_map: &mut TorrentMap = if meta.converted_peer_ip.is_ipv4() {
&mut torrent_maps.ipv4
} else {
&mut torrent_maps.ipv6
@ -269,8 +254,8 @@ pub fn handle_scrape_requests(
// If request.info_hashes is empty, don't return scrape for all
// torrents, even though reference server does it. It is too expensive.
for info_hash in info_hashes.into_iter().take(num_to_take){
if let Some(torrent_data) = torrent_map.get(&info_hash){
for info_hash in info_hashes.into_iter().take(num_to_take) {
if let Some(torrent_data) = torrent_map.get(&info_hash) {
let stats = ScrapeStatistics {
complete: torrent_data.num_seeders,
downloaded: 0, // No implementation planned

81
aquatic_ws/src/lib/lib.rs
View file

@ -1,8 +1,8 @@
use std::time::Duration;
use std::fs::File;
use std::io::Read;
use std::sync::Arc;
use std::thread::Builder;
use std::time::Duration;
use anyhow::Context;
use mio::{Poll, Waker};
@ -19,10 +19,8 @@ pub mod tasks;
use common::*;
use config::Config;
pub const APP_NAME: &str = "aquatic_ws: WebTorrent tracker";
pub fn run(config: Config) -> anyhow::Result<()> {
let state = State::default();
@ -35,7 +33,6 @@ pub fn run(config: Config) -> anyhow::Result<()> {
}
}
pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
let opt_tls_acceptor = create_tls_acceptor(&config)?;
@ -67,17 +64,19 @@ pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
out_message_senders.push(out_message_sender);
wakers.push(waker);
Builder::new().name(format!("socket-{:02}", i + 1)).spawn(move || {
network::run_socket_worker(
config,
i,
socket_worker_statuses,
poll,
in_message_sender,
out_message_receiver,
opt_tls_acceptor
);
})?;
Builder::new()
.name(format!("socket-{:02}", i + 1))
.spawn(move || {
network::run_socket_worker(
config,
i,
socket_worker_statuses,
poll,
in_message_sender,
out_message_receiver,
opt_tls_acceptor,
);
})?;
}
// Wait for socket worker statuses. On error from any, quit program.
@ -86,14 +85,14 @@ pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
loop {
::std::thread::sleep(::std::time::Duration::from_millis(10));
if let Some(statuses) = socket_worker_statuses.try_lock(){
for opt_status in statuses.iter(){
if let Some(statuses) = socket_worker_statuses.try_lock() {
for opt_status in statuses.iter() {
if let Some(Err(err)) = opt_status {
return Err(::anyhow::anyhow!(err.to_owned()));
}
}
if statuses.iter().all(Option::is_some){
if statuses.iter().all(Option::is_some) {
if config.privileges.drop_privileges {
PrivDrop::default()
.chroot(config.privileges.chroot_path.clone())
@ -102,7 +101,7 @@ pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
.context("Couldn't drop root privileges")?;
}
break
break;
}
}
}
@ -116,39 +115,37 @@ pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
let out_message_sender = out_message_sender.clone();
let wakers = wakers.clone();
Builder::new().name(format!("request-{:02}", i + 1)).spawn(move || {
handler::run_request_worker(
config,
state,
in_message_receiver,
out_message_sender,
wakers,
);
})?;
Builder::new()
.name(format!("request-{:02}", i + 1))
.spawn(move || {
handler::run_request_worker(
config,
state,
in_message_receiver,
out_message_sender,
wakers,
);
})?;
}
if config.statistics.interval != 0 {
let state = state.clone();
let config = config.clone();
Builder::new().name("statistics".to_string()).spawn(move ||
loop {
::std::thread::sleep(Duration::from_secs(
config.statistics.interval
));
Builder::new()
.name("statistics".to_string())
.spawn(move || loop {
::std::thread::sleep(Duration::from_secs(config.statistics.interval));
tasks::print_statistics(&state);
}
).expect("spawn statistics thread");
})
.expect("spawn statistics thread");
}
Ok(())
}
pub fn create_tls_acceptor(
config: &Config,
) -> anyhow::Result<Option<TlsAcceptor>> {
pub fn create_tls_acceptor(config: &Config) -> anyhow::Result<Option<TlsAcceptor>> {
if config.network.use_tls {
let mut identity_bytes = Vec::new();
let mut file = File::open(&config.network.tls_pkcs12_path)
@ -157,10 +154,8 @@ pub fn create_tls_acceptor(
file.read_to_end(&mut identity_bytes)
.context("Couldn't read pkcs12 identity file")?;
let identity = Identity::from_pkcs12(
&identity_bytes,
&config.network.tls_pkcs12_password
).context("Couldn't parse pkcs12 identity file")?;
let identity = Identity::from_pkcs12(&identity_bytes, &config.network.tls_pkcs12_password)
.context("Couldn't parse pkcs12 identity file")?;
let acceptor = TlsAcceptor::new(identity)
.context("Couldn't create TlsAcceptor from pkcs12 identity")?;

118
aquatic_ws/src/lib/network/connection.rs
View file

@ -1,26 +1,24 @@
use std::net::{SocketAddr};
use std::io::{Read, Write};
use std::net::SocketAddr;
use either::Either;
use hashbrown::HashMap;
use log::info;
use mio::{Poll, Token};
use mio::net::TcpStream;
use native_tls::{TlsAcceptor, TlsStream, MidHandshakeTlsStream};
use tungstenite::WebSocket;
use tungstenite::handshake::{MidHandshake, HandshakeError, server::NoCallback};
use tungstenite::server::{ServerHandshake};
use mio::{Poll, Token};
use native_tls::{MidHandshakeTlsStream, TlsAcceptor, TlsStream};
use tungstenite::handshake::{server::NoCallback, HandshakeError, MidHandshake};
use tungstenite::protocol::WebSocketConfig;
use tungstenite::server::ServerHandshake;
use tungstenite::WebSocket;
use crate::common::*;
pub enum Stream {
TcpStream(TcpStream),
TlsStream(TlsStream<TcpStream>),
}
impl Stream {
#[inline]
pub fn get_peer_addr(&self) -> SocketAddr {
@ -33,15 +31,12 @@ impl Stream {
#[inline]
pub fn deregister(&mut self, poll: &mut Poll) -> ::std::io::Result<()> {
match self {
Self::TcpStream(stream) =>
poll.registry().deregister(stream),
Self::TlsStream(stream) =>
poll.registry().deregister(stream.get_mut()),
Self::TcpStream(stream) => poll.registry().deregister(stream),
Self::TlsStream(stream) => poll.registry().deregister(stream.get_mut()),
}
}
}
impl Read for Stream {
#[inline]
fn read(&mut self, buf: &mut [u8]) -> Result<usize, ::std::io::Error> {
@ -55,7 +50,7 @@ impl Read for Stream {
#[inline]
fn read_vectored(
&mut self,
bufs: &mut [::std::io::IoSliceMut<'_>]
bufs: &mut [::std::io::IoSliceMut<'_>],
) -> ::std::io::Result<usize> {
match self {
Self::TcpStream(stream) => stream.read_vectored(bufs),
@ -64,7 +59,6 @@ impl Read for Stream {
}
}
impl Write for Stream {
#[inline]
fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
@ -76,10 +70,7 @@ impl Write for Stream {
/// Not used but provided for completeness
#[inline]
fn write_vectored(
&mut self,
bufs: &[::std::io::IoSlice<'_>]
) -> ::std::io::Result<usize> {
fn write_vectored(&mut self, bufs: &[::std::io::IoSlice<'_>]) -> ::std::io::Result<usize> {
match self {
Self::TcpStream(stream) => stream.write_vectored(bufs),
Self::TlsStream(stream) => stream.write_vectored(bufs),
@ -95,7 +86,6 @@ impl Write for Stream {
}
}
enum HandshakeMachine {
TcpStream(TcpStream),
TlsStream(TlsStream<TcpStream>),
@ -103,7 +93,6 @@ enum HandshakeMachine {
WsMidHandshake(MidHandshake<ServerHandshake<Stream, NoCallback>>),
}
impl HandshakeMachine {
#[inline]
fn new(tcp_stream: TcpStream) -> Self {
@ -115,35 +104,32 @@ impl HandshakeMachine {
self,
ws_config: WebSocketConfig,
opt_tls_acceptor: &Option<TlsAcceptor>, // If set, run TLS
) -> (Option<Either<EstablishedWs, Self>>, bool) { // bool = stop looping
) -> (Option<Either<EstablishedWs, Self>>, bool) {
// bool = stop looping
match self {
HandshakeMachine::TcpStream(stream) => {
if let Some(tls_acceptor) = opt_tls_acceptor {
Self::handle_tls_handshake_result(
tls_acceptor.accept(stream)
)
Self::handle_tls_handshake_result(tls_acceptor.accept(stream))
} else {
let handshake_result = ::tungstenite::server::accept_with_config(
Stream::TcpStream(stream),
Some(ws_config)
Some(ws_config),
);
Self::handle_ws_handshake_result(handshake_result)
}
},
}
HandshakeMachine::TlsStream(stream) => {
let handshake_result = ::tungstenite::server::accept(
Stream::TlsStream(stream),
);
let handshake_result = ::tungstenite::server::accept(Stream::TlsStream(stream));
Self::handle_ws_handshake_result(handshake_result)
},
}
HandshakeMachine::TlsMidHandshake(handshake) => {
Self::handle_tls_handshake_result(handshake.handshake())
},
}
HandshakeMachine::WsMidHandshake(handshake) => {
Self::handle_ws_handshake_result(handshake.handshake())
},
}
}
}
@ -152,12 +138,10 @@ impl HandshakeMachine {
result: Result<TlsStream<TcpStream>, ::native_tls::HandshakeError<TcpStream>>,
) -> (Option<Either<EstablishedWs, Self>>, bool) {
match result {
Ok(stream) => {
(Some(Either::Right(Self::TlsStream(stream))), false)
},
Ok(stream) => (Some(Either::Right(Self::TlsStream(stream))), false),
Err(native_tls::HandshakeError::WouldBlock(handshake)) => {
(Some(Either::Right(Self::TlsMidHandshake(handshake))), true)
},
}
Err(native_tls::HandshakeError::Failure(err)) => {
info!("tls handshake error: {}", err);
@ -168,22 +152,20 @@ impl HandshakeMachine {
#[inline]
fn handle_ws_handshake_result(
result: Result<WebSocket<Stream>, HandshakeError<ServerHandshake<Stream, NoCallback>>> ,
result: Result<WebSocket<Stream>, HandshakeError<ServerHandshake<Stream, NoCallback>>>,
) -> (Option<Either<EstablishedWs, Self>>, bool) {
match result {
Ok(mut ws) => {
let peer_addr = ws.get_mut().get_peer_addr();
let established_ws = EstablishedWs {
ws,
peer_addr,
};
let established_ws = EstablishedWs { ws, peer_addr };
(Some(Either::Left(established_ws)), false)
},
Err(HandshakeError::Interrupted(handshake)) => {
(Some(Either::Right(HandshakeMachine::WsMidHandshake(handshake))), true)
},
}
Err(HandshakeError::Interrupted(handshake)) => (
Some(Either::Right(HandshakeMachine::WsMidHandshake(handshake))),
true,
),
Err(HandshakeError::Failure(err)) => {
info!("ws handshake error: {}", err);
@ -193,20 +175,17 @@ impl HandshakeMachine {
}
}
pub struct EstablishedWs {
pub ws: WebSocket<Stream>,
pub peer_addr: SocketAddr,
}
pub struct Connection {
ws_config: WebSocketConfig,
pub valid_until: ValidUntil,
inner: Either<EstablishedWs, HandshakeMachine>,
}
/// Create from TcpStream. Run `advance_handshakes` until `get_established_ws`
/// returns Some(EstablishedWs).
///
@ -219,15 +198,11 @@ pub struct Connection {
/// single method for advancing handshakes and maybe returning a websocket.
impl Connection {
#[inline]
pub fn new(
ws_config: WebSocketConfig,
valid_until: ValidUntil,
tcp_stream: TcpStream,
) -> Self {
pub fn new(ws_config: WebSocketConfig, valid_until: ValidUntil, tcp_stream: TcpStream) -> Self {
Self {
ws_config,
valid_until,
inner: Either::Right(HandshakeMachine::new(tcp_stream))
inner: Either::Right(HandshakeMachine::new(tcp_stream)),
}
}
@ -250,15 +225,12 @@ impl Connection {
Either::Right(machine) => {
let ws_config = self.ws_config;
let (opt_inner, stop_loop) = machine.advance(
ws_config,
opt_tls_acceptor
);
let (opt_inner, stop_loop) = machine.advance(ws_config, opt_tls_acceptor);
let opt_new_self = opt_inner.map(|inner| Self {
ws_config,
valid_until,
inner
inner,
});
(opt_new_self, stop_loop)
@ -267,19 +239,16 @@ impl Connection {
}
#[inline]
pub fn close(&mut self){
pub fn close(&mut self) {
if let Either::Left(ref mut ews) = self.inner {
if ews.ws.can_read(){
if let Err(err) = ews.ws.close(None){
if ews.ws.can_read() {
if let Err(err) = ews.ws.close(None) {
::log::info!("error closing ws: {}", err);
}
// Required after ws.close()
if let Err(err) = ews.ws.write_pending(){
::log::info!(
"error writing pending messages after closing ws: {}",
err
)
if let Err(err) = ews.ws.write_pending() {
::log::info!("error writing pending messages after closing ws: {}", err)
}
}
}
@ -289,24 +258,21 @@ impl Connection {
use Either::{Left, Right};
match self.inner {
Left(EstablishedWs { ref mut ws, .. }) => {
ws.get_mut().deregister(poll)
},
Left(EstablishedWs { ref mut ws, .. }) => ws.get_mut().deregister(poll),
Right(HandshakeMachine::TcpStream(ref mut stream)) => {
poll.registry().deregister(stream)
},
}
Right(HandshakeMachine::TlsMidHandshake(ref mut handshake)) => {
poll.registry().deregister(handshake.get_mut())
},
}
Right(HandshakeMachine::TlsStream(ref mut stream)) => {
poll.registry().deregister(stream.get_mut())
},
}
Right(HandshakeMachine::WsMidHandshake(ref mut handshake)) => {
handshake.get_mut().get_mut().deregister(poll)
},
}
}
}
}
pub type ConnectionMap = HashMap<Token, Connection>;

120
aquatic_ws/src/lib/network/mod.rs
View file

@ -1,12 +1,12 @@
use std::time::Duration;
use std::io::ErrorKind;
use std::time::Duration;
use crossbeam_channel::Receiver;
use hashbrown::HashMap;
use log::{info, debug, error};
use native_tls::TlsAcceptor;
use mio::{Events, Poll, Interest, Token};
use log::{debug, error, info};
use mio::net::TcpListener;
use mio::{Events, Interest, Poll, Token};
use native_tls::TlsAcceptor;
use tungstenite::protocol::WebSocketConfig;
use aquatic_common::convert_ipv4_mapped_ipv6;
@ -21,7 +21,6 @@ pub mod utils;
use connection::*;
use utils::*;
pub fn run_socket_worker(
config: Config,
socket_worker_index: usize,
@ -30,8 +29,8 @@ pub fn run_socket_worker(
in_message_sender: InMessageSender,
out_message_receiver: OutMessageReceiver,
opt_tls_acceptor: Option<TlsAcceptor>,
){
match create_listener(&config){
) {
match create_listener(&config) {
Ok(listener) => {
socket_worker_statuses.lock()[socket_worker_index] = Some(Ok(()));
@ -42,18 +41,16 @@ pub fn run_socket_worker(
in_message_sender,
out_message_receiver,
listener,
opt_tls_acceptor
opt_tls_acceptor,
);
},
}
Err(err) => {
socket_worker_statuses.lock()[socket_worker_index] = Some(
Err(format!("Couldn't open socket: {:#}", err))
);
socket_worker_statuses.lock()[socket_worker_index] =
Some(Err(format!("Couldn't open socket: {:#}", err)));
}
}
}
pub fn run_poll_loop(
config: Config,
socket_worker_index: usize,
@ -62,10 +59,8 @@ pub fn run_poll_loop(
out_message_receiver: OutMessageReceiver,
listener: ::std::net::TcpListener,
opt_tls_acceptor: Option<TlsAcceptor>,
){
let poll_timeout = Duration::from_micros(
config.network.poll_timeout_microseconds
);
) {
let poll_timeout = Duration::from_micros(config.network.poll_timeout_microseconds);
let ws_config = WebSocketConfig {
max_message_size: Some(config.network.websocket_max_message_size),
max_frame_size: Some(config.network.websocket_max_frame_size),
@ -91,7 +86,7 @@ pub fn run_poll_loop(
let valid_until = ValidUntil::new(config.cleaning.max_connection_age);
for event in events.iter(){
for event in events.iter() {
let token = event.token();
if token == LISTENER_TOKEN {
@ -115,11 +110,7 @@ pub fn run_poll_loop(
);
}
send_out_messages(
&mut poll,
&out_message_receiver,
&mut connections
);
send_out_messages(&mut poll, &out_message_receiver, &mut connections);
}
// Remove inactive connections, but not every iteration
@ -131,7 +122,6 @@ pub fn run_poll_loop(
}
}
fn accept_new_streams(
ws_config: WebSocketConfig,
listener: &mut TcpListener,
@ -139,9 +129,9 @@ fn accept_new_streams(
connections: &mut ConnectionMap,
valid_until: ValidUntil,
poll_token_counter: &mut Token,
){
) {
loop {
match listener.accept(){
match listener.accept() {
Ok((mut stream, _)) => {
poll_token_counter.0 = poll_token_counter.0.wrapping_add(1);
@ -160,10 +150,10 @@ fn accept_new_streams(
let connection = Connection::new(ws_config, valid_until, stream);
connections.insert(token, connection);
},
}
Err(err) => {
if err.kind() == ErrorKind::WouldBlock {
break
break;
}
info!("error while accepting streams: {}", err);
@ -172,7 +162,6 @@ fn accept_new_streams(
}
}
/// On the stream given by poll_token, get TLS (if requested) and tungstenite
/// up and running, then read messages and pass on through channel.
pub fn run_handshakes_and_read_messages(
@ -183,10 +172,12 @@ pub fn run_handshakes_and_read_messages(
connections: &mut ConnectionMap,
poll_token: Token,
valid_until: ValidUntil,
){
) {
loop {
if let Some(established_ws) = connections.get_mut(&poll_token)
.map(|c| { // Ugly but works
if let Some(established_ws) = connections
.get_mut(&poll_token)
.map(|c| {
// Ugly but works
c.valid_until = valid_until;
c
@ -195,13 +186,11 @@ pub fn run_handshakes_and_read_messages(
{
use ::tungstenite::Error::Io;
match established_ws.ws.read_message(){
match established_ws.ws.read_message() {
Ok(ws_message) => {
if let Ok(in_message) = InMessage::from_ws_message(ws_message){
if let Ok(in_message) = InMessage::from_ws_message(ws_message) {
let naive_peer_addr = established_ws.peer_addr;
let converted_peer_ip = convert_ipv4_mapped_ipv6(
naive_peer_addr.ip()
);
let converted_peer_ip = convert_ipv4_mapped_ipv6(naive_peer_addr.ip());
let meta = ConnectionMeta {
worker_index: socket_worker_index,
@ -212,24 +201,19 @@ pub fn run_handshakes_and_read_messages(
debug!("read message");
if let Err(err) = in_message_sender
.send((meta, in_message))
{
error!(
"InMessageSender: couldn't send message: {:?}",
err
);
if let Err(err) = in_message_sender.send((meta, in_message)) {
error!("InMessageSender: couldn't send message: {:?}", err);
}
}
},
}
Err(Io(err)) if err.kind() == ErrorKind::WouldBlock => {
break;
},
}
Err(tungstenite::Error::ConnectionClosed) => {
remove_connection_if_exists(poll, connections, poll_token);
break
},
break;
}
Err(err) => {
info!("error reading messages: {}", err);
@ -238,11 +222,9 @@ pub fn run_handshakes_and_read_messages(
break;
}
}
} else if let Some(connection) = connections.remove(&poll_token){
let (opt_new_connection, stop_loop) = connection.advance_handshakes(
opt_tls_acceptor,
valid_until
);
} else if let Some(connection) = connections.remove(&poll_token) {
let (opt_new_connection, stop_loop) =
connection.advance_handshakes(opt_tls_acceptor, valid_until);
if let Some(connection) = opt_new_connection {
connections.insert(poll_token, connection);
@ -252,22 +234,22 @@ pub fn run_handshakes_and_read_messages(
break;
}
} else {
break
break;
}
}
}
/// Read messages from channel, send to peers
pub fn send_out_messages(
poll: &mut Poll,
out_message_receiver: &Receiver<(ConnectionMeta, OutMessage)>,
connections: &mut ConnectionMap,
){
) {
let len = out_message_receiver.len();
for (meta, out_message) in out_message_receiver.try_iter().take(len){
let opt_established_ws = connections.get_mut(&meta.poll_token)
for (meta, out_message) in out_message_receiver.try_iter().take(len) {
let opt_established_ws = connections
.get_mut(&meta.poll_token)
.and_then(Connection::get_established_ws);
if let Some(established_ws) = opt_established_ws {
@ -281,27 +263,19 @@ pub fn send_out_messages(
let ws_message = out_message.to_ws_message();
match established_ws.ws.write_message(ws_message){
match established_ws.ws.write_message(ws_message) {
Ok(()) => {
debug!("sent message");
},
Err(Io(err)) if err.kind() == ErrorKind::WouldBlock => {},
}
Err(Io(err)) if err.kind() == ErrorKind::WouldBlock => {}
Err(tungstenite::Error::ConnectionClosed) => {
remove_connection_if_exists(
poll,
connections,
meta.poll_token
);
},
remove_connection_if_exists(poll, connections, meta.poll_token);
}
Err(err) => {
info!("error writing ws message: {}", err);
remove_connection_if_exists(
poll,
connections,
meta.poll_token
);
},
remove_connection_if_exists(poll, connections, meta.poll_token);
}
}
}
}

46
aquatic_ws/src/lib/network/utils.rs
View file

@ -2,60 +2,54 @@ use std::time::Instant;
use anyhow::Context;
use mio::{Poll, Token};
use socket2::{Socket, Domain, Type, Protocol};
use socket2::{Domain, Protocol, Socket, Type};
use crate::config::Config;
use super::connection::*;
pub fn create_listener(
config: &Config
) -> ::anyhow::Result<::std::net::TcpListener> {
let builder = if config.network.address.is_ipv4(){
pub fn create_listener(config: &Config) -> ::anyhow::Result<::std::net::TcpListener> {
let builder = if config.network.address.is_ipv4() {
Socket::new(Domain::ipv4(), Type::stream(), Some(Protocol::tcp()))
} else {
Socket::new(Domain::ipv6(), Type::stream(), Some(Protocol::tcp()))
}.context("Couldn't create socket2::Socket")?;
}
.context("Couldn't create socket2::Socket")?;
if config.network.ipv6_only {
builder.set_only_v6(true)
builder
.set_only_v6(true)
.context("Couldn't put socket in ipv6 only mode")?
}
builder.set_nonblocking(true)
builder
.set_nonblocking(true)
.context("Couldn't put socket in non-blocking mode")?;
builder.set_reuse_port(true)
builder
.set_reuse_port(true)
.context("Couldn't put socket in reuse_port mode")?;
builder.bind(&config.network.address.into()).with_context(||
format!("Couldn't bind socket to address {}", config.network.address)
)?;
builder.listen(128)
builder
.bind(&config.network.address.into())
.with_context(|| format!("Couldn't bind socket to address {}", config.network.address))?;
builder
.listen(128)
.context("Couldn't listen for connections on socket")?;
Ok(builder.into_tcp_listener())
}
pub fn remove_connection_if_exists(
poll: &mut Poll,
connections: &mut ConnectionMap,
token: Token,
){
if let Some(mut connection) = connections.remove(&token){
pub fn remove_connection_if_exists(poll: &mut Poll, connections: &mut ConnectionMap, token: Token) {
if let Some(mut connection) = connections.remove(&token) {
connection.close();
if let Err(err) = connection.deregister(poll){
if let Err(err) = connection.deregister(poll) {
::log::error!("couldn't deregister stream: {}", err);
}
}
}
// Close and remove inactive connections
pub fn remove_inactive_connections(
connections: &mut ConnectionMap,
){
pub fn remove_inactive_connections(connections: &mut ConnectionMap) {
let now = Instant::now();
connections.retain(|_, connection| {

22
aquatic_ws/src/lib/tasks.rs
View file

@ -4,11 +4,8 @@ use histogram::Histogram;
use crate::common::*;
pub fn clean_torrents(state: &State){
fn clean_torrent_map(
torrent_map: &mut TorrentMap,
){
pub fn clean_torrents(state: &State) {
fn clean_torrent_map(torrent_map: &mut TorrentMap) {
let now = Instant::now();
torrent_map.retain(|_, torrent_data| {
@ -22,10 +19,10 @@ pub fn clean_torrents(state: &State){
match peer.status {
PeerStatus::Seeding => {
*num_seeders -= 1;
},
}
PeerStatus::Leeching => {
*num_leechers -= 1;
},
}
_ => (),
};
}
@ -45,24 +42,23 @@ pub fn clean_torrents(state: &State){
clean_torrent_map(&mut torrent_maps.ipv6);
}
pub fn print_statistics(state: &State){
pub fn print_statistics(state: &State) {
let mut peers_per_torrent = Histogram::new();
{
let torrents = &mut state.torrent_maps.lock();
for torrent in torrents.ipv4.values(){
for torrent in torrents.ipv4.values() {
let num_peers = (torrent.num_seeders + torrent.num_leechers) as u64;
if let Err(err) = peers_per_torrent.increment(num_peers){
if let Err(err) = peers_per_torrent.increment(num_peers) {
eprintln!("error incrementing peers_per_torrent histogram: {}", err)
}
}
for torrent in torrents.ipv6.values(){
for torrent in torrents.ipv6.values() {
let num_peers = (torrent.num_seeders + torrent.num_leechers) as u64;
if let Err(err) = peers_per_torrent.increment(num_peers){
if let Err(err) = peers_per_torrent.increment(num_peers) {
eprintln!("error incrementing peers_per_torrent histogram: {}", err)
}
}

7
aquatic_ws_load_test/src/common.rs
View file

@ -1,10 +1,9 @@
use std::sync::{Arc, atomic::AtomicUsize};
use std::sync::{atomic::AtomicUsize, Arc};
use rand_distr::Pareto;
pub use aquatic_ws_protocol::*;
#[derive(Default)]
pub struct Statistics {
pub requests: AtomicUsize,
@ -15,7 +14,6 @@ pub struct Statistics {
pub responses_scrape: AtomicUsize,
}
#[derive(Clone)]
pub struct LoadTestState {
pub info_hashes: Arc<Vec<InfoHash>>,
@ -23,9 +21,8 @@ pub struct LoadTestState {
pub pareto: Arc<Pareto<f64>>,
}
#[derive(PartialEq, Eq, Clone, Copy)]
pub enum RequestType {
Announce,
Scrape
Scrape,
}

8
aquatic_ws_load_test/src/config.rs
View file

@ -1,7 +1,6 @@
use std::net::SocketAddr;
use serde::{Serialize, Deserialize};
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
@ -14,10 +13,8 @@ pub struct Config {
pub torrents: TorrentConfig,
}
impl aquatic_cli_helpers::Config for Config {}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct NetworkConfig {
@ -26,7 +23,6 @@ pub struct NetworkConfig {
pub poll_event_capacity: usize,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct TorrentConfig {
@ -46,7 +42,6 @@ pub struct TorrentConfig {
pub weight_scrape: usize,
}
impl Default for Config {
fn default() -> Self {
Self {
@ -70,7 +65,6 @@ impl Default for NetworkConfig {
}
}
impl Default for TorrentConfig {
fn default() -> Self {
Self {

65
aquatic_ws_load_test/src/main.rs
View file

@ -1,5 +1,5 @@
use std::sync::{atomic::Ordering, Arc};
use std::thread;
use std::sync::{Arc, atomic::Ordering};
use std::time::{Duration, Instant};
use rand::prelude::*;
@ -14,20 +14,17 @@ use common::*;
use config::*;
use network::*;
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
pub fn main(){
pub fn main() {
aquatic_cli_helpers::run_app_with_cli_and_config::<Config>(
"aquatic_ws_load_test: WebTorrent load tester",
run,
None
None,
)
}
fn run(config: Config) -> ::anyhow::Result<()> {
if config.torrents.weight_announce + config.torrents.weight_scrape == 0 {
panic!("Error: at least one weight must be larger than zero.");
@ -43,10 +40,7 @@ fn run(config: Config) -> ::anyhow::Result<()> {
info_hashes.push(InfoHash(rng.gen()));
}
let pareto = Pareto::new(
1.0,
config.torrents.torrent_selection_pareto_shape
).unwrap();
let pareto = Pareto::new(1.0, config.torrents.torrent_selection_pareto_shape).unwrap();
let state = LoadTestState {
info_hashes: Arc::new(info_hashes),
@ -58,22 +52,15 @@ fn run(config: Config) -> ::anyhow::Result<()> {
let config = config.clone();
let state = state.clone();
thread::spawn(move || run_socket_thread(&config, state,));
thread::spawn(move || run_socket_thread(&config, state));
}
monitor_statistics(
state,
&config
);
monitor_statistics(state, &config);
Ok(())
}
fn monitor_statistics(
state: LoadTestState,
config: &Config,
){
fn monitor_statistics(state: LoadTestState, config: &Config) {
let start_time = Instant::now();
let mut report_avg_response_vec: Vec<f64> = Vec::new();
@ -85,34 +72,36 @@ fn monitor_statistics(
let statistics = state.statistics.as_ref();
let responses_announce = statistics.responses_announce
.fetch_and(0, Ordering::SeqCst) as f64;
let responses_announce =
statistics.responses_announce.fetch_and(0, Ordering::SeqCst) as f64;
// let response_peers = statistics.response_peers
// .fetch_and(0, Ordering::SeqCst) as f64;
let requests_per_second = statistics.requests
.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_offer_per_second = statistics.responses_offer
.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_answer_per_second = statistics.responses_answer
.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_scrape_per_second = statistics.responses_scrape
.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let requests_per_second =
statistics.requests.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_offer_per_second =
statistics.responses_offer.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_answer_per_second =
statistics.responses_answer.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_scrape_per_second =
statistics.responses_scrape.fetch_and(0, Ordering::SeqCst) as f64 / interval_f64;
let responses_announce_per_second = responses_announce / interval_f64;
let responses_announce_per_second = responses_announce / interval_f64;
let responses_per_second =
responses_announce_per_second +
responses_offer_per_second +
responses_answer_per_second +
responses_scrape_per_second;
let responses_per_second = responses_announce_per_second
+ responses_offer_per_second
+ responses_answer_per_second
+ responses_scrape_per_second;
report_avg_response_vec.push(responses_per_second);
println!();
println!("Requests out: {:.2}/second", requests_per_second);
println!("Responses in: {:.2}/second", responses_per_second);
println!(" - Announce responses: {:.2}", responses_announce_per_second);
println!(
" - Announce responses: {:.2}",
responses_announce_per_second
);
println!(" - Offer responses: {:.2}", responses_offer_per_second);
println!(" - Answer responses: {:.2}", responses_answer_per_second);
println!(" - Scrape responses: {:.2}", responses_scrape_per_second);
@ -136,7 +125,7 @@ fn monitor_statistics(
config
);
break
break;
}
}
}

172
aquatic_ws_load_test/src/network.rs
View file

@ -1,26 +1,24 @@
use std::io::ErrorKind;
use std::sync::atomic::Ordering;
use std::time::Duration;
use std::io::ErrorKind;
use hashbrown::HashMap;
use mio::{net::TcpStream, Events, Poll, Interest, Token};
use rand::{rngs::SmallRng, prelude::*};
use tungstenite::{WebSocket, HandshakeError, ClientHandshake, handshake::MidHandshake};
use mio::{net::TcpStream, Events, Interest, Poll, Token};
use rand::{prelude::*, rngs::SmallRng};
use tungstenite::{handshake::MidHandshake, ClientHandshake, HandshakeError, WebSocket};
use crate::common::*;
use crate::config::*;
use crate::utils::create_random_request;
// Allow large enum variant WebSocket because it should be very common
#[allow(clippy::large_enum_variant)]
pub enum ConnectionState {
TcpStream(TcpStream),
WebSocket(WebSocket<TcpStream>),
MidHandshake(MidHandshake<ClientHandshake<TcpStream>>)
MidHandshake(MidHandshake<ClientHandshake<TcpStream>>),
}
impl ConnectionState {
fn advance(self, config: &Config) -> Option<Self> {
match self {
@ -31,33 +29,27 @@ impl ConnectionState {
config.server_address.port()
);
match ::tungstenite::client(req, stream){
Ok((ws, _)) => {
Some(ConnectionState::WebSocket(ws))
},
match ::tungstenite::client(req, stream) {
Ok((ws, _)) => Some(ConnectionState::WebSocket(ws)),
Err(HandshakeError::Interrupted(handshake)) => {
Some(ConnectionState::MidHandshake(handshake))
},
}
Err(HandshakeError::Failure(err)) => {
eprintln!("handshake error: {:?}", err);
None
}
}
},
Self::MidHandshake(handshake) => {
match handshake.handshake() {
Ok((ws, _)) => {
Some(ConnectionState::WebSocket(ws))
},
Err(HandshakeError::Interrupted(handshake)) => {
Some(ConnectionState::MidHandshake(handshake))
},
Err(HandshakeError::Failure(err)) => {
eprintln!("handshake error: {:?}", err);
}
Self::MidHandshake(handshake) => match handshake.handshake() {
Ok((ws, _)) => Some(ConnectionState::WebSocket(ws)),
Err(HandshakeError::Interrupted(handshake)) => {
Some(ConnectionState::MidHandshake(handshake))
}
Err(HandshakeError::Failure(err)) => {
eprintln!("handshake error: {:?}", err);
None
}
None
}
},
Self::WebSocket(ws) => Some(Self::WebSocket(ws)),
@ -65,7 +57,6 @@ impl ConnectionState {
}
}
pub struct Connection {
stream: ConnectionState,
peer_id: PeerId,
@ -73,7 +64,6 @@ pub struct Connection {
send_answer: Option<(PeerId, OfferId)>,
}
impl Connection {
pub fn create_and_register(
config: &Config,
@ -85,10 +75,14 @@ impl Connection {
let mut stream = TcpStream::connect(config.server_address)?;
poll.registry()
.register(&mut stream, Token(*token_counter), Interest::READABLE | Interest::WRITABLE)
.register(
&mut stream,
Token(*token_counter),
Interest::READABLE | Interest::WRITABLE,
)
.unwrap();
let connection = Connection {
let connection = Connection {
stream: ConnectionState::TcpStream(stream),
peer_id: PeerId(rng.gen()),
can_send: false,
@ -103,7 +97,7 @@ impl Connection {
}
pub fn advance(self, config: &Config) -> Option<Self> {
if let Some(stream) = self.stream.advance(config){
if let Some(stream) = self.stream.advance(config) {
let can_send = matches!(stream, ConnectionState::WebSocket(_));
Some(Self {
@ -117,52 +111,53 @@ impl Connection {
}
}
pub fn read_responses(
&mut self,
state: &LoadTestState,
) -> bool { // bool = drop connection
pub fn read_responses(&mut self, state: &LoadTestState) -> bool {
// bool = drop connection
if let ConnectionState::WebSocket(ref mut ws) = self.stream {
loop {
match ws.read_message(){
Ok(message) => {
match OutMessage::from_ws_message(message){
Ok(OutMessage::Offer(offer)) => {
state.statistics.responses_offer
.fetch_add(1, Ordering::SeqCst);
match ws.read_message() {
Ok(message) => match OutMessage::from_ws_message(message) {
Ok(OutMessage::Offer(offer)) => {
state
.statistics
.responses_offer
.fetch_add(1, Ordering::SeqCst);
self.send_answer = Some((
offer.peer_id,
offer.offer_id
));
self.send_answer = Some((offer.peer_id, offer.offer_id));
self.can_send = true;
},
Ok(OutMessage::Answer(_)) => {
state.statistics.responses_answer
.fetch_add(1, Ordering::SeqCst);
self.can_send = true;
}
Ok(OutMessage::Answer(_)) => {
state
.statistics
.responses_answer
.fetch_add(1, Ordering::SeqCst);
self.can_send = true;
},
Ok(OutMessage::AnnounceResponse(_)) => {
state.statistics.responses_announce
.fetch_add(1, Ordering::SeqCst);
self.can_send = true;
}
Ok(OutMessage::AnnounceResponse(_)) => {
state
.statistics
.responses_announce
.fetch_add(1, Ordering::SeqCst);
self.can_send = true;
},
Ok(OutMessage::ScrapeResponse(_)) => {
state.statistics.responses_scrape
.fetch_add(1, Ordering::SeqCst);
self.can_send = true;
}
Ok(OutMessage::ScrapeResponse(_)) => {
state
.statistics
.responses_scrape
.fetch_add(1, Ordering::SeqCst);
self.can_send = true;
},
Err(err) => {
eprintln!("error deserializing offer: {:?}", err);
}
self.can_send = true;
}
Err(err) => {
eprintln!("error deserializing offer: {:?}", err);
}
},
Err(tungstenite::Error::Io(err)) if err.kind() == ErrorKind::WouldBlock => {
return false;
},
}
Err(_) => {
return true;
}
@ -178,18 +173,14 @@ impl Connection {
config: &Config,
state: &LoadTestState,
rng: &mut impl Rng,
) -> bool { // bool = remove connection
) -> bool {
// bool = remove connection
if !self.can_send {
return false;
}
if let ConnectionState::WebSocket(ref mut ws) = self.stream {
let request = create_random_request(
&config,
&state,
rng,
self.peer_id
);
let request = create_random_request(&config, &state, rng, self.peer_id);
// If self.send_answer is set and request is announce request, make
// the request an offer answer
@ -211,20 +202,16 @@ impl Connection {
request
};
match ws.write_message(request.to_ws_message()){
match ws.write_message(request.to_ws_message()) {
Ok(()) => {
state.statistics.requests.fetch_add(1, Ordering::SeqCst);
self.can_send = false;
false
},
Err(tungstenite::Error::Io(err)) if err.kind() == ErrorKind::WouldBlock => {
false
}
Err(_) => {
true
}
Err(tungstenite::Error::Io(err)) if err.kind() == ErrorKind::WouldBlock => false,
Err(_) => true,
}
} else {
println!("send request can't send to non-ws stream");
@ -234,14 +221,9 @@ impl Connection {
}
}
pub type ConnectionMap = HashMap<usize, Connection>;
pub fn run_socket_thread(
config: &Config,
state: LoadTestState,
) {
pub fn run_socket_thread(config: &Config, state: LoadTestState) {
let timeout = Duration::from_micros(config.network.poll_timeout_microseconds);
let create_conn_interval = 2 ^ config.network.connection_creation_interval;
@ -259,11 +241,11 @@ pub fn run_socket_thread(
poll.poll(&mut events, Some(timeout))
.expect("failed polling");
for event in events.iter(){
for event in events.iter() {
let token = event.token();
if event.is_readable(){
if let Some(connection) = connections.get_mut(&token.0){
if event.is_readable() {
if let Some(connection) = connections.get_mut(&token.0) {
if let ConnectionState::WebSocket(_) = connection.stream {
let drop_connection = connection.read_responses(&state);
@ -276,26 +258,22 @@ pub fn run_socket_thread(
}
}
if let Some(connection) = connections.remove(&token.0){
if let Some(connection) = connection.advance(config){
if let Some(connection) = connections.remove(&token.0) {
if let Some(connection) = connection.advance(config) {
connections.insert(token.0, connection);
}
}
}
for (k, connection) in connections.iter_mut(){
let drop_connection = connection.send_request(
config,
&state,
&mut rng,
);
for (k, connection) in connections.iter_mut() {
let drop_connection = connection.send_request(config, &state, &mut rng);
if drop_connection {
drop_keys.push(*k)
}
}
for k in drop_keys.drain(..){
for k in drop_keys.drain(..) {
connections.remove(&k);
}

53
aquatic_ws_load_test/src/utils.rs
View file

@ -1,49 +1,33 @@
use std::sync::Arc;
use rand::distributions::WeightedIndex;
use rand_distr::Pareto;
use rand::prelude::*;
use rand_distr::Pareto;
use crate::common::*;
use crate::config::*;
pub fn create_random_request(
config: &Config,
state: &LoadTestState,
rng: &mut impl Rng,
peer_id: PeerId
peer_id: PeerId,
) -> InMessage {
let weights = [
config.torrents.weight_announce as u32,
config.torrents.weight_scrape as u32,
];
let items = [
RequestType::Announce,
RequestType::Scrape,
];
let items = [RequestType::Announce, RequestType::Scrape];
let dist = WeightedIndex::new(&weights)
.expect("random request weighted index");
let dist = WeightedIndex::new(&weights).expect("random request weighted index");
match items[dist.sample(rng)] {
RequestType::Announce => create_announce_request(
config,
state,
rng,
peer_id
),
RequestType::Scrape => create_scrape_request(
config,
state,
rng,
)
RequestType::Announce => create_announce_request(config, state, rng, peer_id),
RequestType::Scrape => create_scrape_request(config, state, rng),
}
}
#[inline]
fn create_announce_request(
config: &Config,
@ -61,9 +45,7 @@ fn create_announce_request(
let info_hash_index = select_info_hash_index(config, &state, rng);
let mut offers = Vec::with_capacity(
config.torrents.offers_per_request
);
let mut offers = Vec::with_capacity(config.torrents.offers_per_request);
for _ in 0..config.torrents.offers_per_request {
offers.push(AnnounceRequestOffer {
@ -88,13 +70,8 @@ fn create_announce_request(
})
}
#[inline]
fn create_scrape_request(
config: &Config,
state: &LoadTestState,
rng: &mut impl Rng,
) -> InMessage {
fn create_scrape_request(config: &Config, state: &LoadTestState, rng: &mut impl Rng) -> InMessage {
let mut scrape_hashes = Vec::with_capacity(5);
for _ in 0..5 {
@ -109,23 +86,13 @@ fn create_scrape_request(
})
}
#[inline]
fn select_info_hash_index(
config: &Config,
state: &LoadTestState,
rng: &mut impl Rng,
) -> usize {
fn select_info_hash_index(config: &Config, state: &LoadTestState, rng: &mut impl Rng) -> usize {
pareto_usize(rng, &state.pareto, config.torrents.number_of_torrents - 1)
}
#[inline]
fn pareto_usize(
rng: &mut impl Rng,
pareto: &Arc<Pareto<f64>>,
max: usize,
) -> usize {
fn pareto_usize(rng: &mut impl Rng, pareto: &Arc<Pareto<f64>>, max: usize) -> usize {
let p: f64 = pareto.sample(rng);
let p = (p.min(101.0f64) - 1.0) / 100.0;

36
aquatic_ws_protocol/benches/bench_deserialize_announce_request.rs
View file

@ -1,27 +1,25 @@
use std::time::Duration;
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use std::time::Duration;
use aquatic_ws_protocol::*;
pub fn bench(c: &mut Criterion) {
let info_hash = InfoHash([
b'a', b'b', b'c', b'd', b'e',
b'?', b'\n', b'1', b'2', b'3',
0, 1, 2, 3, 4,
0, 1, 2, 3, 4,
b'a', b'b', b'c', b'd', b'e', b'?', b'\n', b'1', b'2', b'3', 0, 1, 2, 3, 4, 0, 1, 2, 3, 4,
]);
let peer_id = PeerId(info_hash.0);
let offers: Vec<AnnounceRequestOffer> = (0..10).map(|i| {
let mut offer_id = OfferId(info_hash.0);
let offers: Vec<AnnounceRequestOffer> = (0..10)
.map(|i| {
let mut offer_id = OfferId(info_hash.0);
offer_id.0[i] = i as u8;
offer_id.0[i] = i as u8;
AnnounceRequestOffer {
offer: JsonValue(::serde_json::json!({ "sdp": "abcdef" })),
offer_id,
}
}).collect();
AnnounceRequestOffer {
offer: JsonValue(::serde_json::json!({ "sdp": "abcdef" })),
offer_id,
}
})
.collect();
let offers_len = offers.len();
let request = InMessage::AnnounceRequest(AnnounceRequest {
@ -34,17 +32,17 @@ pub fn bench(c: &mut Criterion) {
numwant: Some(offers_len),
answer: Some(JsonValue(::serde_json::json!({ "sdp": "abcdef" }))),
to_peer_id: Some(peer_id),
offer_id: Some(OfferId(info_hash.0))
offer_id: Some(OfferId(info_hash.0)),
});
let ws_message = request.to_ws_message();
c.bench_function("deserialize-announce-request", |b| b.iter(||
InMessage::from_ws_message(black_box(ws_message.clone()))
));
c.bench_function("deserialize-announce-request", |b| {
b.iter(|| InMessage::from_ws_message(black_box(ws_message.clone())))
});
}
criterion_group!{
criterion_group! {
name = benches;
config = Criterion::default()
.sample_size(1000)

147
aquatic_ws_protocol/src/lib.rs
View file

@ -1,28 +1,23 @@
use anyhow::Context;
use hashbrown::HashMap;
use serde::{Serialize, Deserialize, Serializer, Deserializer};
use serde::{Deserialize, Deserializer, Serialize, Serializer};
mod serde_helpers;
use serde_helpers::*;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AnnounceAction;
impl Serialize for AnnounceAction {
fn serialize<S>(
&self,
serializer: S
) -> Result<S::Ok, S::Error>
where S: Serializer
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str("announce")
}
}
impl<'de> Deserialize<'de> for AnnounceAction {
fn deserialize<D>(deserializer: D) -> Result<AnnounceAction, D::Error>
where
@ -32,23 +27,18 @@ impl<'de> Deserialize<'de> for AnnounceAction {
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ScrapeAction;
impl Serialize for ScrapeAction {
fn serialize<S>(
&self,
serializer: S
) -> Result<S::Ok, S::Error>
where S: Serializer
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str("scrape")
}
}
impl<'de> Deserialize<'de> for ScrapeAction {
fn deserialize<D>(deserializer: D) -> Result<ScrapeAction, D::Error>
where
@ -58,7 +48,6 @@ impl<'de> Deserialize<'de> for ScrapeAction {
}
}
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct PeerId(
@ -66,10 +55,9 @@ pub struct PeerId(
deserialize_with = "deserialize_20_bytes",
serialize_with = "serialize_20_bytes"
)]
pub [u8; 20]
pub [u8; 20],
);
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct InfoHash(
@ -77,10 +65,9 @@ pub struct InfoHash(
deserialize_with = "deserialize_20_bytes",
serialize_with = "serialize_20_bytes"
)]
pub [u8; 20]
pub [u8; 20],
);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct OfferId(
@ -88,33 +75,29 @@ pub struct OfferId(
deserialize_with = "deserialize_20_bytes",
serialize_with = "serialize_20_bytes"
)]
pub [u8; 20]
pub [u8; 20],
);
/// Some kind of nested structure from https://www.npmjs.com/package/simple-peer
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct JsonValue(pub ::serde_json::Value);
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum AnnounceEvent {
Started,
Stopped,
Completed,
Update
Update,
}
impl Default for AnnounceEvent {
fn default() -> Self {
Self::Update
}
}
/// Apparently, these are sent to a number of peers when they are set
/// in an AnnounceRequest
/// action = "announce"
@ -131,7 +114,6 @@ pub struct MiddlemanOfferToPeer {
pub offer_id: OfferId,
}
/// If announce request has answer = true, send this to peer with
/// peer id == "to_peer_id" field
/// Action field should be 'announce'
@ -145,7 +127,6 @@ pub struct MiddlemanAnswerToPeer {
pub offer_id: OfferId,
}
/// Element of AnnounceRequest.offers
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AnnounceRequestOffer {
@ -153,7 +134,6 @@ pub struct AnnounceRequestOffer {
pub offer_id: OfferId,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AnnounceRequest {
pub action: AnnounceAction,
@ -189,7 +169,6 @@ pub struct AnnounceRequest {
pub offer_id: Option<OfferId>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AnnounceResponse {
pub action: AnnounceAction,
@ -201,7 +180,6 @@ pub struct AnnounceResponse {
pub announce_interval: usize, // Default 2 min probably
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ScrapeRequestInfoHashes {
@ -209,7 +187,6 @@ pub enum ScrapeRequestInfoHashes {
Multiple(Vec<InfoHash>),
}
impl ScrapeRequestInfoHashes {
pub fn as_vec(self) -> Vec<InfoHash> {
match self {
@ -219,7 +196,6 @@ impl ScrapeRequestInfoHashes {
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ScrapeRequest {
pub action: ScrapeAction,
@ -230,7 +206,6 @@ pub struct ScrapeRequest {
pub info_hashes: Option<ScrapeRequestInfoHashes>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ScrapeStatistics {
pub complete: usize,
@ -238,7 +213,6 @@ pub struct ScrapeStatistics {
pub downloaded: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ScrapeResponse {
pub action: ScrapeAction,
@ -247,7 +221,6 @@ pub struct ScrapeResponse {
// pub flags: HashMap<String, usize>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum InMessage {
@ -255,7 +228,6 @@ pub enum InMessage {
ScrapeRequest(ScrapeRequest),
}
impl InMessage {
#[inline]
pub fn to_ws_message(&self) -> ::tungstenite::Message {
@ -263,9 +235,7 @@ impl InMessage {
}
#[inline]
pub fn from_ws_message(
ws_message: tungstenite::Message
) -> ::anyhow::Result<Self> {
pub fn from_ws_message(ws_message: tungstenite::Message) -> ::anyhow::Result<Self> {
use tungstenite::Message::Text;
let mut text = if let Text(text) = ws_message {
@ -274,12 +244,10 @@ impl InMessage {
return Err(anyhow::anyhow!("Message is not text"));
};
return ::simd_json::serde::from_str(&mut text)
.context("deserialize with serde");
return ::simd_json::serde::from_str(&mut text).context("deserialize with serde");
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum OutMessage {
@ -289,7 +257,6 @@ pub enum OutMessage {
ScrapeResponse(ScrapeResponse),
}
impl OutMessage {
#[inline]
pub fn to_ws_message(&self) -> tungstenite::Message {
@ -297,10 +264,8 @@ impl OutMessage {
}
#[inline]
pub fn from_ws_message(
message: ::tungstenite::Message
) -> ::anyhow::Result<Self> {
use tungstenite::Message::{Text, Binary};
pub fn from_ws_message(message: ::tungstenite::Message) -> ::anyhow::Result<Self> {
use tungstenite::Message::{Binary, Text};
let mut text = match message {
Text(text) => text,
@ -312,7 +277,6 @@ impl OutMessage {
}
}
#[cfg(test)]
mod tests {
use quickcheck::Arbitrary;
@ -354,7 +318,7 @@ mod tests {
impl Arbitrary for AnnounceEvent {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
match (bool::arbitrary(g), bool::arbitrary(g)){
match (bool::arbitrary(g), bool::arbitrary(g)) {
(false, false) => Self::Started,
(true, false) => Self::Started,
(false, true) => Self::Completed,
@ -370,7 +334,7 @@ mod tests {
peer_id: Arbitrary::arbitrary(g),
info_hash: Arbitrary::arbitrary(g),
offer_id: Arbitrary::arbitrary(g),
offer: sdp_json_value()
offer: sdp_json_value(),
}
}
}
@ -382,7 +346,7 @@ mod tests {
peer_id: Arbitrary::arbitrary(g),
info_hash: Arbitrary::arbitrary(g),
offer_id: Arbitrary::arbitrary(g),
answer: sdp_json_value()
answer: sdp_json_value(),
}
}
}
@ -391,7 +355,7 @@ mod tests {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
Self {
offer_id: Arbitrary::arbitrary(g),
offer: sdp_json_value()
offer: sdp_json_value(),
}
}
}
@ -408,17 +372,16 @@ mod tests {
match has_offers_or_answer_or_neither {
Some(true) => {
offers = Some(Arbitrary::arbitrary(g));
},
}
Some(false) => {
answer = Some(sdp_json_value());
to_peer_id = Some(Arbitrary::arbitrary(g));
offer_id = Some(Arbitrary::arbitrary(g));
},
}
None => (),
}
let numwant = offers.as_ref()
.map(|offers| offers.len());
let numwant = offers.as_ref().map(|offers| offers.len());
Self {
action: AnnounceAction,
@ -456,7 +419,6 @@ mod tests {
}
}
impl Arbitrary for ScrapeRequestInfoHashes {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
if Arbitrary::arbitrary(g) {
@ -490,7 +452,7 @@ mod tests {
impl Arbitrary for InMessage {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
if Arbitrary::arbitrary(g){
if Arbitrary::arbitrary(g) {
Self::AnnounceRequest(Arbitrary::arbitrary(g))
} else {
Self::ScrapeRequest(Arbitrary::arbitrary(g))
@ -500,7 +462,7 @@ mod tests {
impl Arbitrary for OutMessage {
fn arbitrary(g: &mut quickcheck::Gen) -> Self {
match (Arbitrary::arbitrary(g), Arbitrary::arbitrary(g)){
match (Arbitrary::arbitrary(g), Arbitrary::arbitrary(g)) {
(false, false) => Self::AnnounceResponse(Arbitrary::arbitrary(g)),
(true, false) => Self::ScrapeResponse(Arbitrary::arbitrary(g)),
(false, true) => Self::Offer(Arbitrary::arbitrary(g)),
@ -513,11 +475,9 @@ mod tests {
fn quickcheck_serde_identity_in_message(in_message_1: InMessage) -> bool {
let ws_message = in_message_1.to_ws_message();
let in_message_2 = InMessage::from_ws_message(
ws_message.clone()
).unwrap();
let in_message_2 = InMessage::from_ws_message(ws_message.clone()).unwrap();
let success = in_message_1 == in_message_2;
let success = in_message_1 == in_message_2;
if !success {
dbg!(in_message_1);
@ -534,11 +494,9 @@ mod tests {
fn quickcheck_serde_identity_out_message(out_message_1: OutMessage) -> bool {
let ws_message = out_message_1.to_ws_message();
let out_message_2 = OutMessage::from_ws_message(
ws_message.clone()
).unwrap();
let out_message_2 = OutMessage::from_ws_message(ws_message.clone()).unwrap();
let success = out_message_1 == out_message_2;
let success = out_message_1 == out_message_2;
if !success {
dbg!(out_message_1);
@ -562,22 +520,21 @@ mod tests {
}
#[test]
fn test_deserialize_info_hashes_vec(){
fn test_deserialize_info_hashes_vec() {
let mut input: String = r#"{
"action": "scrape",
"info_hash": ["aaaabbbbccccddddeeee", "aaaabbbbccccddddeeee"]
}"#.into();
}"#
.into();
let info_hashes = ScrapeRequestInfoHashes::Multiple(
vec![
info_hash_from_bytes(b"aaaabbbbccccddddeeee"),
info_hash_from_bytes(b"aaaabbbbccccddddeeee"),
]
);
let info_hashes = ScrapeRequestInfoHashes::Multiple(vec![
info_hash_from_bytes(b"aaaabbbbccccddddeeee"),
info_hash_from_bytes(b"aaaabbbbccccddddeeee"),
]);
let expected = ScrapeRequest {
action: ScrapeAction,
info_hashes: Some(info_hashes)
info_hashes: Some(info_hashes),
};
let observed: ScrapeRequest = ::simd_json::serde::from_str(&mut input).unwrap();
@ -586,19 +543,19 @@ mod tests {
}
#[test]
fn test_deserialize_info_hashes_str(){
fn test_deserialize_info_hashes_str() {
let mut input: String = r#"{
"action": "scrape",
"info_hash": "aaaabbbbccccddddeeee"
}"#.into();
}"#
.into();
let info_hashes = ScrapeRequestInfoHashes::Single(
info_hash_from_bytes(b"aaaabbbbccccddddeeee")
);
let info_hashes =
ScrapeRequestInfoHashes::Single(info_hash_from_bytes(b"aaaabbbbccccddddeeee"));
let expected = ScrapeRequest {
action: ScrapeAction,
info_hashes: Some(info_hashes)
info_hashes: Some(info_hashes),
};
let observed: ScrapeRequest = ::simd_json::serde::from_str(&mut input).unwrap();
@ -607,15 +564,16 @@ mod tests {
}
#[test]
fn test_deserialize_info_hashes_null(){
fn test_deserialize_info_hashes_null() {
let mut input: String = r#"{
"action": "scrape",
"info_hash": null
}"#.into();
}"#
.into();
let expected = ScrapeRequest {
action: ScrapeAction,
info_hashes: None
info_hashes: None,
};
let observed: ScrapeRequest = ::simd_json::serde::from_str(&mut input).unwrap();
@ -624,14 +582,15 @@ mod tests {
}
#[test]
fn test_deserialize_info_hashes_missing(){
fn test_deserialize_info_hashes_missing() {
let mut input: String = r#"{
"action": "scrape"
}"#.into();
}"#
.into();
let expected = ScrapeRequest {
action: ScrapeAction,
info_hashes: None
info_hashes: None,
};
let observed: ScrapeRequest = ::simd_json::serde::from_str(&mut input).unwrap();
@ -645,12 +604,12 @@ mod tests {
println!("{}", json);
let deserialized: ScrapeRequestInfoHashes = ::simd_json::serde::from_str(&mut json).unwrap();
let deserialized: ScrapeRequestInfoHashes =
::simd_json::serde::from_str(&mut json).unwrap();
let success = info_hashes == deserialized;
if !success {
}
if !success {}
success
}

45
aquatic_ws_protocol/src/serde_helpers.rs
View file

@ -1,11 +1,9 @@
use serde::{Serializer, Deserializer, de::Visitor};
use serde::{de::Visitor, Deserializer, Serializer};
use super::{AnnounceAction, ScrapeAction};
pub struct AnnounceActionVisitor;
impl<'de> Visitor<'de> for AnnounceActionVisitor {
type Value = AnnounceAction;
@ -14,7 +12,9 @@ impl<'de> Visitor<'de> for AnnounceActionVisitor {
}
fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
where E: ::serde::de::Error, {
where
E: ::serde::de::Error,
{
if v == "announce" {
Ok(AnnounceAction)
} else {
@ -23,10 +23,8 @@ impl<'de> Visitor<'de> for AnnounceActionVisitor {
}
}
pub struct ScrapeActionVisitor;
impl<'de> Visitor<'de> for ScrapeActionVisitor {
type Value = ScrapeAction;
@ -35,7 +33,9 @@ impl<'de> Visitor<'de> for ScrapeActionVisitor {
}
fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
where E: ::serde::de::Error, {
where
E: ::serde::de::Error,
{
if v == "scrape" {
Ok(ScrapeAction)
} else {
@ -44,17 +44,15 @@ impl<'de> Visitor<'de> for ScrapeActionVisitor {
}
}
pub fn serialize_20_bytes<S>(
data: &[u8; 20],
serializer: S
) -> Result<S::Ok, S::Error> where S: Serializer {
pub fn serialize_20_bytes<S>(data: &[u8; 20], serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let text: String = data.iter().map(|byte| char::from(*byte)).collect();
serializer.serialize_str(&text)
}
struct TwentyByteVisitor;
impl<'de> Visitor<'de> for TwentyByteVisitor {
@ -66,7 +64,8 @@ impl<'de> Visitor<'de> for TwentyByteVisitor {
#[inline]
fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
where E: ::serde::de::Error,
where
E: ::serde::de::Error,
{
// Value is encoded in nodejs reference client something as follows:
// ```
@ -83,8 +82,8 @@ impl<'de> Visitor<'de> for TwentyByteVisitor {
let mut arr = [0u8; 20];
let mut char_iter = value.chars();
for a in arr.iter_mut(){
if let Some(c) = char_iter.next(){
for a in arr.iter_mut() {
if let Some(c) = char_iter.next() {
if c as u32 > 255 {
return Err(E::custom(format!(
"character not in single byte range: {:#?}",
@ -102,17 +101,14 @@ impl<'de> Visitor<'de> for TwentyByteVisitor {
}
}
#[inline]
pub fn deserialize_20_bytes<'de, D>(
deserializer: D
) -> Result<[u8; 20], D::Error>
where D: Deserializer<'de>
pub fn deserialize_20_bytes<'de, D>(deserializer: D) -> Result<[u8; 20], D::Error>
where
D: Deserializer<'de>,
{
deserializer.deserialize_any(TwentyByteVisitor)
}
#[cfg(test)]
mod tests {
use quickcheck_macros::quickcheck;
@ -130,7 +126,7 @@ mod tests {
}
#[test]
fn test_deserialize_20_bytes(){
fn test_deserialize_20_bytes() {
let mut input = r#""aaaabbbbccccddddeeee""#.to_string();
let expected = info_hash_from_bytes(b"aaaabbbbccccddddeeee");
@ -150,7 +146,7 @@ mod tests {
}
#[test]
fn test_serde_20_bytes(){
fn test_serde_20_bytes() {
let info_hash = info_hash_from_bytes(b"aaaabbbbccccddddeeee");
let mut out = ::simd_json::serde::to_string(&info_hash).unwrap();
@ -166,5 +162,4 @@ mod tests {
info_hash == info_hash_2
}
}