ws: use file naming more in line with other impls; other small fixes

aquatic_ws/src/mio/socket/connection.rs

@@ -0,0 +1,298 @@
+use std::io::{Read, Write};
+use std::net::SocketAddr;
+
+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 crate::common::*;
+
+pub enum Stream {
+    TcpStream(TcpStream),
+    TlsStream(TlsStream<TcpStream>),
+}
+
+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(),
+        }
+    }
+
+    #[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()),
+        }
+    }
+}
+
+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,
+        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),
+                    );
+
+                    Self::handle_ws_handshake_result(handshake_result)
+                }
+            }
+            HandshakeMachine::TlsStream(stream) => {
+                let handshake_result = ::tungstenite::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())
+            }
+        }
+    }
+
+    #[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()
+                );
+
+                (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)
+                }
+                Err(err) => {
+                    ::log::info!(
+                        "get_peer_addr failed during handshake, removing connection: {:?}",
+                        err
+                    );
+
+                    (None, false)
+                }
+            },
+            Err(HandshakeError::Interrupted(handshake)) => (
+                Some(Either::Right(HandshakeMachine::WsMidHandshake(handshake))),
+                true,
+            ),
+            Err(HandshakeError::Failure(err)) => {
+                info!("ws handshake error: {}", err);
+
+                (None, false)
+            }
+        }
+    }
+}
+
+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)),
+        }
+    }
+
+    #[inline]
+    pub fn get_established_ws(&mut self) -> Option<&mut EstablishedWs> {
+        match self.inner {
+            Either::Left(ref mut ews) => Some(ews),
+            Either::Right(_) => None,
+        }
+    }
+
+    #[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;
+
+                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,
+                });
+
+                (opt_new_self, stop_loop)
+            }
+        }
+    }
+
+    #[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);
+                }
+
+                // Required after ws.close()
+                if let Err(err) = ews.ws.write_pending() {
+                    ::log::info!("error writing pending messages after closing ws: {}", err)
+                }
+            }
+        }
+    }
+
+    pub fn deregister(&mut self, poll: &mut Poll) -> ::std::io::Result<()> {
+        use Either::{Left, Right};
+
+        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)
+            }
+        }
+    }
+}
+
+pub type ConnectionMap = HashMap<Token, Connection>;

aquatic_ws/src/mio/socket/mod.rs

@@ -0,0 +1,332 @@
+use std::io::ErrorKind;
+use std::time::Duration;
+use std::vec::Drain;
+
+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 tungstenite::protocol::WebSocketConfig;
+
+use aquatic_common::convert_ipv4_mapped_ipv6;
+use aquatic_ws_protocol::*;
+
+use crate::common::*;
+use crate::config::Config;
+
+use super::common::*;
+
+pub mod connection;
+pub mod utils;
+
+use connection::*;
+use utils::*;
+
+pub fn run_socket_worker(
+    config: Config,
+    state: State,
+    socket_worker_index: usize,
+    socket_worker_statuses: SocketWorkerStatuses,
+    poll: Poll,
+    in_message_sender: InMessageSender,
+    out_message_receiver: OutMessageReceiver,
+    opt_tls_acceptor: Option<TlsAcceptor>,
+) {
+    match create_listener(&config) {
+        Ok(listener) => {
+            socket_worker_statuses.lock()[socket_worker_index] = Some(Ok(()));
+
+            run_poll_loop(
+                config,
+                &state,
+                socket_worker_index,
+                poll,
+                in_message_sender,
+                out_message_receiver,
+                listener,
+                opt_tls_acceptor,
+            );
+        }
+        Err(err) => {
+            socket_worker_statuses.lock()[socket_worker_index] =
+                Some(Err(format!("Couldn't open socket: {:#}", err)));
+        }
+    }
+}
+
+pub fn run_poll_loop(
+    config: Config,
+    state: &State,
+    socket_worker_index: usize,
+    mut poll: Poll,
+    in_message_sender: InMessageSender,
+    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 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,
+        ..Default::default()
+    };
+
+    let mut listener = TcpListener::from_std(listener);
+    let mut events = Events::with_capacity(config.network.poll_event_capacity);
+
+    poll.registry()
+        .register(&mut listener, LISTENER_TOKEN, Interest::READABLE)
+        .unwrap();
+
+    let mut connections: ConnectionMap = HashMap::new();
+    let mut local_responses = Vec::new();
+
+    let mut poll_token_counter = Token(0usize);
+    let mut iter_counter = 0usize;
+
+    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() {
+            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,
+                );
+            }
+
+            send_out_messages(
+                &mut poll,
+                local_responses.drain(..),
+                &out_message_receiver,
+                &mut connections,
+            );
+        }
+
+        // Remove inactive connections, but not every iteration
+        if iter_counter % 128 == 0 {
+            remove_inactive_connections(&mut connections);
+        }
+
+        iter_counter = iter_counter.wrapping_add(1);
+    }
+}
+
+fn accept_new_streams(
+    ws_config: WebSocketConfig,
+    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);
+
+                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;
+                    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),
+                        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;
+                }
+            }
+        } 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 {
+                break;
+            }
+        } else {
+            break;
+        }
+    }
+}
+
+/// Read messages from channel, send to peers
+pub fn send_out_messages(
+    poll: &mut Poll,
+    local_responses: Drain<(ConnectionMeta, OutMessage)>,
+    out_message_receiver: &Receiver<(ConnectionMeta, OutMessage)>,
+    connections: &mut ConnectionMap,
+) {
+    let len = out_message_receiver.len();
+
+    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(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) {
+                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));
+                }
+            }
+        }
+    }
+}

aquatic_ws/src/mio/socket/utils.rs

@@ -0,0 +1,66 @@
+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();
+}