Jon Häggblad
499 lines
19 KiB
Rust
499 lines
19 KiB
Rust
// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
use futures::channel::mpsc;
|
|
use futures::{SinkExt, StreamExt};
|
|
use log::*;
|
|
use nym_client_core::client::replies::reply_controller::requests::ReplyControllerSender;
|
|
use nym_client_core::client::{
|
|
inbound_messages::{InputMessage, InputMessageSender},
|
|
received_buffer::{
|
|
ReceivedBufferMessage, ReceivedBufferRequestSender, ReconstructedMessagesReceiver,
|
|
},
|
|
};
|
|
use nym_client_websocket_requests::{requests::ClientRequest, responses::ServerResponse};
|
|
use nym_sphinx::addressing::clients::Recipient;
|
|
use nym_sphinx::anonymous_replies::requests::AnonymousSenderTag;
|
|
use nym_sphinx::params::PacketType;
|
|
use nym_sphinx::receiver::ReconstructedMessage;
|
|
use nym_task::connections::{
|
|
ConnectionCommand, ConnectionCommandSender, ConnectionId, LaneQueueLengths, TransmissionLane,
|
|
};
|
|
use nym_task::TaskClient;
|
|
use std::time::Duration;
|
|
use tokio::net::TcpStream;
|
|
use tokio::time::Instant;
|
|
use tokio_tungstenite::{
|
|
accept_async,
|
|
tungstenite::{protocol::Message as WsMessage, Error as WsError},
|
|
WebSocketStream,
|
|
};
|
|
|
|
#[derive(Default)]
|
|
enum ReceivedResponseType {
|
|
#[default]
|
|
Binary,
|
|
Text,
|
|
}
|
|
|
|
pub(crate) struct HandlerBuilder {
|
|
msg_input: InputMessageSender,
|
|
client_connection_tx: ConnectionCommandSender,
|
|
buffer_requester: ReceivedBufferRequestSender,
|
|
self_full_address: Recipient,
|
|
lane_queue_lengths: LaneQueueLengths,
|
|
reply_controller_sender: ReplyControllerSender,
|
|
packet_type: Option<PacketType>,
|
|
task_client: TaskClient,
|
|
}
|
|
|
|
impl HandlerBuilder {
|
|
#[allow(clippy::too_many_arguments)]
|
|
pub(crate) fn new(
|
|
msg_input: InputMessageSender,
|
|
client_connection_tx: ConnectionCommandSender,
|
|
buffer_requester: ReceivedBufferRequestSender,
|
|
self_full_address: &Recipient,
|
|
lane_queue_lengths: LaneQueueLengths,
|
|
reply_controller_sender: ReplyControllerSender,
|
|
packet_type: Option<PacketType>,
|
|
task_client: TaskClient,
|
|
) -> Self {
|
|
Self {
|
|
msg_input,
|
|
client_connection_tx,
|
|
buffer_requester,
|
|
self_full_address: *self_full_address,
|
|
lane_queue_lengths,
|
|
reply_controller_sender,
|
|
packet_type,
|
|
task_client,
|
|
}
|
|
}
|
|
|
|
// TODO: make sure we only ever have one active handler
|
|
pub fn create_active_handler(&self) -> Handler {
|
|
let mut task_client = self.task_client.fork("active_handler");
|
|
task_client.disarm();
|
|
Handler {
|
|
msg_input: self.msg_input.clone(),
|
|
client_connection_tx: self.client_connection_tx.clone(),
|
|
buffer_requester: self.buffer_requester.clone(),
|
|
self_full_address: self.self_full_address,
|
|
socket: None,
|
|
received_response_type: Default::default(),
|
|
lane_queue_lengths: self.lane_queue_lengths.clone(),
|
|
reply_controller_sender: self.reply_controller_sender.clone(),
|
|
packet_type: self.packet_type,
|
|
task_client,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) struct Handler {
|
|
msg_input: InputMessageSender,
|
|
client_connection_tx: ConnectionCommandSender,
|
|
buffer_requester: ReceivedBufferRequestSender,
|
|
self_full_address: Recipient,
|
|
socket: Option<WebSocketStream<TcpStream>>,
|
|
received_response_type: ReceivedResponseType,
|
|
lane_queue_lengths: LaneQueueLengths,
|
|
reply_controller_sender: ReplyControllerSender,
|
|
packet_type: Option<PacketType>,
|
|
task_client: TaskClient,
|
|
}
|
|
|
|
impl Drop for Handler {
|
|
fn drop(&mut self) {
|
|
if let Err(err) = self
|
|
.buffer_requester
|
|
.unbounded_send(ReceivedBufferMessage::ReceiverDisconnect)
|
|
{
|
|
if !self.task_client.is_shutdown_poll() {
|
|
error!("failed to disconnect the receiver from the buffer: {err}");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Handler {
|
|
async fn get_lane_queue_length(&self, connection_id: ConnectionId) -> Option<ServerResponse> {
|
|
let req_start = Instant::now();
|
|
|
|
// get the base queue length
|
|
// Note that this does _NOT_ take into account the packets that have been received but not
|
|
// yet reach `OutQueueControl`, so it might be a tad low.
|
|
let conn_lane = TransmissionLane::ConnectionId(connection_id);
|
|
let Ok(base_length) = self
|
|
.lane_queue_lengths
|
|
.lock()
|
|
.map(|guard| guard.get(&conn_lane).unwrap_or_default())
|
|
else {
|
|
// I'd argue we should panic here as this error it not recoverable
|
|
error!("The lane queue length lock is poisoned!!");
|
|
return None;
|
|
};
|
|
|
|
// get the number of pending replies waiting for reply surbs
|
|
let reply_queue_length = match self
|
|
.reply_controller_sender
|
|
.get_lane_queue_length(connection_id)
|
|
.await
|
|
{
|
|
Ok(length) => length,
|
|
Err(err) => {
|
|
if !self.task_client.is_shutdown_poll() {
|
|
error!(
|
|
"Failed to get reply queue length for connection {connection_id}: {err}"
|
|
);
|
|
}
|
|
// We're just going to assume that the queue is empty, and I think that's okay
|
|
// during shutdown.
|
|
0
|
|
}
|
|
};
|
|
|
|
let queue_length = base_length + reply_queue_length;
|
|
|
|
let time_taken = req_start.elapsed();
|
|
let msg =
|
|
format!("it took {time_taken:?} to get lane length for connection {connection_id}. The length is: {queue_length} = {base_length} (already queued up) + {reply_queue_length} (waiting for reply SURBs)");
|
|
|
|
if time_taken > Duration::from_millis(1) {
|
|
info!("{msg}");
|
|
} else if time_taken > Duration::from_millis(10) {
|
|
warn!("{msg}");
|
|
} else if time_taken > Duration::from_millis(50) {
|
|
error!("{msg}");
|
|
}
|
|
|
|
Some(ServerResponse::LaneQueueLength {
|
|
lane: connection_id,
|
|
queue_length,
|
|
})
|
|
}
|
|
|
|
async fn handle_send(
|
|
&mut self,
|
|
recipient: Recipient,
|
|
message: Vec<u8>,
|
|
connection_id: Option<u64>,
|
|
) -> Option<ServerResponse> {
|
|
info!(
|
|
"Attempting to send {:.2} kiB message to {recipient} on connection_id {connection_id:?}",
|
|
message.len() as f64 / 1024.0
|
|
);
|
|
|
|
// We map the absence of a connection id as going into the general lane.
|
|
let lane = connection_id.map_or(TransmissionLane::General, |id| {
|
|
TransmissionLane::ConnectionId(id)
|
|
});
|
|
|
|
// the ack control is now responsible for chunking, etc.
|
|
let input_msg = InputMessage::new_regular(recipient, message, lane, self.packet_type);
|
|
if let Err(err) = self.msg_input.send(input_msg).await {
|
|
if !self.task_client.is_shutdown_poll() {
|
|
error!("Failed to send message to the input buffer: {err}");
|
|
}
|
|
}
|
|
|
|
// Only reply back with a `LaneQueueLength` if the sender providided a connection id
|
|
let TransmissionLane::ConnectionId(connection_id) = lane else {
|
|
return None;
|
|
};
|
|
|
|
self.get_lane_queue_length(connection_id).await
|
|
}
|
|
|
|
async fn handle_send_anonymous(
|
|
&mut self,
|
|
recipient: Recipient,
|
|
message: Vec<u8>,
|
|
reply_surbs: u32,
|
|
connection_id: Option<u64>,
|
|
) -> Option<ServerResponse> {
|
|
info!(
|
|
"Attempting to anonymously send {:.2} kiB message to {recipient} on connection_id {connection_id:?} while attaching {reply_surbs} replySURBs.",
|
|
message.len() as f64 / 1024.0
|
|
);
|
|
|
|
// We map the absence of a connection id as going into the general lane.
|
|
let lane = connection_id.map_or(TransmissionLane::General, |id| {
|
|
TransmissionLane::ConnectionId(id)
|
|
});
|
|
|
|
let input_msg =
|
|
InputMessage::new_anonymous(recipient, message, reply_surbs, lane, self.packet_type);
|
|
if let Err(err) = self.msg_input.send(input_msg).await {
|
|
if !self.task_client.is_shutdown_poll() {
|
|
error!("Failed to send anonymous message to the input buffer: {err}");
|
|
}
|
|
}
|
|
|
|
// Only reply back with a `LaneQueueLength` if the sender providided a connection id
|
|
let TransmissionLane::ConnectionId(connection_id) = lane else {
|
|
return None;
|
|
};
|
|
|
|
self.get_lane_queue_length(connection_id).await
|
|
}
|
|
|
|
async fn handle_reply(
|
|
&mut self,
|
|
recipient_tag: AnonymousSenderTag,
|
|
message: Vec<u8>,
|
|
connection_id: Option<u64>,
|
|
) -> Option<ServerResponse> {
|
|
info!("Attempting to send {:.2} kiB reply message to {recipient_tag} on connection_id {connection_id:?}", message.len() as f64 / 1024.0);
|
|
|
|
// We map the absence of a connection id as going into the general lane.
|
|
let lane = connection_id.map_or(TransmissionLane::General, |id| {
|
|
TransmissionLane::ConnectionId(id)
|
|
});
|
|
|
|
let input_msg = InputMessage::new_reply(recipient_tag, message, lane, self.packet_type);
|
|
if let Err(err) = self.msg_input.send(input_msg).await {
|
|
if !self.task_client.is_shutdown_poll() {
|
|
error!("Failed to send reply message to the input buffer: {err}");
|
|
}
|
|
}
|
|
|
|
// Only reply back with a `LaneQueueLength` if the sender providided a connection id
|
|
let TransmissionLane::ConnectionId(connection_id) = lane else {
|
|
return None;
|
|
};
|
|
|
|
self.get_lane_queue_length(connection_id).await
|
|
}
|
|
|
|
fn handle_self_address(&self) -> ServerResponse {
|
|
ServerResponse::SelfAddress(Box::new(self.self_full_address))
|
|
}
|
|
|
|
fn handle_closed_connection(&self, connection_id: u64) -> Option<ServerResponse> {
|
|
if let Err(err) = self
|
|
.client_connection_tx
|
|
.unbounded_send(ConnectionCommand::Close(connection_id))
|
|
{
|
|
if !self.task_client.is_shutdown_poll() {
|
|
error!("Failed to send close connection command: {err}");
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
async fn handle_get_lane_queue_length(&self, connection_id: u64) -> Option<ServerResponse> {
|
|
self.get_lane_queue_length(connection_id).await
|
|
}
|
|
|
|
async fn handle_request(&mut self, request: ClientRequest) -> Option<ServerResponse> {
|
|
match request {
|
|
ClientRequest::Send {
|
|
recipient,
|
|
message,
|
|
connection_id,
|
|
} => self.handle_send(recipient, message, connection_id).await,
|
|
|
|
ClientRequest::SendAnonymous {
|
|
recipient,
|
|
message,
|
|
reply_surbs,
|
|
connection_id,
|
|
} => {
|
|
self.handle_send_anonymous(recipient, message, reply_surbs, connection_id)
|
|
.await
|
|
}
|
|
|
|
ClientRequest::Reply {
|
|
message,
|
|
sender_tag,
|
|
connection_id,
|
|
} => self.handle_reply(sender_tag, message, connection_id).await,
|
|
|
|
ClientRequest::SelfAddress => Some(self.handle_self_address()),
|
|
ClientRequest::ClosedConnection(id) => self.handle_closed_connection(id),
|
|
ClientRequest::GetLaneQueueLength(id) => self.handle_get_lane_queue_length(id).await,
|
|
}
|
|
}
|
|
|
|
async fn handle_text_message(&mut self, msg: String) -> Option<WsMessage> {
|
|
debug!("Handling text message request");
|
|
trace!("Content: {:?}", msg);
|
|
|
|
self.received_response_type = ReceivedResponseType::Text;
|
|
let client_request = ClientRequest::try_from_text(msg);
|
|
|
|
let response = match client_request {
|
|
Err(err) => Some(ServerResponse::Error(err)),
|
|
Ok(req) => self.handle_request(req).await,
|
|
};
|
|
|
|
response.map(|resp| WsMessage::text(resp.into_text()))
|
|
}
|
|
|
|
async fn handle_binary_message(&mut self, msg: &[u8]) -> Option<WsMessage> {
|
|
debug!("Handling binary message request");
|
|
|
|
self.received_response_type = ReceivedResponseType::Binary;
|
|
let client_request = ClientRequest::try_from_binary(msg);
|
|
|
|
let response = match client_request {
|
|
Err(err) => Some(ServerResponse::Error(err)),
|
|
Ok(req) => self.handle_request(req).await,
|
|
};
|
|
|
|
response.map(|resp| WsMessage::Binary(resp.into_binary()))
|
|
}
|
|
|
|
async fn handle_ws_request(&mut self, raw_request: WsMessage) -> Option<WsMessage> {
|
|
// apparently tungstenite auto-handles ping/pong/close messages so for now let's ignore
|
|
// them and let's test that claim. If that's not the case, just copy code from
|
|
// old version of this file.
|
|
match raw_request {
|
|
WsMessage::Text(text_message) => self.handle_text_message(text_message).await,
|
|
WsMessage::Binary(binary_message) => self.handle_binary_message(&binary_message).await,
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
async fn push_websocket_received_plaintexts(
|
|
&mut self,
|
|
reconstructed_messages: Vec<ReconstructedMessage>,
|
|
) -> Result<(), WsError> {
|
|
// TODO: later there might be a flag on the reconstructed message itself to tell us
|
|
// if it's text or binary, but for time being we use the naive assumption that if
|
|
// client is sending Message::Text it expects text back. Same for Message::Binary
|
|
let response_messages = match self.received_response_type {
|
|
ReceivedResponseType::Binary => prepare_reconstructed_binary(reconstructed_messages),
|
|
ReceivedResponseType::Text => prepare_reconstructed_text(reconstructed_messages),
|
|
};
|
|
|
|
let mut send_stream = futures::stream::iter(response_messages);
|
|
self.socket
|
|
.as_mut()
|
|
.unwrap()
|
|
.send_all(&mut send_stream)
|
|
.await
|
|
}
|
|
|
|
async fn send_websocket_response(&mut self, msg: WsMessage) -> Result<(), WsError> {
|
|
match self.socket {
|
|
// TODO: more closely investigate difference between `Sink::send` and `Sink::send_all`
|
|
// it got something to do with batching and flushing - it might be important if it
|
|
// turns out somehow we've got a bottleneck here
|
|
Some(ref mut ws_stream) => ws_stream.send(msg).await,
|
|
_ => panic!("impossible state - websocket handshake was somehow reverted"),
|
|
}
|
|
}
|
|
|
|
async fn next_websocket_request(&mut self) -> Option<Result<WsMessage, WsError>> {
|
|
match self.socket {
|
|
Some(ref mut ws_stream) => ws_stream.next().await,
|
|
None => None,
|
|
}
|
|
}
|
|
|
|
async fn listen_for_requests(&mut self, mut msg_receiver: ReconstructedMessagesReceiver) {
|
|
let mut task_client = self.task_client.fork("select");
|
|
task_client.disarm();
|
|
|
|
while !task_client.is_shutdown() {
|
|
tokio::select! {
|
|
// we can either get a client request from the websocket
|
|
socket_msg = self.next_websocket_request() => {
|
|
if socket_msg.is_none() {
|
|
break;
|
|
}
|
|
let socket_msg = match socket_msg.unwrap() {
|
|
Ok(socket_msg) => socket_msg,
|
|
Err(err) => {
|
|
warn!("failed to obtain message from websocket stream! stopping connection handler: {err}");
|
|
break;
|
|
}
|
|
};
|
|
|
|
if socket_msg.is_close() {
|
|
break;
|
|
}
|
|
|
|
if let Some(response) = self.handle_ws_request(socket_msg).await {
|
|
if let Err(err) = self.send_websocket_response(response).await {
|
|
warn!(
|
|
"Failed to send message over websocket: {err}. Assuming the connection is dead.",
|
|
);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
// or a reconstructed mix message that we need to push back to the client
|
|
mix_messages = msg_receiver.next() => {
|
|
let Some(mix_messages) = mix_messages else {
|
|
error!("mix messages sender was unexpectedly closed! this shouldn't have ever happened! (unless we're shutting down - TODO: implement proper graceful shutdown handler)");
|
|
return
|
|
};
|
|
if let Err(err) = self.push_websocket_received_plaintexts(mix_messages).await {
|
|
warn!("failed to send sphinx packets back to the client - {err}, assuming the connection is dead");
|
|
break;
|
|
}
|
|
}
|
|
_ = task_client.recv() => {
|
|
log::trace!("Websocket handler: Received shutdown");
|
|
}
|
|
}
|
|
}
|
|
log::debug!("Websocket handler: Exiting");
|
|
}
|
|
|
|
// consume self to make sure `drop` is called after this is done
|
|
pub(crate) async fn handle_connection(mut self, socket: TcpStream) {
|
|
let ws_stream = match accept_async(socket).await {
|
|
Ok(ws_stream) => ws_stream,
|
|
Err(err) => {
|
|
warn!("error while performing the websocket handshake - {err}");
|
|
return;
|
|
}
|
|
};
|
|
self.socket = Some(ws_stream);
|
|
|
|
let (reconstructed_sender, reconstructed_receiver) = mpsc::unbounded();
|
|
|
|
// tell the buffer to start sending stuff to us
|
|
if let Err(err) =
|
|
self.buffer_requester
|
|
.unbounded_send(ReceivedBufferMessage::ReceiverAnnounce(
|
|
reconstructed_sender,
|
|
))
|
|
{
|
|
if !self.task_client.is_shutdown_poll() {
|
|
error!("failed to announce the receiver to the buffer: {err}");
|
|
}
|
|
}
|
|
|
|
self.listen_for_requests(reconstructed_receiver).await;
|
|
}
|
|
}
|
|
|
|
// I'm still not entirely sure why `send_all` requires `TryStream` rather than `Stream`, but
|
|
// let's just play along for now
|
|
fn prepare_reconstructed_binary(
|
|
reconstructed_messages: Vec<ReconstructedMessage>,
|
|
) -> Vec<Result<WsMessage, WsError>> {
|
|
reconstructed_messages
|
|
.into_iter()
|
|
.map(ServerResponse::Received)
|
|
.map(|resp| Ok(WsMessage::Binary(resp.into_binary())))
|
|
.collect()
|
|
}
|
|
|
|
// I'm still not entirely sure why `send_all` requires `TryStream` rather than `Stream`, but
|
|
// let's just play along for now
|
|
fn prepare_reconstructed_text(
|
|
reconstructed_messages: Vec<ReconstructedMessage>,
|
|
) -> Vec<Result<WsMessage, WsError>> {
|
|
reconstructed_messages
|
|
.into_iter()
|
|
.map(ServerResponse::Received)
|
|
.map(|resp| Ok(WsMessage::Text(resp.into_text())))
|
|
.collect()
|
|
}
|