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

Merge branch 'master' into ws-file-transfer-ci

Browse source
This commit is contained in:
Joakim Frostegård 2021-08-16 19:27:13 +02:00
commit 82a36041b3
78 changed files with 2044 additions and 2910 deletions
Download patch file Download diff file Expand all files Collapse all files

556
Cargo.lock generated
View file

File diff suppressed because it is too large Load diff

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

2
aquatic_cli_helpers/Cargo.toml
View file

@ -10,5 +10,5 @@ repository = "https://github.com/greatest-ape/aquatic"
[dependencies]
anyhow = "1"
serde = { version = "1", features = ["derive"] }
simplelog = "0.9"
simplelog = "0.10.0"
toml = "0.5"

102
aquatic_cli_helpers/src/lib.rs
View file

@ -2,9 +2,8 @@ use std::fs::File;
use std::io::Read;
use anyhow::Context;
use serde::{Serialize, Deserialize, de::DeserializeOwned};
use simplelog::{ConfigBuilder, LevelFilter, TermLogger, TerminalMode};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use simplelog::{ColorChoice, ConfigBuilder, LevelFilter, TermLogger, TerminalMode};
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
@ -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,
@ -209,8 +193,10 @@ fn start_logger(log_level: LogLevel) -> ::anyhow::Result<()> {
TermLogger::init(
level_filter,
simplelog_config,
TerminalMode::Stderr
).context("Couldn't initialize logger")?;
TerminalMode::Stderr,
ColorChoice::Auto,
)
.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

4
aquatic_http/Cargo.toml
View file

@ -22,7 +22,7 @@ aquatic_common = "0.1.0"
aquatic_http_protocol = "0.1.0"
crossbeam-channel = "0.5"
either = "1"
hashbrown = "0.9"
hashbrown = "0.11.2"
histogram = "0.6"
indexmap = "1"
itoa = "0.4"
@ -36,7 +36,7 @@ privdrop = "0.5"
rand = { version = "0.8", features = ["small_rng"] }
serde = { version = "1", features = ["derive"] }
smartstring = "0.2"
socket2 = { version = "0.3", features = ["reuseport"] }
socket2 = { version = "0.4.1", features = ["all"] }
[dev-dependencies]
quickcheck = "1.0"

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),

49
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,31 +29,35 @@ 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(){
Socket::new(Domain::ipv4(), Type::stream(), Some(Protocol::tcp()))
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")?;
Socket::new(Domain::IPV6, Type::STREAM, Some(Protocol::TCP))
}
.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())
Ok(builder.into())
}

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)
}
}

2
aquatic_http_load_test/Cargo.toml
View file

@ -13,7 +13,7 @@ name = "aquatic_http_load_test"
anyhow = "1"
aquatic_cli_helpers = "0.1.0"
aquatic_http_protocol = "0.1.0"
hashbrown = "0.9"
hashbrown = "0.11.2"
mimalloc = { version = "0.1", default-features = false }
mio = { version = "0.7", features = ["udp", "os-poll", "os-util"] }
rand = { version = "0.8", features = ["small_rng"] }

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;

9
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"
@ -28,7 +23,7 @@ harness = false
[dependencies]
anyhow = "1"
hashbrown = "0.9"
hashbrown = "0.11.2"
hex = { version = "0.4", default-features = false }
httparse = "1"
itoa = "0.4"
@ -38,7 +33,7 @@ rand = { version = "0.8", features = ["small_rng"] }
serde = { version = "1", features = ["derive"] }
serde_bencode = "0.2"
smartstring = "0.2"
urlencoding = "1"
urlencoding = "2.1.0"
[dev-dependencies]
bendy = { version = "0.3", features = ["std", "serde"] }

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
}

4
aquatic_udp/Cargo.toml
View file

@ -20,7 +20,7 @@ aquatic_cli_helpers = "0.1.0"
aquatic_common = "0.1.0"
aquatic_udp_protocol = "0.1.0"
crossbeam-channel = "0.5"
hashbrown = "0.9"
hashbrown = "0.11.2"
histogram = "0.6"
indexmap = "1"
log = "0.4"
@ -30,7 +30,7 @@ parking_lot = "0.11"
privdrop = "0.5"
rand = { version = "0.8", features = ["small_rng"] }
serde = { version = "1", features = ["derive"] }
socket2 = { version = "0.3", features = ["reuseport"] }
socket2 = { version = "0.4.1", features = ["all"] }
[dev-dependencies]
quickcheck = "1.0"

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)

108
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(){
Socket::new(Domain::ipv4(), Type::dgram(), Some(Protocol::udp()))
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");
Socket::new(Domain::IPV6, Type::DGRAM, Some(Protocol::UDP))
}
.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,
@ -111,10 +116,9 @@ fn create_socket(config: &Config) -> ::std::net::UdpSocket {
}
}
socket.into_udp_socket()
socket.into()
}
#[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)
}
}

2
aquatic_udp_bench/Cargo.toml
View file

