Websocket handling 'Ping', 'Pong' and 'Close' messages + general improvements

Also fixes https://github.com/nymtech/nym/issues/56
This commit is contained in:
Jedrzej Stuczynski
2020-01-20 16:49:13 +00:00
parent 35086ec63e
commit 33dafd9bba
+143 -38
View File
@@ -9,9 +9,11 @@ use futures::{SinkExt, StreamExt};
use log::{debug, error, info, trace, warn};
use serde::{Deserialize, Serialize};
use sphinx::route::{Destination, DestinationAddressBytes};
use std::convert::TryFrom;
use std::io;
use std::net::SocketAddr;
use tungstenite::protocol::Message;
use tungstenite::protocol::frame::coding::CloseCode;
use tungstenite::protocol::{CloseFrame, Message};
struct Connection {
address: SocketAddr,
@@ -23,6 +25,122 @@ struct Connection {
tx: UnboundedSender<Message>,
}
impl Connection {
async fn handle_text_message(&self, msg: String) -> ServerResponse {
debug!("Handling text message request");
trace!("Content: {:?}", msg.clone());
let request = match ClientRequest::try_from(msg) {
Ok(req) => req,
Err(err) => {
return ServerResponse::Error {
// we failed to parse the request
message: format!("received invalid request. err: {:?}", err),
};
}
};
match request {
ClientRequest::Send {
message,
recipient_address,
} => {
ClientRequest::handle_send(message, recipient_address, self.msg_input.clone()).await
}
ClientRequest::Fetch => ClientRequest::handle_fetch(self.msg_query.clone()).await,
ClientRequest::GetClients => {
ClientRequest::handle_get_clients(self.topology.clone()).await
}
ClientRequest::OwnDetails => ClientRequest::handle_own_details(self.self_address).await,
}
}
// Currently our websocket cannot handle binary data, so just close the connection
// with unsupported close code.
async fn handle_binary_message(&self, _msg: Vec<u8>) -> Message {
debug!("Handling binary message request");
Message::Close(Some(CloseFrame {
code: CloseCode::Unsupported,
reason: "binary messages aren't yet supported".into(),
}))
}
// As per RFC6455 5.5.2. and 5.5.3.:
// Upon receipt of a Ping frame, an endpoint MUST send a Pong frame in response.
// A Pong frame sent in response to a Ping frame must have identical
// "Application data" as found in the message body of the Ping frame
// being replied to.
async fn handle_ping_message(&self, msg: Vec<u8>) -> Message {
debug!("Handling binary ping request");
// As per RFC6455 5.5:
// All control frames MUST have a payload length of 125 bytes or less
if msg.len() > 125 {
return Message::Close(Some(CloseFrame {
code: CloseCode::Protocol,
reason: format!("ping message of length {} sent", msg.len()).into(),
}));
}
Message::Pong(msg)
}
// As per RFC6455 5.5.3.:
// A Pong frame MAY be sent unsolicited. This serves as a
// unidirectional heartbeat. A response to an unsolicited Pong frame is
// not expected.
// Realistically this handler should never be used,
// but since we're nice we will reply with a Pong containing original content
async fn handle_pong_message(&self, msg: Vec<u8>) -> Message {
debug!("Handling pong message request");
// As per RFC6455 5.5:
// All control frames MUST have a payload length of 125 bytes or less
if msg.len() > 125 {
return Message::Close(Some(CloseFrame {
code: CloseCode::Protocol,
reason: format!("ping message of length {} sent", msg.len()).into(),
}));
}
Message::Pong(msg)
}
// As per RFC6455 5.5.1.:
// If an endpoint receives a Close frame and did not previously send a
// Close frame, the endpoint MUST send a Close frame in response. (When
// sending a Close frame in response, the endpoint typically echos the
// status code it received.)
async fn handle_close_message(&self, close_frame: Option<CloseFrame<'static>>) -> Message {
debug!("Handling close message request");
Message::Close(close_frame)
}
async fn handle(mut self) {
while let Some(msg) = self.rx.next().await {
trace!("Received a message from {}: {}", self.address, msg);
let response_message = match msg {
Message::Text(text_message) => self.handle_text_message(text_message).await.into(),
Message::Binary(binary_message) => self.handle_binary_message(binary_message).await,
Message::Ping(ping_message) => self.handle_ping_message(ping_message).await,
Message::Pong(pong_message) => self.handle_pong_message(pong_message).await,
Message::Close(close_frame) => self.handle_close_message(close_frame).await,
};
if let Err(err) = self.tx.unbounded_send(response_message) {
error!(
"Failed to send response message to accepted connections handler: {}.\n\
Shutting off the connection handler",
err
);
return;
}
}
}
}
#[derive(Debug)]
pub enum WebSocketError {
FailedToStartSocketError,
@@ -37,7 +155,6 @@ impl From<io::Error> for WebSocketError {
io::ErrorKind::ConnectionReset => FailedToStartSocketError,
io::ErrorKind::ConnectionAborted => FailedToStartSocketError,
io::ErrorKind::NotConnected => FailedToStartSocketError,
io::ErrorKind::AddrInUse => FailedToStartSocketError,
io::ErrorKind::AddrNotAvailable => FailedToStartSocketError,
_ => UnknownSocketError,
@@ -57,15 +174,11 @@ enum ClientRequest {
OwnDetails,
}
impl From<Message> for ClientRequest {
fn from(msg: Message) -> Self {
let text_msg = match msg {
Message::Text(msg) => msg,
Message::Binary(_) => panic!("binary messages are not supported!"),
Message::Close(_) => panic!("todo: handle close!"),
_ => panic!("Other types of messages are also unsupported!"),
};
serde_json::from_str(&text_msg).expect("unable to deserialize From<Message> json")
impl TryFrom<String> for ClientRequest {
type Error = serde_json::Error;
fn try_from(msg: String) -> Result<Self, Self::Error> {
serde_json::from_str(&msg)
}
}
@@ -193,32 +306,6 @@ impl Into<Message> for ServerResponse {
}
}
async fn handle_connection(conn: Connection) {
let mut conn = conn;
while let Some(msg) = conn.rx.next().await {
println!("Received a message from {}: {}", conn.address, msg);
let request: ClientRequest = msg.into();
let response = match request {
ClientRequest::Send {
message,
recipient_address,
} => {
ClientRequest::handle_send(message, recipient_address, conn.msg_input.clone()).await
}
ClientRequest::Fetch => ClientRequest::handle_fetch(conn.msg_query.clone()).await,
ClientRequest::GetClients => {
ClientRequest::handle_get_clients(conn.topology.clone()).await
}
ClientRequest::OwnDetails => ClientRequest::handle_own_details(conn.self_address).await,
};
conn.tx
.unbounded_send(response.into())
.expect("Failed to forward message");
}
}
async fn accept_connection(
stream: tokio::net::TcpStream,
msg_input: mpsc::UnboundedSender<InputMessage>,
@@ -250,7 +337,7 @@ async fn accept_connection(
msg_query,
self_address,
};
tokio::spawn(handle_connection(conn));
tokio::spawn(conn.handle());
while let Some(message) = ws_stream.next().await {
let message = match message {
@@ -260,6 +347,12 @@ async fn accept_connection(
return;
}
};
let mut should_close = false;
if message.is_close() {
should_close = true;
}
if let Err(err) = msg_tx.unbounded_send(message) {
error!(
"Failed to forward request. Closing the socket connection: {}",
@@ -267,6 +360,14 @@ async fn accept_connection(
);
return;
}
// if the received message is a close message it means we will reply with a close
// or it is a reply to our close, either way as per RFC6455 5.5.1 we should close the
// underlying TCP connection:
// After both sending and receiving a Close message, an endpoint
// considers the WebSocket connection closed and MUST close the
// underlying TCP connection. The server MUST close the underlying TCP
// connection immediately;
if let Some(resp) = response_rx.next().await {
if let Err(err) = ws_stream.send(resp).await {
warn!(
@@ -276,6 +377,10 @@ async fn accept_connection(
return;
}
}
if should_close {
info!("Closing the websocket connection");
return;
}
}
}