Add configurable gRPC and WebSocket message size limit (#487)

crates/yaak-ws/src/manager.rs

@@ -1,4 +1,5 @@
-use crate::connect::ws_connect;
+use crate::connect::{message_size_limit, ws_connect};
+use crate::error::Error::GenericError;
 use crate::error::Result;
 use futures_util::stream::SplitSink;
 use futures_util::{SinkExt, StreamExt};
@@ -15,10 +16,16 @@ use tokio_tungstenite::tungstenite::http::HeaderValue;
 use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
 use yaak_tls::ClientCertificateConfig;
 
+type WebsocketSink = SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>;
+
+struct WebsocketConnection {
+    max_message_size: Option<usize>,
+    sink: WebsocketSink,
+}
+
 #[derive(Clone)]
 pub struct WebsocketManager {
-    connections:
-        Arc<Mutex<HashMap<String, SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>>>>,
+    connections: Arc<Mutex<HashMap<String, WebsocketConnection>>>,
     read_tasks: Arc<Mutex<HashMap<String, tokio::task::JoinHandle<()>>>>,
 }
 
@@ -35,14 +42,20 @@ impl WebsocketManager {
         receive_tx: mpsc::Sender<Message>,
         validate_certificates: bool,
         client_cert: Option<ClientCertificateConfig>,
+        request_message_size: i32,
     ) -> Result<Response> {
         let tx = receive_tx.clone();
+        let max_message_size = message_size_limit(request_message_size);
 
         let (stream, response) =
-            ws_connect(url, headers, validate_certificates, client_cert).await?;
+            ws_connect(url, headers, validate_certificates, client_cert, request_message_size)
+                .await?;
         let (write, mut read) = stream.split();
 
-        self.connections.lock().await.insert(id.to_string(), write);
+        self.connections
+            .lock()
+            .await
+            .insert(id.to_string(), WebsocketConnection { max_message_size, sink: write });
 
         let handle = {
             let connection_id = id.to_string();
@@ -70,13 +83,20 @@ impl WebsocketManager {
     }
 
     pub async fn send(&mut self, id: &str, msg: Message) -> Result<()> {
-        debug!("Send websocket message {msg:?}");
         let mut connections = self.connections.lock().await;
         let connection = match connections.get_mut(id) {
             None => return Ok(()),
             Some(c) => c,
         };
-        connection.send(msg).await?;
+        if let Some(limit) = connection.max_message_size {
+            let message_size = msg.len();
+            if message_size > limit {
+                return Err(GenericError(format!(
+                    "WebSocket message too large: found {message_size} bytes, the limit is {limit} bytes"
+                )));
+            }
+        }
+        connection.sink.send(msg).await?;
         Ok(())
     }
 
@@ -84,7 +104,7 @@ impl WebsocketManager {
         info!("Closing websocket");
         if let Some(mut connection) = self.connections.lock().await.remove(id) {
             // Wait a maximum of 1 second for the connection to close
-            if let Err(e) = connection.close().await {
+            if let Err(e) = connection.sink.close().await {
                 warn!("Failed to close websocket connection {e:?}");
             };
         }