mirror of
https://github.com/YGGverse/aquatic.git
synced 2026-04-02 02:35:31 +00:00
WIP: start work on http tracker
This commit is contained in:
Joakim Frostegård
parent
ebe4d4357b
commit
404e528616
15 changed files with 1845 additions and 0 deletions
278
aquatic_http/src/lib/network/connection.rs
Normal file
278
aquatic_http/src/lib/network/connection.rs
Normal file
|
|
@ -0,0 +1,278 @@
|
|||
use std::net::{SocketAddr};
|
||||
use std::io::{Read, Write};
|
||||
|
||||
use either::Either;
|
||||
use hashbrown::HashMap;
|
||||
use log::info;
|
||||
use mio::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 tungstenite::protocol::WebSocketConfig;
|
||||
|
||||
use crate::common::*;
|
||||
|
||||
|
||||
pub enum Stream {
|
||||
TcpStream(TcpStream),
|
||||
TlsStream(TlsStream<TcpStream>),
|
||||
}
|
||||
|
||||
|
||||
impl Stream {
|
||||
#[inline]
|
||||
pub fn get_peer_addr(&self) -> SocketAddr {
|
||||
match self {
|
||||
Self::TcpStream(stream) => stream.peer_addr().unwrap(),
|
||||
Self::TlsStream(stream) => stream.get_ref().peer_addr().unwrap(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
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::server::accept_with_config(
|
||||
Stream::TcpStream(stream),
|
||||
Some(ws_config)
|
||||
);
|
||||
|
||||
Self::handle_ws_handshake_result(handshake_result)
|
||||
}
|
||||
},
|
||||
HandshakeMachine::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())
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn handle_tls_handshake_result(
|
||||
result: Result<TlsStream<TcpStream>, ::native_tls::HandshakeError<TcpStream>>,
|
||||
) -> (Option<Either<EstablishedWs, Self>>, bool) {
|
||||
match result {
|
||||
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);
|
||||
|
||||
(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) => {
|
||||
let peer_addr = ws.get_mut().get_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::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<'a>(&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(){
|
||||
ews.ws.close(None).unwrap();
|
||||
|
||||
// Required after ws.close()
|
||||
if let Err(err) = ews.ws.write_pending(){
|
||||
info!(
|
||||
"error writing pending messages after closing ws: {}",
|
||||
err
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
pub type ConnectionMap = HashMap<Token, Connection>;
|
||||
279
aquatic_http/src/lib/network/mod.rs
Normal file
279
aquatic_http/src/lib/network/mod.rs
Normal file
|
|
@ -0,0 +1,279 @@
|
|||
use std::time::Duration;
|
||||
use std::io::ErrorKind;
|
||||
|
||||
use hashbrown::HashMap;
|
||||
use log::{info, debug, error};
|
||||
use native_tls::TlsAcceptor;
|
||||
use mio::{Events, Poll, Interest, Token};
|
||||
use mio::net::TcpListener;
|
||||
use tungstenite::protocol::WebSocketConfig;
|
||||
|
||||
use crate::common::*;
|
||||
use crate::config::Config;
|
||||
use crate::protocol::*;
|
||||
|
||||
pub mod connection;
|
||||
pub mod utils;
|
||||
|
||||
use connection::*;
|
||||
use utils::*;
|
||||
|
||||
|
||||
// will be almost identical to ws version
|
||||
pub fn run_socket_worker(
|
||||
config: Config,
|
||||
socket_worker_index: usize,
|
||||
socket_worker_statuses: SocketWorkerStatuses,
|
||||
request_channel_sender: InMessageSender,
|
||||
response_channel_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,
|
||||
socket_worker_index,
|
||||
request_channel_sender,
|
||||
response_channel_receiver,
|
||||
listener,
|
||||
opt_tls_acceptor
|
||||
);
|
||||
},
|
||||
Err(err) => {
|
||||
socket_worker_statuses.lock()[socket_worker_index] = Some(
|
||||
Err(format!("Couldn't open socket: {:#}", err))
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// will be almost identical to ws version
|
||||
pub fn run_poll_loop(
|
||||
config: Config,
|
||||
socket_worker_index: usize,
|
||||
request_channel_sender: InMessageSender,
|
||||
response_channel_receiver: OutMessageReceiver,
|
||||
listener: ::std::net::TcpListener,
|
||||
opt_tls_acceptor: Option<TlsAcceptor>,
|
||||
){
|
||||
let poll_timeout = Duration::from_millis(
|
||||
config.network.poll_timeout_milliseconds
|
||||
);
|
||||
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,
|
||||
};
|
||||
|
||||
let mut listener = TcpListener::from_std(listener);
|
||||
let mut poll = Poll::new().expect("create poll");
|
||||
let mut events = Events::with_capacity(config.network.poll_event_capacity);
|
||||
|
||||
poll.registry()
|
||||
.register(&mut listener, Token(0), Interest::READABLE)
|
||||
.unwrap();
|
||||
|
||||
let mut connections: ConnectionMap = HashMap::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.0 == 0 {
|
||||
accept_new_streams(
|
||||
ws_config,
|
||||
&mut listener,
|
||||
&mut poll,
|
||||
&mut connections,
|
||||
valid_until,
|
||||
&mut poll_token_counter,
|
||||
);
|
||||
} else {
|
||||
run_handshake_and_read_requests(
|
||||
socket_worker_index,
|
||||
&request_channel_sender,
|
||||
&opt_tls_acceptor,
|
||||
&mut connections,
|
||||
token,
|
||||
valid_until,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
send_out_messages(
|
||||
response_channel_receiver.drain(),
|
||||
&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);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// will be identical to ws version
|
||||
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 == 0 {
|
||||
poll_token_counter.0 = 1;
|
||||
}
|
||||
|
||||
let token = *poll_token_counter;
|
||||
|
||||
remove_connection_if_exists(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_handshake_and_read_requests(
|
||||
socket_worker_index: usize,
|
||||
request_channel_sender: &InMessageSender,
|
||||
opt_tls_acceptor: &Option<TlsAcceptor>, // If set, run TLS
|
||||
connections: &mut ConnectionMap,
|
||||
poll_token: Token,
|
||||
valid_until: ValidUntil,
|
||||
){
|
||||
loop {
|
||||
if let Some(established_ws) = connections.get_mut(&poll_token)
|
||||
.and_then(Connection::get_established_ws)
|
||||
{
|
||||
use ::tungstenite::Error::Io;
|
||||
|
||||
match established_ws.ws.read_message(){
|
||||
Ok(ws_message) => {
|
||||
if let Some(in_message) = InMessage::from_ws_message(ws_message){
|
||||
let meta = ConnectionMeta {
|
||||
worker_index: socket_worker_index,
|
||||
poll_token,
|
||||
peer_addr: established_ws.peer_addr
|
||||
};
|
||||
|
||||
debug!("read message");
|
||||
|
||||
if let Err(err) = request_channel_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(connections, poll_token);
|
||||
|
||||
break
|
||||
},
|
||||
Err(err) => {
|
||||
info!("error reading messages: {}", err);
|
||||
|
||||
remove_connection_if_exists(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;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/// Read messages from channel, send to peers
|
||||
pub fn send_out_messages(
|
||||
response_channel_receiver: ::flume::Drain<(ConnectionMeta, OutMessage)>,
|
||||
connections: &mut ConnectionMap,
|
||||
){
|
||||
for (meta, out_message) in response_channel_receiver {
|
||||
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.peer_addr {
|
||||
info!("socket worker error: peer socket addrs didn't match");
|
||||
|
||||
continue;
|
||||
}
|
||||
|
||||
use ::tungstenite::Error::Io;
|
||||
|
||||
match established_ws.ws.write_message(out_message.to_ws_message()){
|
||||
Ok(()) => {
|
||||
debug!("sent message");
|
||||
},
|
||||
Err(Io(err)) if err.kind() == ErrorKind::WouldBlock => {},
|
||||
Err(tungstenite::Error::ConnectionClosed) => {
|
||||
remove_connection_if_exists(connections, meta.poll_token);
|
||||
},
|
||||
Err(err) => {
|
||||
info!("error writing ws message: {}", err);
|
||||
|
||||
remove_connection_if_exists(
|
||||
connections,
|
||||
meta.poll_token
|
||||
);
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
74
aquatic_http/src/lib/network/utils.rs
Normal file
74
aquatic_http/src/lib/network/utils.rs
Normal file
|
|
@ -0,0 +1,74 @@
|
|||
use std::time::Instant;
|
||||
|
||||
use anyhow::Context;
|
||||
use mio::Token;
|
||||
use socket2::{Socket, Domain, Type, Protocol};
|
||||
|
||||
use crate::config::Config;
|
||||
|
||||
use super::connection::*;
|
||||
|
||||
|
||||
// will be identical to ws version
|
||||
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_tcp_listener())
|
||||
}
|
||||
|
||||
|
||||
// will be identical to ws version
|
||||
/// Don't bother with deregistering from Poll. In my understanding, this is
|
||||
/// done automatically when the stream is dropped, as long as there are no
|
||||
/// other references to the file descriptor, such as when it is accessed
|
||||
/// in multiple threads.
|
||||
pub fn remove_connection_if_exists(
|
||||
connections: &mut ConnectionMap,
|
||||
token: Token,
|
||||
){
|
||||
if let Some(mut connection) = connections.remove(&token){
|
||||
connection.close();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// will be identical to ws version
|
||||
// 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();
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue