diff --git a/Cargo.toml b/Cargo.toml index 81bdc6dccb..b0c2e5fece 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -57,7 +57,7 @@ actix-web-rust-embed-responder = { version = "2.2", default-features = false, fe "support-rust-embed-for-web", "base64", ] } -actix-ws = "0.2.5" +actix-ws = "0.3.0" actix = "0.13.5" actix-web-actors = "4.3.0" actix-tls.workspace = true diff --git a/src/config/src/config.rs b/src/config/src/config.rs index b538ba165a..0a744471a6 100644 --- a/src/config/src/config.rs +++ b/src/config/src/config.rs @@ -431,6 +431,18 @@ pub struct WebSocket { pub session_gc_interval_secs: i64, #[env_config(name = "ZO_WEBSOCKET_PING_INTERVAL_SECS", default = 30)] pub ping_interval_secs: i64, + #[env_config( + name = "ZO_WEBSOCKET_MAX_FRAME_SIZE", + default = 1, + help = "Maximum allowed frame size in MB" + )] + pub max_frame_size: usize, + #[env_config( + name = "ZO_WEBSOCKET_MAX_CONTINUATION_SIZE", + default = 256, + help = "Maximum allowed continuation size in MB" + )] + pub max_continuation_size: usize, } #[derive(EnvConfig)] diff --git a/src/handler/http/request/ws_v2/mod.rs b/src/handler/http/request/ws_v2/mod.rs index 834a7f8a41..fb586225dd 100644 --- a/src/handler/http/request/ws_v2/mod.rs +++ b/src/handler/http/request/ws_v2/mod.rs @@ -70,6 +70,13 @@ pub async fn websocket( let (res, session, msg_stream) = actix_ws::handle(&req, stream)?; + // increase the maximum allowed frame size to 1MB actix-ws and + // aggregate continuation frames + let msg_stream = msg_stream + .max_frame_size(cfg.websocket.max_frame_size * 1024 * 1024) + .aggregate_continuations() + .max_continuation_size(cfg.websocket.max_continuation_size * 1024 * 1024); + let ws_session = WsSession::new(session, cookie_expiry); sessions_cache_utils::insert_session(&router_id, Arc::new(RwLock::new(ws_session))).await; log::info!( diff --git a/src/handler/http/request/ws_v2/session.rs b/src/handler/http/request/ws_v2/session.rs index 3e21e3c44b..9e09660fcb 100644 --- a/src/handler/http/request/ws_v2/session.rs +++ b/src/handler/http/request/ws_v2/session.rs @@ -16,7 +16,7 @@ use std::time::Duration; use actix_http::ws::{CloseCode, CloseReason}; -use actix_ws::{MessageStream, Session}; +use actix_ws::{AggregatedMessageStream, Session}; use config::{ get_config, meta::websocket::{SearchEventReq, SearchResultType}, @@ -133,7 +133,12 @@ impl WsSession { } } -pub async fn run(mut msg_stream: MessageStream, user_id: String, req_id: String, path: String) { +pub async fn run( + mut msg_stream: AggregatedMessageStream, + user_id: String, + req_id: String, + path: String, +) { log::info!( "[WS_HANDLER]: Starting WebSocket session for req_id: {}, path: {}", req_id, @@ -161,7 +166,7 @@ pub async fn run(mut msg_stream: MessageStream, user_id: String, req_id: String, } match msg { - Ok(actix_ws::Message::Ping(bytes)) => { + Ok(actix_ws::AggregatedMessage::Ping(bytes)) => { log::debug!( "[WS_HANDLER]: Received ping from client for req_id: {}", req_id @@ -183,13 +188,13 @@ pub async fn run(mut msg_stream: MessageStream, user_id: String, req_id: String, break; } } - Ok(actix_ws::Message::Pong(_)) => { + Ok(actix_ws::AggregatedMessage::Pong(_)) => { log::debug!( "[WS_HANDLER]: Received pong from client for req_id: {}", req_id ); } - Ok(actix_ws::Message::Text(msg)) => { + Ok(actix_ws::AggregatedMessage::Text(msg)) => { log::info!("[WS_HANDLER]: Request Id: {} Node Role: {} Received message: {}", req_id, get_config().common.node_role, @@ -197,7 +202,7 @@ pub async fn run(mut msg_stream: MessageStream, user_id: String, req_id: String, ); handle_text_message(&user_id, &req_id, msg.to_string(), path.clone()).await; } - Ok(actix_ws::Message::Close(reason)) => { + Ok(actix_ws::AggregatedMessage::Close(reason)) => { if let Some(reason) = reason.as_ref() { match reason.code { CloseCode::Normal | CloseCode::Error => { diff --git a/src/router/http/ws_v2/handler.rs b/src/router/http/ws_v2/handler.rs index 80c51401f3..c7bd82d0aa 100644 --- a/src/router/http/ws_v2/handler.rs +++ b/src/router/http/ws_v2/handler.rs @@ -77,7 +77,13 @@ impl WsHandler { log::info!("[WS::Router::Handler] new client incoming: {}", client_id); // Client -> Router connection - let (response, mut ws_session, mut msg_stream) = actix_ws::handle(&req, stream)?; + let (response, mut ws_session, msg_stream) = actix_ws::handle(&req, stream)?; + // increase the maximum allowed frame size to 128KiB and aggregate continuation frames + let mut msg_stream = msg_stream + .max_frame_size(cfg.websocket.max_frame_size * 1024 * 1024) + .aggregate_continuations() + .max_continuation_size(cfg.websocket.max_continuation_size * 1024 * 1024); + log::info!("[WS::Router::Handler] new client connected: {}", client_id); // Create session by registering the client & extract the user_id from the auth @@ -156,7 +162,7 @@ impl WsHandler { break; } } - Ok(actix_ws::Message::Text(text)) => { + Ok(actix_ws::AggregatedMessage::Text(text)) => { match json::from_str::(&text) { Err(e) => { log::error!( @@ -275,7 +281,7 @@ impl WsHandler { } } } - Ok(actix_ws::Message::Close(close_reason)) => { + Ok(actix_ws::AggregatedMessage::Close(close_reason)) => { log::info!( "[WS::Router::Handler] disconnect signal received from client_id: {}, close_reason: {:?}", client_id, close_reason ); @@ -287,10 +293,10 @@ impl WsHandler { log::debug!("[WS::Router::Handler] Stop handle_incoming for client_id: {}", client_id); break; } - Ok(actix_ws::Message::Ping(ping)) => { + Ok(actix_ws::AggregatedMessage::Ping(ping)) => { let _ = response_tx.send(WsServerEvents::Ping(ping.to_vec())).await; } - Ok(actix_ws::Message::Pong(pong)) => { + Ok(actix_ws::AggregatedMessage::Pong(pong)) => { log::debug!("[WS::Router::Handler] Pong received from client_id: {}, pong: {:?}", client_id, pong); } Ok(_) => {}