[codex] Fix relay WebSocket admission limits

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 }

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<Semaphore>,
+) -> Option<OwnedSemaphorePermit> {
+    Arc::clone(conn_semaphore).try_acquire_owned().ok()
+}
+
+fn limit_audio_websocket<F>(ws: WebSocketUpgrade<F>) -> WebSocketUpgrade<F> {
+    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<AppState>,
     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"
+        );
+    }
+}

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<F>(
+    ws: WebSocketUpgrade<F>,
+    max_frame_bytes: usize,
+) -> WebSocketUpgrade<F> {
+    // 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"
+        );
+    }
+}