@ -14,7 +14,7 @@ anyhow = "1"
aquatic_cli_helpers = "0.1.0"
aquatic_udp = "0.1.0"
crossbeam-channel = "0.5"
indicatif = "0.15"
indicatif = "0.16.2"
mimalloc = { version = "0.1", default-features = false }
num-format = "0.4"
rand = { version = "0.8", features = ["small_rng"] }

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 {

4
aquatic_udp_load_test/Cargo.toml
View file

@ -14,14 +14,14 @@ anyhow = "1"
aquatic_cli_helpers = "0.1.0"
aquatic_udp_protocol = "0.1.0"
crossbeam-channel = "0.5"
hashbrown = "0.9"
hashbrown = "0.11.2"
mimalloc = { version = "0.1", default-features = false }
mio = { version = "0.7", features = ["udp", "os-poll", "os-util"] }
parking_lot = "0.11"
rand = { version = "0.8", features = ["small_rng"] }
rand_distr = "0.4"
serde = { version = "1", features = ["derive"] }
socket2 = { version = "0.3", features = ["reuseport"] }
socket2 = { version = "0.4.1", features = ["all"] }
[dev-dependencies]
quickcheck = "1.0"

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;
}
}
}

107
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(){
Socket::new(Domain::ipv4(), Type::dgram(), Some(Protocol::udp()))
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");
Socket::new(Domain::IPV6, Type::DGRAM, Some(Protocol::UDP))
}
.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()
socket.into()
}
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)
}
}

4
aquatic_ws/Cargo.toml
View file

@ -22,7 +22,7 @@ aquatic_common = "0.1.0"
aquatic_ws_protocol = "0.1.0"
crossbeam-channel = "0.5"
either = "1"
hashbrown = { version = "0.9", features = ["serde"] }
hashbrown = { version = "0.11.2", features = ["serde"] }
histogram = "0.6"
indexmap = "1"
log = "0.4"
@ -33,7 +33,7 @@ parking_lot = "0.11"
privdrop = "0.5"
rand = { version = "0.8", features = ["small_rng"] }
serde = { version = "1", features = ["derive"] }
socket2 = { version = "0.3", features = ["reuseport"] }
socket2 = { version = "0.4.1", features = ["all"] }
tungstenite = "0.13"
[dev-dependencies]

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;
}
}
@ -129,7 +124,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 {
@ -143,24 +138,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;
},
}
_ => {}
}
}
@ -168,8 +161,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 {
@ -181,12 +173,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,
@ -197,7 +187,7 @@ pub fn handle_announce_requests(
out_message_sender.send(
offer_receiver.connection_meta,
OutMessage::Offer(middleman_offer)
OutMessage::Offer(middleman_offer),
);
::log::trace!(
"sent middleman offer to {:?}",
@ -211,9 +201,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,
@ -224,7 +212,7 @@ pub fn handle_announce_requests(
out_message_sender.send(
answer_receiver.connection_meta,
OutMessage::Answer(middleman_answer)
OutMessage::Answer(middleman_answer),
);
::log::trace!(
"sent middleman answer to {:?}",
@ -247,31 +235,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
@ -279,8 +264,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")?;

109
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())
},
}
}
}
@ -162,7 +148,7 @@ impl HandshakeMachine {
},
Err(native_tls::HandshakeError::WouldBlock(handshake)) => {
(Some(Either::Right(Self::TlsMidHandshake(handshake))), true)
},
}
Err(native_tls::HandshakeError::Failure(err)) => {
info!("tls handshake error: {}", err);
@ -173,7 +159,7 @@ 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) => {
@ -190,10 +176,11 @@ impl HandshakeMachine {
};
(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);
@ -203,20 +190,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).
///
@ -229,15 +213,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)),
}
}
@ -260,15 +240,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)
@ -277,19 +254,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)
}
}
}
@ -299,24 +273,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);
}
}
}
}

52
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(){
Socket::new(Domain::ipv4(), Type::stream(), Some(Protocol::tcp()))
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")?;
Socket::new(Domain::IPV6, Type::STREAM, Some(Protocol::TCP))
}
.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())
Ok(builder.into())
}
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)
}
}

2
aquatic_ws_load_test/Cargo.toml
View file

@ -13,7 +13,7 @@ name = "aquatic_ws_load_test"
anyhow = "1"
aquatic_cli_helpers = "0.1.0"
aquatic_ws_protocol = "0.1.0"
hashbrown = { version = "0.9", features = ["serde"] }
hashbrown = { version = "0.11.2", features = ["serde"] }
mimalloc = { version = "0.1", default-features = false }
mio = { version = "0.7", features = ["udp", "os-poll", "os-util"] }
rand = { version = "0.8", features = ["small_rng"] }

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;

4
aquatic_ws_protocol/Cargo.toml
View file

@ -18,10 +18,10 @@ harness = false
[dependencies]
anyhow = "1"
hashbrown = { version = "0.9", features = ["serde"] }
hashbrown = { version = "0.11.2", features = ["serde"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
simd-json = { version = "0.3", features = ["allow-non-simd"] }
simd-json = { version = "0.4.7", features = ["allow-non-simd"] }
tungstenite = "0.13"
[dev-dependencies]

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
}
}