diff --git a/services/commtest/tests/tunnelbroker_websocket_messages_tests.rs b/services/commtest/tests/tunnelbroker_websocket_messages_tests.rs new file mode 100644 index 000000000..53cb52a66 --- /dev/null +++ b/services/commtest/tests/tunnelbroker_websocket_messages_tests.rs @@ -0,0 +1,49 @@ +use commtest::identity::device::create_device; +use commtest::identity::olm_account_infos::MOCK_CLIENT_KEYS_1; +use commtest::tunnelbroker::socket::create_socket; +use futures_util::{SinkExt, StreamExt}; +use tokio_tungstenite::tungstenite::{Error, Message, Message::Close}; + +/// Tests for message types defined in tungstenite crate + +#[tokio::test] +async fn test_ping_pong() { + let device = create_device(Some(&MOCK_CLIENT_KEYS_1)).await; + + let ping_message = vec![1, 2, 3, 4, 5]; + + let mut socket = create_socket(&device).await; + socket + .send(Message::Ping(ping_message.clone())) + .await + .expect("Failed to send message"); + + if let Some(Ok(response)) = socket.next().await { + let received_payload = match response { + Message::Pong(received_payload) => received_payload, + unexpected => panic!( + "Unexpected message type or result. Expected Pong, got: {:?}. ", + unexpected + ), + }; + assert_eq!(ping_message.clone(), received_payload); + }; +} + +#[tokio::test] +async fn test_close_message() { + let device = create_device(Some(&MOCK_CLIENT_KEYS_1)).await; + + let mut socket = create_socket(&device).await; + socket + .send(Close(None)) + .await + .expect("Failed to send message"); + + if let Some(response) = socket.next().await { + assert!(matches!( + response, + Err(Error::AlreadyClosed | Error::ConnectionClosed) | Ok(Close(None)) + )); + }; +} diff --git a/services/tunnelbroker/src/websockets/mod.rs b/services/tunnelbroker/src/websockets/mod.rs index e1df464da..8b93b0857 100644 --- a/services/tunnelbroker/src/websockets/mod.rs +++ b/services/tunnelbroker/src/websockets/mod.rs @@ -1,219 +1,236 @@ pub mod session; use crate::database::DatabaseClient; use crate::websockets::session::SessionError; use crate::CONFIG; use futures_util::stream::SplitSink; use futures_util::StreamExt; use hyper::{Body, Request, Response, StatusCode}; use hyper_tungstenite::tungstenite::Message; use hyper_tungstenite::HyperWebsocket; use hyper_tungstenite::WebSocketStream; use std::env; use std::future::Future; use std::net::SocketAddr; use std::pin::Pin; use tokio::io::{AsyncRead, AsyncWrite}; use tokio::net::TcpListener; use tracing::{debug, error, info}; type BoxedError = Box; use self::session::WebsocketSession; /// Hyper HTTP service that handles incoming HTTP and websocket connections /// It handles the initial websocket upgrade request and spawns a task to /// handle the websocket connection. /// It also handles regular HTTP requests (currently health check) struct WebsocketService { addr: SocketAddr, channel: lapin::Channel, db_client: DatabaseClient, } impl hyper::service::Service> for WebsocketService { type Response = Response; type Error = BoxedError; type Future = Pin> + Send>>; // This function is called to check if the service is ready to accept // connections. Since we don't have any state to check, we're always ready. fn poll_ready( &mut self, _: &mut std::task::Context<'_>, ) -> std::task::Poll> { std::task::Poll::Ready(Ok(())) } fn call(&mut self, mut req: Request) -> Self::Future { let addr = self.addr; let db_client = self.db_client.clone(); let channel = self.channel.clone(); let future = async move { // Check if the request is a websocket upgrade request. if hyper_tungstenite::is_upgrade_request(&req) { let (response, websocket) = hyper_tungstenite::upgrade(&mut req, None)?; // Spawn a task to handle the websocket connection. tokio::spawn(async move { accept_connection(websocket, addr, db_client, channel).await; }); // Return the response so the spawned future can continue. return Ok(response); } debug!( "Incoming HTTP request on WebSocket port: {} {}", req.method(), req.uri().path() ); // A simple router for regular HTTP requests let response = match req.uri().path() { "/health" => Response::new(Body::from("OK")), _ => Response::builder() .status(StatusCode::NOT_FOUND) .body(Body::from("Not found"))?, }; Ok(response) }; Box::pin(future) } } pub async fn run_server( db_client: DatabaseClient, amqp_connection: &lapin::Connection, ) -> Result<(), BoxedError> { let addr = env::var("COMM_TUNNELBROKER_WEBSOCKET_ADDR") .unwrap_or_else(|_| format!("0.0.0.0:{}", &CONFIG.http_port)); let listener = TcpListener::bind(&addr).await.expect("Failed to bind"); info!("WebSocket listening on: {}", addr); let mut http = hyper::server::conn::Http::new(); http.http1_only(true); http.http1_keep_alive(true); while let Ok((stream, addr)) = listener.accept().await { let channel = amqp_connection .create_channel() .await .expect("Failed to create AMQP channel"); let connection = http .serve_connection( stream, WebsocketService { channel, db_client: db_client.clone(), addr, }, ) .with_upgrades(); tokio::spawn(async move { if let Err(err) = connection.await { error!("Error serving HTTP/WebSocket connection: {:?}", err); } }); } Ok(()) } /// Handler for any incoming websocket connections async fn accept_connection( hyper_ws: HyperWebsocket, addr: SocketAddr, db_client: DatabaseClient, amqp_channel: lapin::Channel, ) { debug!("Incoming connection from: {}", addr); let ws_stream = match hyper_ws.await { Ok(stream) => stream, Err(e) => { info!( "Failed to establish connection with {}. Reason: {}", addr, e ); return; } }; let (outgoing, mut incoming) = ws_stream.split(); // We don't know the identity of the device until it sends the session // request over the websocket connection let mut session = if let Some(Ok(first_msg)) = incoming.next().await { match initiate_session(outgoing, first_msg, db_client, amqp_channel).await { Ok(session) => session, Err(_) => { error!("Failed to create session with device"); return; } } } else { error!("Failed to create session with device"); return; }; // Poll for messages either being sent to the device (rx) // or messages being received from the device (incoming) loop { debug!("Polling for messages from: {}", addr); tokio::select! { Some(Ok(delivery)) = session.next_amqp_message() => { if let Ok(message) = std::str::from_utf8(&delivery.data) { - session.send_message_to_device(message.to_string()).await; + session.send_message_to_device(Message::Text(message.to_string())).await; } else { error!("Invalid payload"); } }, device_message = incoming.next() => { - match device_message { - Some(Ok(msg)) => { - session::consume_error(session.handle_websocket_frame_from_device(msg).await); - } + let message: Message = match device_message { + Some(Ok(msg)) => msg, _ => { debug!("Connection to {} closed remotely.", addr); break; } + }; + match message { + Message::Close(_) => { + debug!("Connection to {} closed.", addr); + break; + } + Message::Pong(_) => { + debug!("Received Pong message from {}", addr); + } + Message::Ping(msg) => { + debug!("Received Ping message from {}", addr); + session.send_message_to_device(Message::Pong(msg)).await; + } + Message::Text(msg) => { + session::consume_error(session.handle_websocket_frame_from_device(msg).await); + } + _ => { + error!("Client sent invalid message type"); + } } }, else => { debug!("Unhealthy connection for: {}", addr); break; }, } } info!("Unregistering connection to: {}", addr); session.close().await } async fn initiate_session( outgoing: SplitSink, Message>, frame: Message, db_client: DatabaseClient, amqp_channel: lapin::Channel, ) -> Result, session::SessionError> { let mut session = session::WebsocketSession::from_frame( outgoing, db_client.clone(), frame, &amqp_channel, ) .await .map_err(|_| { error!("Device failed to send valid connection request."); SessionError::InvalidMessage })?; session::consume_error(session.deliver_persisted_messages().await); Ok(session) } diff --git a/services/tunnelbroker/src/websockets/session.rs b/services/tunnelbroker/src/websockets/session.rs index 0b4bc2e75..3d59bf0aa 100644 --- a/services/tunnelbroker/src/websockets/session.rs +++ b/services/tunnelbroker/src/websockets/session.rs @@ -1,275 +1,269 @@ use aws_sdk_dynamodb::error::SdkError; use aws_sdk_dynamodb::operation::put_item::PutItemError; use derive_more; use futures_util::stream::SplitSink; use futures_util::SinkExt; use futures_util::StreamExt; use hyper_tungstenite::{tungstenite::Message, WebSocketStream}; use lapin::message::Delivery; use lapin::options::{ BasicCancelOptions, BasicConsumeOptions, BasicPublishOptions, QueueDeclareOptions, QueueDeleteOptions, }; use lapin::types::FieldTable; use lapin::BasicProperties; use tokio::io::AsyncRead; use tokio::io::AsyncWrite; use tracing::{debug, error, info}; use tunnelbroker_messages::{session::DeviceTypes, Messages}; use crate::database::{self, DatabaseClient, DeviceMessage}; use crate::error::Error; use crate::identity; pub struct DeviceInfo { pub device_id: String, pub notify_token: Option, pub device_type: DeviceTypes, pub device_app_version: Option, pub device_os: Option, } pub struct WebsocketSession { tx: SplitSink, Message>, db_client: DatabaseClient, pub device_info: DeviceInfo, amqp_channel: lapin::Channel, // Stream of messages from AMQP endpoint amqp_consumer: lapin::Consumer, } #[derive( Debug, derive_more::Display, derive_more::From, derive_more::Error, )] pub enum SessionError { InvalidMessage, SerializationError(serde_json::Error), MessageError(database::MessageErrors), AmqpError(lapin::Error), InternalError, UnauthorizedDevice, PersistenceError(SdkError), } pub fn consume_error(result: Result) { if let Err(e) = result { error!("{}", e) } } // Parse a session request and retrieve the device information pub async fn handle_first_message_from_device( message: &str, ) -> Result { let serialized_message = serde_json::from_str::(message)?; match serialized_message { Messages::ConnectionInitializationMessage(mut session_info) => { let device_info = DeviceInfo { device_id: session_info.device_id.clone(), notify_token: session_info.notify_token.take(), device_type: session_info.device_type, device_app_version: session_info.device_app_version.take(), device_os: session_info.device_os.take(), }; // Authenticate device debug!("Authenticating device: {}", &session_info.device_id); let auth_request = identity::verify_user_access_token( &session_info.user_id, &device_info.device_id, &session_info.access_token, ) .await; match auth_request { Err(e) => { error!("Failed to complete request to identity service: {:?}", e); return Err(SessionError::InternalError.into()); } Ok(false) => { info!("Device failed authentication: {}", &session_info.device_id); return Err(SessionError::UnauthorizedDevice.into()); } Ok(true) => { debug!( "Successfully authenticated device: {}", &session_info.device_id ); } } Ok(device_info) } _ => { debug!("Received invalid request"); Err(SessionError::InvalidMessage.into()) } } } impl WebsocketSession { pub async fn from_frame( tx: SplitSink, Message>, db_client: DatabaseClient, frame: Message, amqp_channel: &lapin::Channel, ) -> Result, Error> { let device_info = match frame { Message::Text(payload) => { handle_first_message_from_device(&payload).await? } _ => { error!("Client sent wrong frame type for establishing connection"); return Err(SessionError::InvalidMessage.into()); } }; // We don't currently have a use case to interact directly with the queue, // however, we need to declare a queue for a given device amqp_channel .queue_declare( &device_info.device_id, QueueDeclareOptions::default(), FieldTable::default(), ) .await?; let amqp_consumer = amqp_channel .basic_consume( &device_info.device_id, "tunnelbroker", BasicConsumeOptions::default(), FieldTable::default(), ) .await?; Ok(WebsocketSession { tx, db_client, device_info, amqp_channel: amqp_channel.clone(), amqp_consumer, }) } pub async fn handle_websocket_frame_from_device( &self, - msg: Message, + msg: String, ) -> Result<(), SessionError> { - let text_msg = match msg { - Message::Text(payload) => payload, - _ => { - error!("Client sent invalid message type"); - return Err(SessionError::InvalidMessage); - } - }; - - let serialized_message = serde_json::from_str::(&text_msg)?; + let serialized_message = serde_json::from_str::(&msg)?; match serialized_message { Messages::MessageToDevice(message_to_device) => { debug!("Received message for {}", message_to_device.device_id); self .db_client .persist_message( message_to_device.device_id.as_str(), message_to_device.payload.as_str(), "message_id", ) .await?; self .amqp_channel .basic_publish( "", &message_to_device.device_id, BasicPublishOptions::default(), message_to_device.payload.as_bytes(), BasicProperties::default(), ) .await?; } _ => { error!("Client sent invalid message type"); return Err(SessionError::InvalidMessage); } } Ok(()) } pub async fn next_amqp_message( &mut self, ) -> Option> { self.amqp_consumer.next().await } pub async fn deliver_persisted_messages( &mut self, ) -> Result<(), SessionError> { // Check for persisted messages let messages = self .db_client .retrieve_messages(&self.device_info.device_id) .await .unwrap_or_else(|e| { error!("Error while retrieving messages: {}", e); Vec::new() }); for message in messages { let device_message = DeviceMessage::from_hashmap(message)?; - self.send_message_to_device(device_message.payload).await; + self + .send_message_to_device(Message::Text(device_message.payload)) + .await; if let Err(e) = self .db_client .delete_message(&self.device_info.device_id, &device_message.message_id) .await { error!("Failed to delete message: {}:", e); } } debug!( "Flushed messages for device: {}", &self.device_info.device_id ); Ok(()) } - pub async fn send_message_to_device(&mut self, incoming_payload: String) { - if let Err(e) = self.tx.send(Message::Text(incoming_payload)).await { + pub async fn send_message_to_device(&mut self, message: Message) { + if let Err(e) = self.tx.send(message).await { error!("Failed to send message to device: {}", e); } } // Release WebSocket and remove from active connections pub async fn close(&mut self) { if let Err(e) = self.tx.close().await { debug!("Failed to close WebSocket session: {}", e); } if let Err(e) = self .amqp_channel .basic_cancel( self.amqp_consumer.tag().as_str(), BasicCancelOptions::default(), ) .await { error!("Failed to cancel consumer: {}", e); } if let Err(e) = self .amqp_channel .queue_delete( self.device_info.device_id.as_str(), QueueDeleteOptions::default(), ) .await { error!("Failed to delete queue: {}", e); } } }