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

Run rustfmt, clean up aquatic_http_protocol/Cargo.toml

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

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

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

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

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