YGGverse/aquatic
1
0
Fork 0
mirror of https://github.com/YGGverse/aquatic.git synced 2026-04-01 18:25:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

aquatic_ws: mio: replace native_tls with rustls, rewrite connection state logic completely (#38)

Browse source 
* ws: mio: use rustls, rewrite Connection logic

* ws: mio: improve poll register/deregister handling

* ws: mio: work on type-level poll registry safety

* ws: mio: use stronger poll registry type-level guarantees

* ws: mio: fix stream reading

* ws: clean up, run fmt

* ws: mio: don't require registered connection for writing, improve docs

* ws: mio: add Connection::get_meta(), make Connection::meta private

* ws: mio: add ConnectionMap struct; remove utils.rs

* ws: mio: move token counter into ConnectionMap, improve docs

* ws: mio: connection: move Connection struct above state structs

* Update TODO

* ws: fix build errors

* ws: upgrade to tungstenite 0.16

* ws load test: don't panic on Close message; print shorter errors

* ws: fix socket worker bugs, add log statements

* ws: mio: wait for write availability if would block for ws messages

* Update README

* ws: mio: limit channels & queues; read 1 message only; other fixes

* ws: mio: send local responses each event; decrease channel size

* Update TODO

* ws: mio: limit ws send queue, fixing memory leak; limit pending messages

Also change some log output levels and run rustfmt

* Update TODO

* Update TODO
This commit is contained in:
Joakim Frostegård 2021-12-16 15:09:36 +00:00 • committed by GitHub
parent 222fac0e09
commit 67c4c02bbd
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
13 changed files with 893 additions and 776 deletions
Download patch file Download diff file Expand all files Collapse all files

40
aquatic_ws/src/mio/mod.rs
View file

@ -1,5 +1,3 @@
use std::fs::File;
use std::io::Read;
use std::sync::Arc;
use std::thread::Builder;
use std::time::Duration;
@ -9,7 +7,6 @@ use anyhow::Context;
use aquatic_common::cpu_pinning::{pin_current_if_configured_to, WorkerIndex};
use histogram::Histogram;
use mio::{Poll, Waker};
use native_tls::{Identity, TlsAcceptor};
use parking_lot::Mutex;
use privdrop::PrivDrop;
@ -17,11 +14,13 @@ pub mod common;
pub mod request;
pub mod socket;
use crate::config::Config;
use crate::{common::create_tls_config, config::Config};
use common::*;
pub const APP_NAME: &str = "aquatic_ws: WebTorrent tracker";
const SHARED_IN_CHANNEL_SIZE: usize = 1024;
pub fn run(config: Config, state: State) -> anyhow::Result<()> {
start_workers(config.clone(), state.clone()).expect("couldn't start workers");
@ -44,9 +43,10 @@ pub fn run(config: Config, state: State) -> anyhow::Result<()> {
}
pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
let opt_tls_acceptor = create_tls_acceptor(&config)?;
let tls_config = Arc::new(create_tls_config(&config)?);
let (in_message_sender, in_message_receiver) = ::crossbeam_channel::unbounded();
let (in_message_sender, in_message_receiver) =
::crossbeam_channel::bounded(SHARED_IN_CHANNEL_SIZE);
let mut out_message_senders = Vec::new();
let mut wakers = Vec::new();
@ -66,11 +66,12 @@ pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
let state = state.clone();
let socket_worker_statuses = socket_worker_statuses.clone();
let in_message_sender = in_message_sender.clone();
let opt_tls_acceptor = opt_tls_acceptor.clone();
let tls_config = tls_config.clone();
let poll = Poll::new()?;
let waker = Arc::new(Waker::new(poll.registry(), CHANNEL_TOKEN)?);
let (out_message_sender, out_message_receiver) = ::crossbeam_channel::unbounded();
let (out_message_sender, out_message_receiver) =
::crossbeam_channel::bounded(SHARED_IN_CHANNEL_SIZE * 16);
out_message_senders.push(out_message_sender);
wakers.push(waker);
@ -93,7 +94,7 @@ pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
poll,
in_message_sender,
out_message_receiver,
opt_tls_acceptor,
tls_config,
);
})?;
}
@ -180,27 +181,6 @@ pub fn start_workers(config: Config, state: State) -> anyhow::Result<()> {
Ok(())
}
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)
.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.network.tls_pkcs12_password)
.context("Couldn't parse pkcs12 identity file")?;
let acceptor = TlsAcceptor::new(identity)
.context("Couldn't create TlsAcceptor from pkcs12 identity")?;
Ok(Some(acceptor))
} else {
Ok(None)
}
}
fn print_statistics(state: &State) {
let mut peers_per_torrent = Histogram::new();

789
aquatic_ws/src/mio/socket/connection.rs
View file

@ -1,298 +1,577 @@
use std::io::{Read, Write};
use std::net::SocketAddr;
use std::{collections::VecDeque, io::ErrorKind, marker::PhantomData, net::Shutdown, sync::Arc};
use either::Either;
use hashbrown::HashMap;
use log::info;
use mio::net::TcpStream;
use mio::{Poll, Token};
use native_tls::{MidHandshakeTlsStream, TlsAcceptor, TlsStream};
use tungstenite::handshake::{server::NoCallback, HandshakeError, MidHandshake};
use tungstenite::protocol::WebSocketConfig;
use tungstenite::ServerHandshake;
use tungstenite::WebSocket;
use aquatic_common::ValidUntil;
use aquatic_ws_protocol::{InMessage, OutMessage};
use mio::{net::TcpStream, Interest, Poll, Token};
use rustls::{ServerConfig, ServerConnection};
use tungstenite::{
handshake::{server::NoCallback, MidHandshake},
protocol::WebSocketConfig,
HandshakeError, ServerHandshake,
};
use crate::common::*;
use crate::common::ConnectionMeta;
pub enum Stream {
TcpStream(TcpStream),
TlsStream(TlsStream<TcpStream>),
const MAX_PENDING_MESSAGES: usize = 16;
type TlsStream = rustls::StreamOwned<ServerConnection, TcpStream>;
type WsHandshakeResult<S> =
Result<tungstenite::WebSocket<S>, HandshakeError<ServerHandshake<S, NoCallback>>>;
type ConnectionReadResult<T> = ::std::io::Result<ConnectionReadStatus<T>>;
pub trait RegistryStatus {}
pub struct Registered;
impl RegistryStatus for Registered {}
pub struct NotRegistered;
impl RegistryStatus for NotRegistered {}
enum ConnectionReadStatus<T> {
Message(T, InMessage),
Ok(T),
WouldBlock(T),
}
impl Stream {
#[inline]
pub fn get_peer_addr(&self) -> ::std::io::Result<SocketAddr> {
match self {
Self::TcpStream(stream) => stream.peer_addr(),
Self::TlsStream(stream) => stream.get_ref().peer_addr(),
}
}
enum ConnectionState<R: RegistryStatus> {
TlsHandshaking(TlsHandshaking<R>),
WsHandshaking(WsHandshaking<R>),
WsConnection(WsConnection<R>),
}
#[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()),
}
pub struct Connection<R: RegistryStatus> {
pub valid_until: ValidUntil,
meta: ConnectionMeta,
state: ConnectionState<R>,
pub message_queue: VecDeque<OutMessage>,
pub interest: Interest,
phantom_data: PhantomData<R>,
}
impl<R: RegistryStatus> Connection<R> {
pub fn get_meta(&self) -> ConnectionMeta {
self.meta
}
}
impl Read for Stream {
#[inline]
fn read(&mut self, buf: &mut [u8]) -> Result<usize, ::std::io::Error> {
match self {
Self::TcpStream(stream) => stream.read(buf),
Self::TlsStream(stream) => stream.read(buf),
}
}
/// Not used but provided for completeness
#[inline]
fn read_vectored(
&mut self,
bufs: &mut [::std::io::IoSliceMut<'_>],
) -> ::std::io::Result<usize> {
match self {
Self::TcpStream(stream) => stream.read_vectored(bufs),
Self::TlsStream(stream) => stream.read_vectored(bufs),
}
}
}
impl Write for Stream {
#[inline]
fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
match self {
Self::TcpStream(stream) => stream.write(buf),
Self::TlsStream(stream) => stream.write(buf),
}
}
/// Not used but provided for completeness
#[inline]
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),
}
}
#[inline]
fn flush(&mut self) -> ::std::io::Result<()> {
match self {
Self::TcpStream(stream) => stream.flush(),
Self::TlsStream(stream) => stream.flush(),
}
}
}
enum HandshakeMachine {
TcpStream(TcpStream),
TlsStream(TlsStream<TcpStream>),
TlsMidHandshake(MidHandshakeTlsStream<TcpStream>),
WsMidHandshake(MidHandshake<ServerHandshake<Stream, NoCallback>>),
}
impl HandshakeMachine {
#[inline]
fn new(tcp_stream: TcpStream) -> Self {
Self::TcpStream(tcp_stream)
}
#[inline]
fn advance(
self,
impl Connection<NotRegistered> {
pub fn new(
tls_config: Arc<ServerConfig>,
ws_config: WebSocketConfig,
opt_tls_acceptor: &Option<TlsAcceptor>, // If set, run TLS
) -> (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))
} else {
let handshake_result = ::tungstenite::accept_with_config(
Stream::TcpStream(stream),
Some(ws_config),
);
tcp_stream: TcpStream,
valid_until: ValidUntil,
meta: ConnectionMeta,
) -> Self {
let state =
ConnectionState::TlsHandshaking(TlsHandshaking::new(tls_config, ws_config, tcp_stream));
Self::handle_ws_handshake_result(handshake_result)
Self {
valid_until,
meta,
state,
message_queue: Default::default(),
interest: Interest::READABLE,
phantom_data: PhantomData::default(),
}
}
/// Read until stream blocks (or error occurs)
///
/// Requires Connection not to be registered, since it might be dropped on errors
pub fn read<F>(
mut self,
message_handler: &mut F,
) -> ::std::io::Result<Connection<NotRegistered>>
where
F: FnMut(ConnectionMeta, InMessage),
{
loop {
let result = match self.state {
ConnectionState::TlsHandshaking(inner) => inner.read(),
ConnectionState::WsHandshaking(inner) => inner.read(),
ConnectionState::WsConnection(inner) => inner.read(),
};
match result {
Ok(ConnectionReadStatus::Message(state, message)) => {
self.state = state;
message_handler(self.meta, message);
// Stop looping even if WouldBlock wasn't necessarily reached. Otherwise,
// we might get stuck reading from this connection only. Since we register
// the connection again upon reinsertion into the ConnectionMap, we should
// be getting new events anyway.
return Ok(self);
}
}
HandshakeMachine::TlsStream(stream) => {
let handshake_result = ::tungstenite::accept(Stream::TlsStream(stream));
Ok(ConnectionReadStatus::Ok(state)) => {
self.state = state;
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())
}
}
}
::log::debug!("read connection");
}
Ok(ConnectionReadStatus::WouldBlock(state)) => {
self.state = state;
#[inline]
fn handle_tls_handshake_result(
result: Result<TlsStream<TcpStream>, ::native_tls::HandshakeError<TcpStream>>,
) -> (Option<Either<EstablishedWs, Self>>, bool) {
match result {
Ok(stream) => {
::log::trace!(
"established tls handshake with peer with addr: {:?}",
stream.get_ref().peer_addr()
);
::log::debug!("reading connection would block");
(Some(Either::Right(Self::TlsStream(stream))), false)
}
Err(native_tls::HandshakeError::WouldBlock(handshake)) => {
(Some(Either::Right(Self::TlsMidHandshake(handshake))), true)
}
Err(native_tls::HandshakeError::Failure(err)) => {
info!("tls handshake error: {}", err);
(None, false)
}
}
}
#[inline]
fn handle_ws_handshake_result(
result: Result<WebSocket<Stream>, HandshakeError<ServerHandshake<Stream, NoCallback>>>,
) -> (Option<Either<EstablishedWs, Self>>, bool) {
match result {
Ok(mut ws) => match ws.get_mut().get_peer_addr() {
Ok(peer_addr) => {
::log::trace!(
"established ws handshake with peer with addr: {:?}",
peer_addr
);
let established_ws = EstablishedWs { ws, peer_addr };
(Some(Either::Left(established_ws)), false)
return Ok(self);
}
Err(err) => {
::log::info!(
"get_peer_addr failed during handshake, removing connection: {:?}",
err
);
::log::debug!("Connection::read error: {}", err);
(None, false)
return Err(err);
}
}
}
}
pub fn register(self, poll: &mut Poll, token: Token) -> Connection<Registered> {
let state = match self.state {
ConnectionState::TlsHandshaking(inner) => {
ConnectionState::TlsHandshaking(inner.register(poll, token, self.interest))
}
ConnectionState::WsHandshaking(inner) => {
ConnectionState::WsHandshaking(inner.register(poll, token, self.interest))
}
ConnectionState::WsConnection(inner) => {
ConnectionState::WsConnection(inner.register(poll, token, self.interest))
}
};
Connection {
valid_until: self.valid_until,
meta: self.meta,
state,
message_queue: self.message_queue,
interest: self.interest,
phantom_data: PhantomData::default(),
}
}
pub fn close(self) {
::log::debug!("will close connection to {}", self.meta.naive_peer_addr);
match self.state {
ConnectionState::TlsHandshaking(inner) => inner.close(),
ConnectionState::WsHandshaking(inner) => inner.close(),
ConnectionState::WsConnection(inner) => inner.close(),
}
}
}
impl Connection<Registered> {
pub fn write_or_queue_message(
&mut self,
poll: &mut Poll,
message: OutMessage,
) -> ::std::io::Result<()> {
let message_clone = message.clone();
match self.write_message(message) {
Ok(()) => Ok(()),
Err(err) if err.kind() == ErrorKind::WouldBlock => {
if self.message_queue.len() < MAX_PENDING_MESSAGES {
self.message_queue.push_back(message_clone);
if !self.interest.is_writable() {
self.interest = Interest::WRITABLE;
self.reregister(poll)?;
}
} else {
::log::info!("Connection::message_queue is full, dropping message");
}
Ok(())
}
Err(err) => Err(err),
}
}
pub fn write(&mut self, poll: &mut Poll) -> ::std::io::Result<()> {
if let ConnectionState::WsConnection(_) = self.state {
while let Some(message) = self.message_queue.pop_front() {
let message_clone = message.clone();
match self.write_message(message) {
Ok(()) => {}
Err(err) if err.kind() == ErrorKind::WouldBlock => {
// Can't make message queue longer than it was before pop_front
self.message_queue.push_front(message_clone);
return Ok(());
}
Err(err) => {
return Err(err);
}
}
}
if self.message_queue.is_empty() {
self.interest = Interest::READABLE;
}
self.reregister(poll)?;
Ok(())
} else {
Err(std::io::Error::new(
ErrorKind::NotConnected,
"WebSocket connection not established",
))
}
}
fn write_message(&mut self, message: OutMessage) -> ::std::io::Result<()> {
if let ConnectionState::WsConnection(WsConnection {
ref mut web_socket, ..
}) = self.state
{
match web_socket.write_message(message.to_ws_message()) {
Ok(_) => {}
Err(tungstenite::Error::SendQueueFull(_message)) => {
return Err(std::io::Error::new(
ErrorKind::WouldBlock,
"Send queue full",
))
}
Err(tungstenite::Error::Io(err)) => return Err(err),
Err(err) => return Err(std::io::Error::new(ErrorKind::Other, err))?,
}
match web_socket.write_pending() {
Ok(()) => Ok(()),
Err(tungstenite::Error::Io(err)) => Err(err),
Err(err) => Err(std::io::Error::new(ErrorKind::Other, err))?,
}
} else {
Err(std::io::Error::new(
ErrorKind::NotConnected,
"WebSocket connection not established",
))
}
}
pub fn reregister(&mut self, poll: &mut Poll) -> ::std::io::Result<()> {
let token = Token(self.meta.connection_id.0);
match self.state {
ConnectionState::TlsHandshaking(ref mut inner) => {
inner.reregister(poll, token, self.interest)
}
ConnectionState::WsHandshaking(ref mut inner) => {
inner.reregister(poll, token, self.interest)
}
ConnectionState::WsConnection(ref mut inner) => {
inner.reregister(poll, token, self.interest)
}
}
}
pub fn deregister(self, poll: &mut Poll) -> Connection<NotRegistered> {
let state = match self.state {
ConnectionState::TlsHandshaking(inner) => {
ConnectionState::TlsHandshaking(inner.deregister(poll))
}
ConnectionState::WsHandshaking(inner) => {
ConnectionState::WsHandshaking(inner.deregister(poll))
}
ConnectionState::WsConnection(inner) => {
ConnectionState::WsConnection(inner.deregister(poll))
}
};
Connection {
valid_until: self.valid_until,
meta: self.meta,
state,
message_queue: self.message_queue,
interest: self.interest,
phantom_data: PhantomData::default(),
}
}
}
struct TlsHandshaking<R: RegistryStatus> {
tls_conn: ServerConnection,
ws_config: WebSocketConfig,
tcp_stream: TcpStream,
phantom_data: PhantomData<R>,
}
impl TlsHandshaking<NotRegistered> {
fn new(tls_config: Arc<ServerConfig>, ws_config: WebSocketConfig, stream: TcpStream) -> Self {
Self {
tls_conn: ServerConnection::new(tls_config).unwrap(),
ws_config,
tcp_stream: stream,
phantom_data: PhantomData::default(),
}
}
fn read(mut self) -> ConnectionReadResult<ConnectionState<NotRegistered>> {
match self.tls_conn.read_tls(&mut self.tcp_stream) {
Ok(0) => {
return Err(::std::io::Error::new(
ErrorKind::ConnectionReset,
"Connection closed",
))
}
Ok(_) => match self.tls_conn.process_new_packets() {
Ok(_) => {
while self.tls_conn.wants_write() {
self.tls_conn.write_tls(&mut self.tcp_stream)?;
}
if self.tls_conn.is_handshaking() {
Ok(ConnectionReadStatus::WouldBlock(
ConnectionState::TlsHandshaking(self),
))
} else {
let tls_stream = TlsStream::new(self.tls_conn, self.tcp_stream);
WsHandshaking::handle_handshake_result(tungstenite::accept_with_config(
tls_stream,
Some(self.ws_config),
))
}
}
Err(err) => {
let _ = self.tls_conn.write_tls(&mut self.tcp_stream);
Err(::std::io::Error::new(ErrorKind::InvalidData, err))
}
},
Err(HandshakeError::Interrupted(handshake)) => (
Some(Either::Right(HandshakeMachine::WsMidHandshake(handshake))),
true,
),
Err(HandshakeError::Failure(err)) => {
info!("ws handshake error: {}", err);
(None, false)
Err(err) if err.kind() == ErrorKind::WouldBlock => {
return Ok(ConnectionReadStatus::WouldBlock(
ConnectionState::TlsHandshaking(self),
))
}
}
}
}
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).
///
/// advance_handshakes takes ownership of self because the TLS and WebSocket
/// handshake methods do. get_established_ws doesn't, since work can be done
/// on a mutable reference to a tungstenite websocket, and this way, the whole
/// Connection doesn't have to be removed from and reinserted into the
/// TorrentMap. This is also the reason for wrapping Container.inner in an
/// Either instead of combining all states into one structure just having a
/// 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 {
Self {
ws_config,
valid_until,
inner: Either::Right(HandshakeMachine::new(tcp_stream)),
Err(err) => return Err(err),
}
}
#[inline]
pub fn get_established_ws(&mut self) -> Option<&mut EstablishedWs> {
match self.inner {
Either::Left(ref mut ews) => Some(ews),
Either::Right(_) => None,
fn register(
mut self,
poll: &mut Poll,
token: Token,
interest: Interest,
) -> TlsHandshaking<Registered> {
poll.registry()
.register(&mut self.tcp_stream, token, interest)
.unwrap();
TlsHandshaking {
tls_conn: self.tls_conn,
ws_config: self.ws_config,
tcp_stream: self.tcp_stream,
phantom_data: PhantomData::default(),
}
}
#[inline]
pub fn advance_handshakes(
self,
opt_tls_acceptor: &Option<TlsAcceptor>,
valid_until: ValidUntil,
) -> (Option<Self>, bool) {
match self.inner {
Either::Left(_) => (Some(self), false),
Either::Right(machine) => {
let ws_config = self.ws_config;
fn close(self) {
::log::debug!("closing connection (TlsHandshaking state)");
let (opt_inner, stop_loop) = machine.advance(ws_config, opt_tls_acceptor);
let _ = self.tcp_stream.shutdown(Shutdown::Both);
}
}
let opt_new_self = opt_inner.map(|inner| Self {
ws_config,
valid_until,
inner,
impl TlsHandshaking<Registered> {
fn deregister(mut self, poll: &mut Poll) -> TlsHandshaking<NotRegistered> {
poll.registry().deregister(&mut self.tcp_stream).unwrap();
TlsHandshaking {
tls_conn: self.tls_conn,
ws_config: self.ws_config,
tcp_stream: self.tcp_stream,
phantom_data: PhantomData::default(),
}
}
fn reregister(
&mut self,
poll: &mut Poll,
token: Token,
interest: Interest,
) -> std::io::Result<()> {
poll.registry()
.reregister(&mut self.tcp_stream, token, interest)
}
}
struct WsHandshaking<R: RegistryStatus> {
mid_handshake: MidHandshake<ServerHandshake<TlsStream, NoCallback>>,
phantom_data: PhantomData<R>,
}
impl WsHandshaking<NotRegistered> {
fn read(self) -> ConnectionReadResult<ConnectionState<NotRegistered>> {
Self::handle_handshake_result(self.mid_handshake.handshake())
}
fn handle_handshake_result(
handshake_result: WsHandshakeResult<TlsStream>,
) -> ConnectionReadResult<ConnectionState<NotRegistered>> {
match handshake_result {
Ok(web_socket) => {
let conn = ConnectionState::WsConnection(WsConnection {
web_socket,
phantom_data: PhantomData::default(),
});
(opt_new_self, stop_loop)
Ok(ConnectionReadStatus::Ok(conn))
}
Err(HandshakeError::Interrupted(mid_handshake)) => {
let conn = ConnectionState::WsHandshaking(WsHandshaking {
mid_handshake,
phantom_data: PhantomData::default(),
});
Ok(ConnectionReadStatus::WouldBlock(conn))
}
Err(HandshakeError::Failure(err)) => {
return Err(std::io::Error::new(ErrorKind::InvalidData, err))
}
}
}
#[inline]
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) {
::log::info!("error closing ws: {}", err);
}
fn register(
mut self,
poll: &mut Poll,
token: Token,
interest: Interest,
) -> WsHandshaking<Registered> {
let tcp_stream = &mut self.mid_handshake.get_mut().get_mut().sock;
// Required after ws.close()
if let Err(err) = ews.ws.write_pending() {
::log::info!("error writing pending messages after closing ws: {}", err)
}
}
poll.registry()
.register(tcp_stream, token, interest)
.unwrap();
WsHandshaking {
mid_handshake: self.mid_handshake,
phantom_data: PhantomData::default(),
}
}
pub fn deregister(&mut self, poll: &mut Poll) -> ::std::io::Result<()> {
use Either::{Left, Right};
fn close(mut self) {
::log::debug!("closing connection (WsHandshaking state)");
match self.inner {
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)
}
}
let tcp_stream = &mut self.mid_handshake.get_mut().get_mut().sock;
let _ = tcp_stream.shutdown(Shutdown::Both);
}
}
pub type ConnectionMap = HashMap<Token, Connection>;
impl WsHandshaking<Registered> {
fn deregister(mut self, poll: &mut Poll) -> WsHandshaking<NotRegistered> {
let tcp_stream = &mut self.mid_handshake.get_mut().get_mut().sock;
poll.registry().deregister(tcp_stream).unwrap();
WsHandshaking {
mid_handshake: self.mid_handshake,
phantom_data: PhantomData::default(),
}
}
fn reregister(
&mut self,
poll: &mut Poll,
token: Token,
interest: Interest,
) -> std::io::Result<()> {
let tcp_stream = &mut self.mid_handshake.get_mut().get_mut().sock;
poll.registry().reregister(tcp_stream, token, interest)
}
}
struct WsConnection<R: RegistryStatus> {
web_socket: tungstenite::WebSocket<TlsStream>,
phantom_data: PhantomData<R>,
}
impl WsConnection<NotRegistered> {
fn read(mut self) -> ConnectionReadResult<ConnectionState<NotRegistered>> {
match self.web_socket.read_message() {
Ok(
message @ tungstenite::Message::Text(_) | message @ tungstenite::Message::Binary(_),
) => match InMessage::from_ws_message(message) {
Ok(message) => {
::log::debug!("received WebSocket message");
Ok(ConnectionReadStatus::Message(
ConnectionState::WsConnection(self),
message,
))
}
Err(err) => Err(std::io::Error::new(ErrorKind::InvalidData, err)),
},
Ok(message) => {
::log::info!("received unexpected WebSocket message: {}", message);
Err(std::io::Error::new(
ErrorKind::InvalidData,
"unexpected WebSocket message type",
))
}
Err(tungstenite::Error::Io(err)) if err.kind() == ErrorKind::WouldBlock => {
let conn = ConnectionState::WsConnection(self);
Ok(ConnectionReadStatus::WouldBlock(conn))
}
Err(tungstenite::Error::Io(err)) => Err(err),
Err(err) => Err(std::io::Error::new(ErrorKind::InvalidData, err)),
}
}
fn register(
mut self,
poll: &mut Poll,
token: Token,
interest: Interest,
) -> WsConnection<Registered> {
poll.registry()
.register(self.web_socket.get_mut().get_mut(), token, interest)
.unwrap();
WsConnection {
web_socket: self.web_socket,
phantom_data: PhantomData::default(),
}
}
fn close(mut self) {
::log::debug!("closing connection (WsConnection state)");
let _ = self.web_socket.close(None);
let _ = self.web_socket.write_pending();
}
}
impl WsConnection<Registered> {
fn deregister(mut self, poll: &mut Poll) -> WsConnection<NotRegistered> {
poll.registry()
.deregister(self.web_socket.get_mut().get_mut())
.unwrap();
WsConnection {
web_socket: self.web_socket,
phantom_data: PhantomData::default(),
}
}
fn reregister(
&mut self,
poll: &mut Poll,
token: Token,
interest: Interest,
) -> std::io::Result<()> {
poll.registry()
.reregister(self.web_socket.get_mut().get_mut(), token, interest)
}
}

474
aquatic_ws/src/mio/socket/mod.rs
View file

@ -1,14 +1,13 @@
use std::io::ErrorKind;
use std::time::Duration;
use std::vec::Drain;
use std::sync::Arc;
use std::time::{Duration, Instant};
use anyhow::Context;
use aquatic_common::access_list::AccessListQuery;
use crossbeam_channel::Receiver;
use hashbrown::HashMap;
use log::{debug, error, info};
use mio::net::TcpListener;
use mio::{Events, Interest, Poll, Token};
use native_tls::TlsAcceptor;
use socket2::{Domain, Protocol, Socket, Type};
use tungstenite::protocol::WebSocketConfig;
use aquatic_common::convert_ipv4_mapped_ipv6;
@ -17,13 +16,101 @@ use aquatic_ws_protocol::*;
use crate::common::*;
use crate::config::Config;
pub mod connection;
use super::common::*;
pub mod connection;
pub mod utils;
use connection::{Connection, NotRegistered, Registered};
use connection::*;
use utils::*;
struct ConnectionMap {
token_counter: Token,
connections: HashMap<Token, Connection<Registered>>,
}
impl Default for ConnectionMap {
fn default() -> Self {
Self {
token_counter: Token(2),
connections: Default::default(),
}
}
}
impl ConnectionMap {
fn insert_and_register_new<F>(&mut self, poll: &mut Poll, connection_creator: F)
where
F: FnOnce(Token) -> Connection<NotRegistered>,
{
self.token_counter.0 = self.token_counter.0.wrapping_add(1);
// Don't assign LISTENER_TOKEN or CHANNEL_TOKEN
if self.token_counter.0 < 2 {
self.token_counter.0 = 2;
}
let token = self.token_counter;
// Remove, deregister and close any existing connection with this token.
// This shouldn't happen in practice.
if let Some(connection) = self.connections.remove(&token) {
::log::warn!(
"removing existing connection {} because of token reuse",
token.0
);
connection.deregister(poll).close();
}
let connection = connection_creator(token);
self.insert_and_register(poll, token, connection);
}
fn insert_and_register(
&mut self,
poll: &mut Poll,
key: Token,
conn: Connection<NotRegistered>,
) {
self.connections.insert(key, conn.register(poll, key));
}
fn remove_and_deregister(
&mut self,
poll: &mut Poll,
key: &Token,
) -> Option<Connection<NotRegistered>> {
if let Some(connection) = self.connections.remove(key) {
Some(connection.deregister(poll))
} else {
None
}
}
fn get_mut(&mut self, key: &Token) -> Option<&mut Connection<Registered>> {
self.connections.get_mut(key)
}
/// Close and remove inactive connections
fn clean(mut self, poll: &mut Poll) -> Self {
let now = Instant::now();
let mut retained_connections = HashMap::default();
for (token, connection) in self.connections.drain() {
if connection.valid_until.0 < now {
connection.deregister(poll).close();
} else {
retained_connections.insert(token, connection);
}
}
ConnectionMap {
connections: retained_connections,
..self
}
}
}
pub fn run_socket_worker(
config: Config,
@ -33,7 +120,7 @@ pub fn run_socket_worker(
poll: Poll,
in_message_sender: InMessageSender,
out_message_receiver: OutMessageReceiver,
opt_tls_acceptor: Option<TlsAcceptor>,
tls_config: Arc<rustls::ServerConfig>,
) {
match create_listener(&config) {
Ok(listener) => {
@ -47,7 +134,7 @@ pub fn run_socket_worker(
in_message_sender,
out_message_receiver,
listener,
opt_tls_acceptor,
tls_config,
);
}
Err(err) => {
@ -57,7 +144,7 @@ pub fn run_socket_worker(
}
}
pub fn run_poll_loop(
fn run_poll_loop(
config: Config,
state: &State,
socket_worker_index: usize,
@ -65,13 +152,13 @@ pub fn run_poll_loop(
in_message_sender: InMessageSender,
out_message_receiver: OutMessageReceiver,
listener: ::std::net::TcpListener,
opt_tls_acceptor: Option<TlsAcceptor>,
tls_config: Arc<rustls::ServerConfig>,
) {
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),
max_send_queue: None,
max_send_queue: Some(2),
..Default::default()
};
@ -82,10 +169,9 @@ pub fn run_poll_loop(
.register(&mut listener, LISTENER_TOKEN, Interest::READABLE)
.unwrap();
let mut connections: ConnectionMap = HashMap::new();
let mut connections = ConnectionMap::default();
let mut local_responses = Vec::new();
let mut poll_token_counter = Token(0usize);
let mut iter_counter = 0usize;
loop {
@ -97,41 +183,68 @@ pub fn run_poll_loop(
for event in events.iter() {
let token = event.token();
if token == LISTENER_TOKEN {
accept_new_streams(
ws_config,
&mut listener,
&mut poll,
&mut connections,
valid_until,
&mut poll_token_counter,
);
} else if token != CHANNEL_TOKEN {
run_handshakes_and_read_messages(
&config,
state,
socket_worker_index,
&mut local_responses,
&in_message_sender,
&opt_tls_acceptor,
&mut poll,
&mut connections,
token,
valid_until,
);
match token {
LISTENER_TOKEN => {
accept_new_streams(
&tls_config,
ws_config,
socket_worker_index,
&mut listener,
&mut poll,
&mut connections,
valid_until,
);
}
CHANNEL_TOKEN => {
write_or_queue_messages(
&mut poll,
out_message_receiver
.try_iter()
.take(out_message_receiver.len()),
&mut connections,
);
}
token => {
if event.is_writable() {
let mut remove_connection = false;
if let Some(connection) = connections.get_mut(&token) {
if let Err(err) = connection.write(&mut poll) {
::log::debug!("Connection::write error: {}", err);
remove_connection = true;
}
}
if remove_connection {
if let Some(connection) =
connections.remove_and_deregister(&mut poll, &token)
{
connection.close();
}
}
}
if event.is_readable() {
handle_stream_read_event(
&config,
state,
&mut local_responses,
&in_message_sender,
&mut poll,
&mut connections,
token,
valid_until,
);
}
}
}
send_out_messages(
&mut poll,
local_responses.drain(..),
&out_message_receiver,
&mut connections,
);
write_or_queue_messages(&mut poll, local_responses.drain(..), &mut connections);
}
// Remove inactive connections, but not every iteration
if iter_counter % 128 == 0 {
remove_inactive_connections(&mut connections);
connections = connections.clean(&mut poll);
}
iter_counter = iter_counter.wrapping_add(1);
@ -139,194 +252,155 @@ pub fn run_poll_loop(
}
fn accept_new_streams(
tls_config: &Arc<rustls::ServerConfig>,
ws_config: WebSocketConfig,
socket_worker_index: usize,
listener: &mut TcpListener,
poll: &mut Poll,
connections: &mut ConnectionMap,
valid_until: ValidUntil,
poll_token_counter: &mut Token,
) {
loop {
match listener.accept() {
Ok((mut stream, _)) => {
poll_token_counter.0 = poll_token_counter.0.wrapping_add(1);
Ok((stream, _)) => {
let naive_peer_addr = if let Ok(peer_addr) = stream.peer_addr() {
peer_addr
} else {
continue;
};
if poll_token_counter.0 < 2 {
poll_token_counter.0 = 2;
}
let token = *poll_token_counter;
remove_connection_if_exists(poll, connections, token);
poll.registry()
.register(&mut stream, token, Interest::READABLE)
.unwrap();
let connection = Connection::new(ws_config, valid_until, stream);
connections.insert(token, connection);
}
Err(err) => {
if err.kind() == ErrorKind::WouldBlock {
break;
}
info!("error while accepting streams: {}", err);
}
}
}
}
/// 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(
config: &Config,
state: &State,
socket_worker_index: usize,
local_responses: &mut Vec<(ConnectionMeta, OutMessage)>,
in_message_sender: &InMessageSender,
opt_tls_acceptor: &Option<TlsAcceptor>, // If set, run TLS
poll: &mut Poll,
connections: &mut ConnectionMap,
poll_token: Token,
valid_until: ValidUntil,
) {
let access_list_mode = config.access_list.mode;
loop {
if let Some(established_ws) = connections
.get_mut(&poll_token)
.map(|c| {
// Ugly but works
c.valid_until = valid_until;
c
})
.and_then(Connection::get_established_ws)
{
use ::tungstenite::Error::Io;
match established_ws.ws.read_message() {
Ok(ws_message) => {
let naive_peer_addr = established_ws.peer_addr;
connections.insert_and_register_new(poll, move |token| {
let converted_peer_ip = convert_ipv4_mapped_ipv6(naive_peer_addr.ip());
let meta = ConnectionMeta {
out_message_consumer_id: ConsumerId(socket_worker_index),
connection_id: ConnectionId(poll_token.0),
connection_id: ConnectionId(token.0),
naive_peer_addr,
converted_peer_ip,
pending_scrape_id: None, // FIXME
};
debug!("read message");
match InMessage::from_ws_message(ws_message) {
Ok(InMessage::AnnounceRequest(ref request))
if !state
.access_list
.allows(access_list_mode, &request.info_hash.0) =>
{
let out_message = OutMessage::ErrorResponse(ErrorResponse {
failure_reason: "Info hash not allowed".into(),
action: Some(ErrorResponseAction::Announce),
info_hash: Some(request.info_hash),
});
local_responses.push((meta, out_message));
}
Ok(in_message) => {
if let Err(err) = in_message_sender.send((meta, in_message)) {
error!("InMessageSender: couldn't send message: {:?}", err);
}
}
Err(_) => {
// FIXME: maybe this condition just occurs when enough data hasn't been recevied?
/*
info!("error parsing message: {:?}", err);
let out_message = OutMessage::ErrorResponse(ErrorResponse {
failure_reason: "Error parsing message".into(),
action: None,
info_hash: None,
});
local_responses.push((meta, out_message));
*/
}
}
}
Err(Io(err)) if err.kind() == ErrorKind::WouldBlock => {
break;
}
Err(tungstenite::Error::ConnectionClosed) => {
remove_connection_if_exists(poll, connections, poll_token);
break;
}
Err(err) => {
info!("error reading messages: {}", err);
remove_connection_if_exists(poll, connections, poll_token);
break;
}
Connection::new(tls_config.clone(), ws_config, stream, valid_until, meta)
});
}
} 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);
}
if stop_loop {
Err(err) if err.kind() == ErrorKind::WouldBlock => {
break;
}
} else {
break;
Err(err) => {
::log::info!("error while accepting streams: {}", err);
}
}
}
}
/// Read messages from channel, send to peers
pub fn send_out_messages(
fn handle_stream_read_event(
config: &Config,
state: &State,
local_responses: &mut Vec<(ConnectionMeta, OutMessage)>,
in_message_sender: &InMessageSender,
poll: &mut Poll,
local_responses: Drain<(ConnectionMeta, OutMessage)>,
out_message_receiver: &Receiver<(ConnectionMeta, OutMessage)>,
connections: &mut ConnectionMap,
token: Token,
valid_until: ValidUntil,
) {
let len = out_message_receiver.len();
let access_list_mode = config.access_list.mode;
for (meta, out_message) in local_responses.chain(out_message_receiver.try_iter().take(len)) {
let opt_established_ws = connections
.get_mut(&Token(meta.connection_id.0))
.and_then(Connection::get_established_ws);
if let Some(mut connection) = connections.remove_and_deregister(poll, &token) {
let message_handler = &mut |meta, message| match message {
InMessage::AnnounceRequest(ref request)
if !state
.access_list
.allows(access_list_mode, &request.info_hash.0) =>
{
let out_message = OutMessage::ErrorResponse(ErrorResponse {
failure_reason: "Info hash not allowed".into(),
action: Some(ErrorResponseAction::Announce),
info_hash: Some(request.info_hash),
});
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;
local_responses.push((meta, out_message));
}
use ::tungstenite::Error::Io;
let ws_message = out_message.to_ws_message();
match established_ws.ws.write_message(ws_message) {
Ok(()) => {
debug!("sent message");
}
Err(Io(err)) if err.kind() == ErrorKind::WouldBlock => {}
Err(tungstenite::Error::ConnectionClosed) => {
remove_connection_if_exists(poll, connections, Token(meta.connection_id.0));
}
Err(err) => {
info!("error writing ws message: {}", err);
remove_connection_if_exists(poll, connections, Token(meta.connection_id.0));
in_message => {
if let Err(err) = in_message_sender.send((meta, in_message)) {
::log::info!("InMessageSender: couldn't send message: {:?}", err);
}
}
};
connection.valid_until = valid_until;
match connection.read(message_handler) {
Ok(connection) => {
connections.insert_and_register(poll, token, connection);
}
Err(_) => {}
}
}
}
fn write_or_queue_messages<I>(poll: &mut Poll, responses: I, connections: &mut ConnectionMap)
where
I: Iterator<Item = (ConnectionMeta, OutMessage)>,
{
for (meta, out_message) in responses {
let token = Token(meta.connection_id.0);
let mut remove_connection = false;
if let Some(connection) = connections.get_mut(&token) {
if connection.get_meta().naive_peer_addr != meta.naive_peer_addr {
::log::warn!(
"socket worker error: connection socket addr {} didn't match channel {}. Token: {}.",
connection.get_meta().naive_peer_addr,
meta.naive_peer_addr,
token.0
);
remove_connection = true;
} else {
match connection.write_or_queue_message(poll, out_message) {
Ok(()) => {}
Err(err) => {
::log::debug!("Connection::write_or_queue_message error: {}", err);
remove_connection = true;
}
}
}
}
if remove_connection {
connections.remove_and_deregister(poll, &token);
}
}
}
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")?;
if config.network.ipv6_only {
builder
.set_only_v6(true)
.context("Couldn't put socket in ipv6 only mode")?
}
builder
.set_nonblocking(true)
.context("Couldn't put socket in non-blocking mode")?;
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)
.context("Couldn't listen for connections on socket")?;
Ok(builder.into())
}

66
aquatic_ws/src/mio/socket/utils.rs
View file

@ -1,66 +0,0 @@
use std::time::Instant;
use anyhow::Context;
use mio::{Poll, Token};
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))
} else {
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)
.context("Couldn't put socket in ipv6 only mode")?
}
builder
.set_nonblocking(true)
.context("Couldn't put socket in non-blocking mode")?;
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)
.context("Couldn't listen for connections on socket")?;
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) {
connection.close();
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) {
let now = Instant::now();
connections.retain(|_, connection| {
if connection.valid_until.0 < now {
connection.close();
false
} else {
true
}
});
connections.shrink_to_fit();
}