From e5ae112bbf66c81383904449f659ff7f52099f79 Mon Sep 17 00:00:00 2001 From: Jordan Mecom Date: Thu, 9 Jul 2026 13:14:46 -0400 Subject: [PATCH 1/2] Fix huddle audio WebSocket admission limits --- Cargo.lock | 1 + crates/buzz-relay/Cargo.toml | 1 + crates/buzz-relay/src/audio/handler.rs | 133 ++++++++++++++++++++++++- 3 files changed, 133 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index afb474a2dc..9abe60f57b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1074,6 +1074,7 @@ dependencies = [ "tempfile", "thiserror 2.0.18", "tokio", + "tokio-tungstenite 0.29.0", "tokio-util", "tower", "tower-http", diff --git a/crates/buzz-relay/Cargo.toml b/crates/buzz-relay/Cargo.toml index 1134b2c703..154f26636f 100644 --- a/crates/buzz-relay/Cargo.toml +++ b/crates/buzz-relay/Cargo.toml @@ -82,3 +82,4 @@ mesh-llm-host-runtime = { git = "https://github.com/Mesh-LLM/mesh-llm.git", rev buzz-core = { workspace = true, features = ["test-utils"] } buzz-auth = { workspace = true, features = ["dev"] } reqwest = { workspace = true } +tokio-tungstenite = { workspace = true } diff --git a/crates/buzz-relay/src/audio/handler.rs b/crates/buzz-relay/src/audio/handler.rs index 9e930ccc40..c394f2b51e 100644 --- a/crates/buzz-relay/src/audio/handler.rs +++ b/crates/buzz-relay/src/audio/handler.rs @@ -25,7 +25,7 @@ use bytes::Bytes; use futures_util::{SinkExt, StreamExt}; use nostr::{EventBuilder, Kind, Tag}; use serde::Deserialize; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, OwnedSemaphorePermit, Semaphore}; use tokio_util::sync::CancellationToken; use tracing::{debug, info, warn}; use uuid::Uuid; @@ -46,6 +46,11 @@ const MAX_AUDIO_FRAME_BYTES: usize = 4096; /// Maximum text frame size: 8 KB bounds auth/control JSON parsing. const MAX_TEXT_FRAME_BYTES: usize = 8192; +/// Parser-level cap for this route. Text auth/control frames are the largest +/// message type audio accepts; binary Opus frames are bounded more tightly +/// after parsing. +const MAX_WEBSOCKET_MESSAGE_BYTES: usize = MAX_TEXT_FRAME_BYTES; + /// Heartbeat interval. const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(30); @@ -81,7 +86,36 @@ pub async fn ws_audio_handler( .into_response(); } }; - ws.on_upgrade(move |socket| handle_audio_connection(socket, state, tenant, channel_id)) + + let permit = match acquire_audio_connection_permit(&state.conn_semaphore) { + Some(permit) => permit, + None => { + warn!(channel_id = %channel_id, "Connection limit reached, rejecting audio WebSocket"); + return ( + StatusCode::SERVICE_UNAVAILABLE, + "relay: connection limit reached", + ) + .into_response(); + } + }; + + // Keep the parser boundary at the largest message this route accepts. The + // checks in the receive loop still distinguish text from binary policy, but + // they run after tungstenite has assembled a message. + limit_audio_websocket(ws).on_upgrade(move |socket| { + handle_audio_connection(socket, state, tenant, channel_id, permit) + }) +} + +fn acquire_audio_connection_permit( + conn_semaphore: &Arc, +) -> Option { + Arc::clone(conn_semaphore).try_acquire_owned().ok() +} + +fn limit_audio_websocket(ws: WebSocketUpgrade) -> WebSocketUpgrade { + ws.max_message_size(MAX_WEBSOCKET_MESSAGE_BYTES) + .max_frame_size(MAX_WEBSOCKET_MESSAGE_BYTES) } /// Highest huddle audio protocol version this relay understands. Clients are @@ -112,6 +146,7 @@ async fn handle_audio_connection( state: Arc, tenant: TenantContext, channel_id: Uuid, + _permit: OwnedSemaphorePermit, ) { let (mut ws_send, mut ws_recv) = socket.split(); @@ -880,3 +915,97 @@ async fn emit_participant_event( ); } } + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use axum::{routing::get, Router}; + use futures_util::SinkExt; + use tokio::net::TcpListener; + use tokio::sync::oneshot; + use tokio_tungstenite::{connect_async, tungstenite::Message}; + + use super::*; + + #[test] + fn audio_connection_permits_share_the_global_websocket_budget() { + let semaphore = Arc::new(Semaphore::new(1)); + let first = acquire_audio_connection_permit(&semaphore).expect("first permit"); + + assert!( + acquire_audio_connection_permit(&semaphore).is_none(), + "audio connections must stop when the global WebSocket budget is exhausted" + ); + + drop(first); + assert!( + acquire_audio_connection_permit(&semaphore).is_some(), + "dropping an audio connection must return its global permit" + ); + } + + async fn handler_receives_message_of_size(size: usize) -> bool { + let (received_tx, received_rx) = oneshot::channel(); + let received_tx = Arc::new(Mutex::new(Some(received_tx))); + let app = Router::new().route( + "/", + get({ + let received_tx = Arc::clone(&received_tx); + move |ws: WebSocketUpgrade| { + let received_tx = Arc::clone(&received_tx); + async move { + limit_audio_websocket(ws).on_upgrade(move |mut socket| async move { + let received = matches!(socket.recv().await, Some(Ok(_))); + if let Some(tx) = + received_tx.lock().expect("result lock poisoned").take() + { + let _ = tx.send(received); + } + }) + } + } + }), + ); + + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind test WebSocket listener"); + let addr = listener.local_addr().expect("test listener address"); + let server = tokio::spawn(async move { + axum::serve(listener, app) + .await + .expect("test WebSocket server"); + }); + + let (mut client, _) = connect_async(format!("ws://{addr}/")) + .await + .expect("connect test WebSocket client"); + client + .send(Message::Text("x".repeat(size).into())) + .await + .expect("send test WebSocket message"); + + let received = tokio::time::timeout(Duration::from_secs(2), received_rx) + .await + .expect("server should process the test message") + .expect("server should report whether it received the message"); + + server.abort(); + let _ = server.await; + + received + } + + #[tokio::test] + async fn audio_websocket_parser_rejects_oversized_messages_before_handler_reads_them() { + assert!( + handler_receives_message_of_size(MAX_WEBSOCKET_MESSAGE_BYTES).await, + "messages at the audio route limit should still reach the handler" + ); + assert!( + !handler_receives_message_of_size(MAX_WEBSOCKET_MESSAGE_BYTES + 1).await, + "oversized messages must be rejected by the WebSocket parser before the handler sees them" + ); + } +} From bf8ebb0c1fe7b7999950e2b63e929f4f5065bc2c Mon Sep 17 00:00:00 2001 From: Jordan Mecom Date: Thu, 9 Jul 2026 14:10:04 -0400 Subject: [PATCH 2/2] Cap main relay WebSocket parser messages --- crates/buzz-relay/src/router.rs | 81 ++++++++++++++++++++++++++++++++- 1 file changed, 80 insertions(+), 1 deletion(-) diff --git a/crates/buzz-relay/src/router.rs b/crates/buzz-relay/src/router.rs index 69e0368ccd..65a7a94dd4 100644 --- a/crates/buzz-relay/src/router.rs +++ b/crates/buzz-relay/src/router.rs @@ -195,8 +195,9 @@ async fn nip11_or_ws_handler( } }; + let max_frame_bytes = state.config.max_frame_bytes; match WebSocketUpgrade::from_request(req, &state).await { - Ok(ws) => ws + Ok(ws) => limit_relay_websocket(ws, max_frame_bytes) .on_upgrade(move |socket| handle_connection(socket, state, addr, tenant)) .into_response(), Err(_) => { @@ -215,6 +216,16 @@ async fn nip11_or_ws_handler( } } +fn limit_relay_websocket( + ws: WebSocketUpgrade, + max_frame_bytes: usize, +) -> WebSocketUpgrade { + // recv_loop keeps the application-level check as defense in depth, but + // parser limits must be set before tungstenite assembles the message. + ws.max_message_size(max_frame_bytes) + .max_frame_size(max_frame_bytes) +} + async fn health_handler() -> impl IntoResponse { (StatusCode::OK, "ok") } @@ -292,3 +303,71 @@ fn build_cors_layer(cors_origins: &[String]) -> CorsLayer { .allow_methods(tower_http::cors::Any) .allow_headers(tower_http::cors::Any) } + +#[cfg(test)] +mod tests { + use axum::{routing::get, Router}; + use futures_util::SinkExt; + use tokio::net::TcpListener; + use tokio::sync::mpsc; + use tokio_tungstenite::{connect_async, tungstenite::Message}; + + use super::*; + + async fn handler_receives_message_with_limit(limit: usize, size: usize) -> bool { + let (received_tx, mut received_rx) = mpsc::unbounded_channel(); + let app = Router::new().route( + "/", + get(move |ws: WebSocketUpgrade| { + let received_tx = received_tx.clone(); + async move { + limit_relay_websocket(ws, limit).on_upgrade(move |mut socket| async move { + let _ = received_tx.send(matches!(socket.recv().await, Some(Ok(_)))); + }) + } + }), + ); + + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind test WebSocket listener"); + let addr = listener.local_addr().expect("test listener address"); + let server = tokio::spawn(async move { + axum::serve(listener, app) + .await + .expect("test WebSocket server"); + }); + + let (mut client, _) = connect_async(format!("ws://{addr}/")) + .await + .expect("connect test WebSocket client"); + client + .send(Message::Text("x".repeat(size).into())) + .await + .expect("send test WebSocket message"); + + let received = tokio::time::timeout(std::time::Duration::from_secs(2), received_rx.recv()) + .await + .expect("server should process the test message") + .expect("server should report whether it received the message"); + + server.abort(); + let _ = server.await; + + received + } + + #[tokio::test] + async fn relay_websocket_parser_rejects_oversized_messages_before_handler_reads_them() { + let limit = 64; + + assert!( + handler_receives_message_with_limit(limit, limit).await, + "messages at the relay limit should still reach the handler" + ); + assert!( + !handler_receives_message_with_limit(limit, limit + 1).await, + "oversized messages must be rejected by the WebSocket parser before the handler sees them" + ); + } +}