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

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]

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

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

23
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 {
@ -185,4 +168,4 @@ impl Default for PrivilegeConfig {
user: "nobody".to_string(),
}
}
}
}

93
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
@ -294,4 +279,4 @@ pub fn handle_scrape_requests(
out_message_sender.send(meta, OutMessage::ScrapeResponse(response));
wake_socket_workers[meta.worker_index] = true;
}
}
}

83
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")?;
@ -169,4 +164,4 @@ pub fn create_tls_acceptor(
} else {
Ok(None)
}
}
}

111
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>;
pub type ConnectionMap = HashMap<Token, Connection>;

134
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),
@ -88,10 +83,10 @@ pub fn run_poll_loop(
loop {
poll.poll(&mut events, Some(poll_timeout))
.expect("failed polling");
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,
@ -211,38 +200,31 @@ 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);
remove_connection_if_exists(poll, connections, poll_token);
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,57 +234,49 @@ 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 {
if established_ws.peer_addr != meta.naive_peer_addr {
info!("socket worker error: peer socket addrs didn't match");
continue;
}
use ::tungstenite::Error::Io;
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);
}
}
}
}
}
}

54
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| {
@ -69,4 +63,4 @@ pub fn remove_inactive_connections(
});
connections.shrink_to_fit();
}
}

24
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)
}
}
@ -80,4 +76,4 @@ pub fn print_statistics(state: &State){
peers_per_torrent.maximum().unwrap(),
);
}
}
}