aquatic/crates/ws/src/workers/socket/connection.rs
2024-01-20 09:56:14 +01:00

648 lines
24 KiB
Rust

use std::borrow::Cow;
use std::cell::RefCell;
use std::collections::BTreeMap;
use std::rc::Rc;
use std::sync::Arc;
use std::time::Duration;
use anyhow::Context;
use aquatic_common::access_list::{create_access_list_cache, AccessListArcSwap, AccessListCache};
use aquatic_common::rustls_config::RustlsConfig;
use aquatic_common::ServerStartInstant;
use aquatic_ws_protocol::common::{InfoHash, PeerId, ScrapeAction};
use aquatic_ws_protocol::incoming::{
AnnounceEvent, AnnounceRequest, InMessage, ScrapeRequest, ScrapeRequestInfoHashes,
};
use aquatic_ws_protocol::outgoing::{
ErrorResponse, ErrorResponseAction, OutMessage, ScrapeResponse, ScrapeStatistics,
};
use arc_swap::ArcSwap;
use async_tungstenite::WebSocketStream;
use futures::stream::{SplitSink, SplitStream};
use futures::{AsyncWriteExt, StreamExt};
use futures_lite::future::race;
use futures_rustls::TlsAcceptor;
use glommio::channels::channel_mesh::Senders;
use glommio::channels::local_channel::{LocalReceiver, LocalSender};
use glommio::net::TcpStream;
use glommio::timer::timeout;
use glommio::{enclose, prelude::*};
use hashbrown::hash_map::Entry;
use hashbrown::HashMap;
use slab::Slab;
#[cfg(feature = "metrics")]
use metrics::{Counter, Gauge};
use crate::common::*;
use crate::config::Config;
use crate::workers::socket::calculate_in_message_consumer_index;
#[cfg(feature = "metrics")]
use crate::workers::socket::{ip_version_to_metrics_str, WORKER_INDEX};
pub struct ConnectionRunner {
pub config: Rc<Config>,
pub access_list: Arc<AccessListArcSwap>,
pub in_message_senders: Rc<Senders<(InMessageMeta, InMessage)>>,
pub connection_valid_until: Rc<RefCell<ValidUntil>>,
pub out_message_sender: Rc<LocalSender<(OutMessageMeta, OutMessage)>>,
pub out_message_receiver: LocalReceiver<(OutMessageMeta, OutMessage)>,
pub server_start_instant: ServerStartInstant,
pub out_message_consumer_id: ConsumerId,
pub connection_id: ConnectionId,
pub opt_tls_config: Option<Arc<ArcSwap<RustlsConfig>>>,
pub ip_version: IpVersion,
}
impl ConnectionRunner {
pub async fn run(
self,
control_message_senders: Rc<Senders<SwarmControlMessage>>,
close_conn_receiver: LocalReceiver<()>,
stream: TcpStream,
) {
let clean_up_data = ConnectionCleanupData {
announced_info_hashes: Default::default(),
ip_version: self.ip_version,
opt_peer_client: Default::default(),
#[cfg(feature = "metrics")]
active_connections_gauge: ::metrics::gauge!(
"aquatic_active_connections",
"ip_version" => ip_version_to_metrics_str(self.ip_version),
"worker_index" => WORKER_INDEX.get().to_string(),
),
};
clean_up_data.before_open();
let config = self.config.clone();
let connection_id = self.connection_id;
race(
async {
if let Err(err) = self.run_inner(clean_up_data.clone(), stream).await {
::log::debug!("connection {:?} closed: {:#}", connection_id, err);
}
},
async {
close_conn_receiver.recv().await;
},
)
.await;
::log::debug!("connection {:?} starting clean up", connection_id);
clean_up_data
.after_close(&config, control_message_senders)
.await;
::log::debug!("connection {:?} finished clean up", connection_id);
}
async fn run_inner(
self,
clean_up_data: ConnectionCleanupData,
mut stream: TcpStream,
) -> anyhow::Result<()> {
if let Some(tls_config) = self.opt_tls_config.as_ref() {
let tls_config = tls_config.load_full();
let tls_acceptor = TlsAcceptor::from(tls_config);
let stream = tls_acceptor.accept(stream).await?;
self.run_inner_stream_agnostic(clean_up_data, stream).await
} else {
// Implementing this over TLS is too cumbersome, since the crate used
// for TLS streams doesn't support peek and tungstenite doesn't
// properly support sending a HTTP error response in accept_hdr
// callback.
if self.config.network.enable_http_health_checks {
let mut peek_buf = [0u8; 11];
stream
.peek(&mut peek_buf)
.await
.map_err(|err| anyhow::anyhow!("error peeking: {:#}", err))?;
if &peek_buf == b"GET /health" {
stream
.write_all(b"HTTP/1.1 200 Ok\r\nContent-Length: 2\r\n\r\nOk")
.await
.map_err(|err| {
anyhow::anyhow!("error sending health check response: {:#}", err)
})?;
stream.flush().await.map_err(|err| {
anyhow::anyhow!("error flushing health check response: {:#}", err)
})?;
return Err(anyhow::anyhow!(
"client requested health check, skipping websocket negotiation"
));
}
}
self.run_inner_stream_agnostic(clean_up_data, stream).await
}
}
async fn run_inner_stream_agnostic<S>(
self,
clean_up_data: ConnectionCleanupData,
stream: S,
) -> anyhow::Result<()>
where
S: futures::AsyncRead + futures::AsyncWrite + Unpin + 'static,
{
let ws_config = tungstenite::protocol::WebSocketConfig {
max_frame_size: Some(self.config.network.websocket_max_frame_size),
max_message_size: Some(self.config.network.websocket_max_message_size),
write_buffer_size: self.config.network.websocket_write_buffer_size,
max_write_buffer_size: self.config.network.websocket_write_buffer_size * 3,
..Default::default()
};
let stream = async_tungstenite::accept_async_with_config(stream, Some(ws_config)).await?;
let (ws_out, ws_in) = futures::StreamExt::split(stream);
let pending_scrape_slab = Rc::new(RefCell::new(Slab::new()));
let access_list_cache = create_access_list_cache(&self.access_list);
let config = self.config.clone();
let reader_future = enclose!((pending_scrape_slab, clean_up_data) async move {
let mut reader = ConnectionReader {
config: self.config.clone(),
access_list_cache,
in_message_senders: self.in_message_senders,
out_message_sender: self.out_message_sender,
pending_scrape_slab,
out_message_consumer_id: self.out_message_consumer_id,
ws_in,
ip_version: self.ip_version,
connection_id: self.connection_id,
clean_up_data: clean_up_data.clone(),
#[cfg(feature = "metrics")]
total_announce_requests_counter: ::metrics::counter!(
"aquatic_requests_total",
"type" => "announce",
"ip_version" => ip_version_to_metrics_str(self.ip_version),
"worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(),
),
#[cfg(feature = "metrics")]
total_scrape_requests_counter: ::metrics::counter!(
"aquatic_requests_total",
"type" => "scrape",
"ip_version" => ip_version_to_metrics_str(self.ip_version),
"worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(),
)
};
reader.run_in_message_loop().await
});
let writer_future = async move {
let mut writer = ConnectionWriter {
config,
out_message_receiver: self.out_message_receiver,
connection_valid_until: self.connection_valid_until,
ws_out,
pending_scrape_slab,
server_start_instant: self.server_start_instant,
ip_version: self.ip_version,
clean_up_data,
};
writer.run_out_message_loop().await
};
race(reader_future, writer_future).await
}
}
struct ConnectionReader<S> {
config: Rc<Config>,
access_list_cache: AccessListCache,
in_message_senders: Rc<Senders<(InMessageMeta, InMessage)>>,
out_message_sender: Rc<LocalSender<(OutMessageMeta, OutMessage)>>,
pending_scrape_slab: Rc<RefCell<Slab<PendingScrapeResponse>>>,
out_message_consumer_id: ConsumerId,
ws_in: SplitStream<WebSocketStream<S>>,
ip_version: IpVersion,
connection_id: ConnectionId,
clean_up_data: ConnectionCleanupData,
#[cfg(feature = "metrics")]
total_announce_requests_counter: Counter,
#[cfg(feature = "metrics")]
total_scrape_requests_counter: Counter,
}
impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
async fn run_in_message_loop(&mut self) -> anyhow::Result<()> {
loop {
let message = self
.ws_in
.next()
.await
.ok_or_else(|| anyhow::anyhow!("Stream ended"))??;
match &message {
tungstenite::Message::Text(_) | tungstenite::Message::Binary(_) => {
match InMessage::from_ws_message(message) {
Ok(InMessage::AnnounceRequest(request)) => {
self.handle_announce_request(request).await?;
}
Ok(InMessage::ScrapeRequest(request)) => {
self.handle_scrape_request(request).await?;
}
Err(err) => {
::log::debug!("Couldn't parse in_message: {:#}", err);
self.send_error_response("Invalid request".into(), None, None)
.await?;
}
}
}
tungstenite::Message::Ping(_) => {
::log::trace!("Received ping message");
// tungstenite sends a pong response by itself
}
tungstenite::Message::Pong(_) => {
::log::trace!("Received pong message");
}
tungstenite::Message::Close(_) => {
::log::debug!("Client sent close frame");
break Ok(());
}
tungstenite::Message::Frame(_) => {
::log::warn!("Read raw websocket frame, this should not happen");
}
}
yield_if_needed().await;
}
}
async fn handle_announce_request(&mut self, request: AnnounceRequest) -> anyhow::Result<()> {
#[cfg(feature = "metrics")]
self.total_announce_requests_counter.increment(1);
let info_hash = request.info_hash;
if self
.access_list_cache
.load()
.allows(self.config.access_list.mode, &info_hash.0)
{
let mut announced_info_hashes = self.clean_up_data.announced_info_hashes.borrow_mut();
// Store peer id / check if stored peer id matches
match announced_info_hashes.entry(request.info_hash) {
Entry::Occupied(entry) => {
if *entry.get() != request.peer_id {
// Drop Rc borrow before awaiting
drop(announced_info_hashes);
self.send_error_response(
"Only one peer id can be used per torrent".into(),
Some(ErrorResponseAction::Announce),
Some(info_hash),
)
.await?;
return Err(anyhow::anyhow!(
"Peer used more than one PeerId for a single torrent"
));
}
}
Entry::Vacant(entry) => {
entry.insert(request.peer_id);
// Set peer client info if not set
#[cfg(feature = "metrics")]
if self.config.metrics.run_prometheus_endpoint
&& self.config.metrics.peer_clients
&& self.clean_up_data.opt_peer_client.borrow().is_none()
{
let peer_id = aquatic_peer_id::PeerId(request.peer_id.0);
let peer_client_gauge = ::metrics::gauge!(
"aquatic_peer_clients",
"client" => peer_id.client().to_string(),
);
peer_client_gauge.increment(1.0);
let opt_peer_id_prefix_gauge =
self.config.metrics.peer_id_prefixes.then(|| {
let g = ::metrics::gauge!(
"aquatic_peer_id_prefixes",
"prefix_hex" => peer_id.first_8_bytes_hex().to_string(),
);
g.increment(1.0);
g
});
*self.clean_up_data.opt_peer_client.borrow_mut() =
Some((peer_client_gauge, opt_peer_id_prefix_gauge));
};
}
}
if let Some(AnnounceEvent::Stopped) = request.event {
announced_info_hashes.remove(&request.info_hash);
}
// Drop Rc borrow before awaiting
drop(announced_info_hashes);
let in_message = InMessage::AnnounceRequest(request);
let consumer_index = calculate_in_message_consumer_index(&self.config, info_hash);
// Only fails when receiver is closed
self.in_message_senders
.send_to(
consumer_index,
(self.make_connection_meta(None), in_message),
)
.await
.unwrap();
} else {
self.send_error_response(
"Info hash not allowed".into(),
Some(ErrorResponseAction::Announce),
Some(info_hash),
)
.await?;
}
Ok(())
}
async fn handle_scrape_request(&mut self, request: ScrapeRequest) -> anyhow::Result<()> {
#[cfg(feature = "metrics")]
self.total_scrape_requests_counter.increment(1);
let info_hashes = if let Some(info_hashes) = request.info_hashes {
info_hashes
} else {
// If request.info_hashes is empty, don't return scrape for all
// torrents, even though reference server does it. It is too expensive.
self.send_error_response(
"Full scrapes are not allowed".into(),
Some(ErrorResponseAction::Scrape),
None,
)
.await?;
return Ok(());
};
let mut info_hashes_by_worker: BTreeMap<usize, Vec<InfoHash>> = BTreeMap::new();
for info_hash in info_hashes.as_vec() {
let info_hashes = info_hashes_by_worker
.entry(calculate_in_message_consumer_index(&self.config, info_hash))
.or_default();
info_hashes.push(info_hash);
}
let pending_worker_out_messages = info_hashes_by_worker.len();
let pending_scrape_response = PendingScrapeResponse {
pending_worker_out_messages,
stats: Default::default(),
};
let pending_scrape_id: u8 = self
.pending_scrape_slab
.borrow_mut()
.insert(pending_scrape_response)
.try_into()
.with_context(|| "Reached 256 pending scrape responses")?;
let meta = self.make_connection_meta(Some(PendingScrapeId(pending_scrape_id)));
for (consumer_index, info_hashes) in info_hashes_by_worker {
let in_message = InMessage::ScrapeRequest(ScrapeRequest {
action: ScrapeAction::Scrape,
info_hashes: Some(ScrapeRequestInfoHashes::Multiple(info_hashes)),
});
// Only fails when receiver is closed
self.in_message_senders
.send_to(consumer_index, (meta, in_message))
.await
.unwrap();
}
Ok(())
}
async fn send_error_response(
&self,
failure_reason: Cow<'static, str>,
action: Option<ErrorResponseAction>,
info_hash: Option<InfoHash>,
) -> anyhow::Result<()> {
let out_message = OutMessage::ErrorResponse(ErrorResponse {
action,
failure_reason,
info_hash,
});
self.out_message_sender
.send((self.make_connection_meta(None).into(), out_message))
.await
.map_err(|err| {
anyhow::anyhow!("ConnectionReader::send_error_response failed: {:#}", err)
})
}
fn make_connection_meta(&self, pending_scrape_id: Option<PendingScrapeId>) -> InMessageMeta {
InMessageMeta {
connection_id: self.connection_id,
out_message_consumer_id: self.out_message_consumer_id,
ip_version: self.ip_version,
pending_scrape_id,
}
}
}
struct ConnectionWriter<S> {
config: Rc<Config>,
out_message_receiver: LocalReceiver<(OutMessageMeta, OutMessage)>,
connection_valid_until: Rc<RefCell<ValidUntil>>,
ws_out: SplitSink<WebSocketStream<S>, tungstenite::Message>,
pending_scrape_slab: Rc<RefCell<Slab<PendingScrapeResponse>>>,
server_start_instant: ServerStartInstant,
ip_version: IpVersion,
clean_up_data: ConnectionCleanupData,
}
impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionWriter<S> {
async fn run_out_message_loop(&mut self) -> anyhow::Result<()> {
loop {
let (meta, out_message) = self.out_message_receiver.recv().await.ok_or_else(|| {
anyhow::anyhow!("ConnectionWriter couldn't receive message, sender is closed")
})?;
match out_message {
OutMessage::ScrapeResponse(out_message) => {
let pending_scrape_id = meta
.pending_scrape_id
.expect("meta.pending_scrape_id not set");
let mut pending_responses = self.pending_scrape_slab.borrow_mut();
let pending_response = pending_responses
.get_mut(pending_scrape_id.0 as usize)
.ok_or(anyhow::anyhow!("pending scrape not found in slab"))?;
pending_response.stats.extend(out_message.files);
pending_response.pending_worker_out_messages -= 1;
if pending_response.pending_worker_out_messages == 0 {
let pending_response =
pending_responses.remove(pending_scrape_id.0 as usize);
pending_responses.shrink_to_fit();
let out_message = OutMessage::ScrapeResponse(ScrapeResponse {
action: ScrapeAction::Scrape,
files: pending_response.stats,
});
// Drop Rc borrow before awaiting
drop(pending_responses);
self.send_out_message(&out_message).await?;
}
}
out_message => {
self.send_out_message(&out_message).await?;
}
};
yield_if_needed().await;
}
}
async fn send_out_message(&mut self, out_message: &OutMessage) -> anyhow::Result<()> {
timeout(Duration::from_secs(10), async {
Ok(futures::SinkExt::send(&mut self.ws_out, out_message.to_ws_message()).await)
})
.await
.map_err(|err| {
anyhow::anyhow!("send_out_message: sending to peer took too long: {:#}", err)
})?
.with_context(|| "send_out_message")?;
if let OutMessage::AnnounceResponse(_) | OutMessage::ScrapeResponse(_) = out_message {
*self.connection_valid_until.borrow_mut() = ValidUntil::new(
self.server_start_instant,
self.config.cleaning.max_connection_idle,
);
}
#[cfg(feature = "metrics")]
{
let out_message_type = match &out_message {
OutMessage::OfferOutMessage(_) => "offer",
OutMessage::AnswerOutMessage(_) => "offer_answer",
OutMessage::AnnounceResponse(_) => "announce",
OutMessage::ScrapeResponse(_) => "scrape",
OutMessage::ErrorResponse(_) => "error",
};
::metrics::counter!(
"aquatic_responses_total",
"type" => out_message_type,
"ip_version" => ip_version_to_metrics_str(self.ip_version),
"worker_index" => WORKER_INDEX.with(|index| index.get()).to_string(),
)
.increment(1);
// As long as connection is still alive, increment peer client
// gauges by zero to prevent them from being removed due to
// idleness
if let Some((peer_client_gauge, opt_peer_id_prefix_gauge)) =
self.clean_up_data.opt_peer_client.borrow().as_ref()
{
peer_client_gauge.increment(0.0);
if let Some(g) = opt_peer_id_prefix_gauge {
g.increment(0.0);
}
}
}
Ok(())
}
}
/// Data stored with connection needed for cleanup after it closes
#[derive(Clone)]
struct ConnectionCleanupData {
announced_info_hashes: Rc<RefCell<HashMap<InfoHash, PeerId>>>,
ip_version: IpVersion,
opt_peer_client: Rc<RefCell<Option<(Gauge, Option<Gauge>)>>>,
#[cfg(feature = "metrics")]
active_connections_gauge: Gauge,
}
impl ConnectionCleanupData {
fn before_open(&self) {
#[cfg(feature = "metrics")]
self.active_connections_gauge.increment(1.0);
}
async fn after_close(
&self,
config: &Config,
control_message_senders: Rc<Senders<SwarmControlMessage>>,
) {
let mut announced_info_hashes = HashMap::new();
for (info_hash, peer_id) in self.announced_info_hashes.take().into_iter() {
let consumer_index = calculate_in_message_consumer_index(config, info_hash);
announced_info_hashes
.entry(consumer_index)
.or_insert(Vec::new())
.push((info_hash, peer_id));
}
for (consumer_index, announced_info_hashes) in announced_info_hashes.into_iter() {
let message = SwarmControlMessage::ConnectionClosed {
ip_version: self.ip_version,
announced_info_hashes,
};
control_message_senders
.send_to(consumer_index, message)
.await
.expect("control message receiver open");
}
#[cfg(feature = "metrics")]
self.active_connections_gauge.decrement(1.0);
#[cfg(feature = "metrics")]
if let Some((peer_client_gauge, opt_peer_id_prefix_gauge)) = self.opt_peer_client.take() {
peer_client_gauge.decrement(1.0);
if let Some(g) = opt_peer_id_prefix_gauge {
g.decrement(1.0);
}
}
}
}
struct PendingScrapeResponse {
pending_worker_out_messages: usize,
stats: HashMap<InfoHash, ScrapeStatistics>,
}