From 91e7b469b82b235c62885f82a6323b7e666e9254 Mon Sep 17 00:00:00 2001 From: Will Pfleger Date: Mon, 19 Oct 2026 16:13:27 -0400 Subject: [PATCH] fix(relay): wait for the peer's Close before dropping any WebSocket (#8426) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 🤖 A NIP-42 login refused for a NIP-FI key mismatch could reach the client as an abnormal close (`1006`, no reason) instead of the `NOTICE "restricted: authorization denied"` followed by close `1008` that every other authorization denial produces. The relay did write the NOTICE and the `1008` Close, but the connection task dropped TCP as soon as the writer finished, before the peer finished the close handshake. The client could then lose frames it hadn't read yet. It depends on timing, so it showed up intermittently and is not specific to key pairing. ## Relay route - `recv_loop` now hands back the read half even while a frame handler is still running. Frames are handled one at a time while the connection is live. Once it is cancelled, an in-flight handler (for example a Nostr `CLOSE` waiting on the subscription lock behind a REQ stuck in Redis) is spawned to finish instead of holding the reader. It is never dropped partway, so subscription retirement stays atomic, and cleanup joins it before the expiry barrier and application cleanup. - As soon as the writer finishes, `await_client_close` keeps the transport open until the client's Close arrives or the stream ends. It runs concurrently with joining the heartbeat, auth-timeout, expiry and auth-cancel tasks and the connection's application cleanup (subscriptions released, registry entry removed, presence cleared). Cleanup never waits on the peer, so a nonresponsive client cannot delay registry counts, presence or `has_admitted_connection`. The socket never waits on cleanup, so a stalled Redis (the pool has no wait timeout) or an expiry join blocked on its quiescence barrier cannot hold the socket and its file descriptor open past the bounded wait. The active-connection gauge and the connection permit are released after both finish. - The connection permit (`BUZZ_MAX_CONNECTIONS`) is now taken at the start of `handle_connection`, before the NIP-FI gate and the community-active bootstrap check, and held through the active handler or the refusal drain and its bounded wait. Previously a socket in bootstrap, or refused there, held no permit, so a burst of refusals could open sockets beyond the limit. Over the limit, the socket is still dropped with no frames. `buzz_ws_connections_active` still counts only connections that reached the active handler. ## Huddle-audio route - The relay and audio bootstrap drains share one bounded helper. A socket refused before admission waits after its community registry entry is released. - Every exit of the audio handler gets the bounded peer wait. Exits whose cleanup follows their terminal frames (`room_full`, `room_ended`, version mismatch, archived channel, owner rejection or unreachable, failed commit confirmation) and the normal teardown hand the socket to a background task that waits for the client's Close, so lease releases, `48102` writes and room archiving on Postgres and Redis never hold the socket open. The remaining exits send their frames last, and `handle_active_audio_connection` waits once they return. The wrapper joins the background task before returning, so the audio connection permit is held until the socket ends, as on the relay route. - Audio error exits now follow their existing JSON error with a Close frame: | Exit | Close code | | --- | --- | | `huddle_relay_draining` | `1012` service restart | | `join_rejected` (owner resolution failed), `huddle_owner_unreachable`, failed channel lookup | `1011` internal error | | `huddle_audio_unavailable`, `unsupported_version`, `upgrade_required`, `room_full`, `room_ended`, archived channel, owner rejection, NIP-FI authorization denial | `1008` policy | Audio login failures and authorization refusals outside NIP-FI `authorization_denied` still send only their existing error frame, then get the bounded wait. - The audio read half is an enum, so the hand-off has no `expect()` that can panic. ## Timing - The peer wait is bounded by `WS_TERMINAL_FLUSH_TIMEOUT` (1s), and drains to end-of-stream so a client-initiated Close gets its reply. The terminal writer flush has its own 1s budget, so the worst case for terminal transport work is about 2s per refused or closed socket. - Terminal frame writes (`flush_terminal_frames`, `send_exit_frames_bounded`) still give up on the first failure, but a failed or timed-out send is now logged at `debug` instead of being discarded. - The graceful-shutdown comment in `main.rs` now describes per-stage budgets (restart jitter, restart-Close acknowledgement, terminal flush, peer wait) rather than a per-socket total. The 30s hard drain bounds every socket, and the 35s process shutdown budget fits within the chart's 60s `terminationGracePeriodSeconds`. ## Tests Real-socket tests in the infra-free relay lane of `just test-unit` and `scripts/run-tests.sh` unless marked PostgreSQL: - The key-mismatch login and a session-expiry denial (as a control) both assert NOTICE then `1008` / `authorization denied` on the wire. - A raw-TCP client that never answers Close: the relay holds the socket for at least the bounded wait, has deregistered the connection by the time the socket ends, and ends it on its own under a 10s watchdog. - A raw-TCP client that withholds its Close until the connection has deregistered: the relay ends the socket in under the 1s cap, which fails if cleanup waits on the peer. - A raw-TCP client that never answers while final cleanup is stalled at a test pause point: the socket still ends after the bounded wait. - A raw-TCP client that never answers Close on the relay bootstrap drain, the audio bootstrap drain and an early audio-handler exit: the relay holds each socket for the bounded wait. - On the relay route with a one-slot budget through the real `handle_connection`: a refused bootstrap holds the slot until its socket ends, and a bootstrap stalled on the community check already holds the slot before the AUTH challenge is sent. - An established relay session whose reader is blocked in a Nostr `CLOSE` handler on a held subscription lock when the session expires: the socket ends within the bounded wait while the lock is held, and cleanup completes with the subscription retired once it is released. - An established relay session whose expiry task is stalled behind a held session effect permit: the socket still ends within the bounded wait. - PostgreSQL: an established audio session whose teardown is stalled after its Close, and a stalled archived-channel exit, both end within the bounded wait from the hand-off. - PostgreSQL: a fast refusal (unsupported protocol version) through the real `handle_audio_connection` with a one-slot permit budget keeps the permit taken until that socket ends. - PostgreSQL: `huddle_audio_unavailable` sends its error and then Close `1008` as the very next frame. --------- Signed-off-by: Will Pfleger Signed-off-by: Duncan Co-authored-by: Duncan --- Cargo.lock | 1 + Justfile | 11 + crates/buzz-relay/Cargo.toml | 3 + crates/buzz-relay/src/audio/handler.rs | 768 ++++++++++++++-- crates/buzz-relay/src/connection.rs | 967 +++++++++++++++++++-- crates/buzz-relay/src/handlers/close.rs | 6 +- crates/buzz-relay/src/main.rs | 38 +- crates/buzz-relay/src/nip_fi_test_hooks.rs | 13 + scripts/run-tests.sh | 11 + 9 files changed, 1645 insertions(+), 173 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d266b8580a..38ad23b30b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1489,6 +1489,7 @@ dependencies = [ "tracing", "tracing-opentelemetry", "tracing-subscriber", + "tungstenite 0.29.0", "url", "uuid", ] diff --git a/Justfile b/Justfile index e479c467f4..59e98d9c11 100644 --- a/Justfile +++ b/Justfile @@ -577,6 +577,17 @@ test-unit: + test(=connection::tests::expiry_notice_queued_on_ctrl_before_cancel) + test(=connection::tests::f3_root_outer_wrapper_delivers_denial_on_bootstrap_cancellation) + test(=connection::tests::f3_root_pre_built_expired_gate_terminates_connection) + + test(=connection::terminal_close_delivery_tests::key_pairing_mismatch_closes_with_1008_on_real_socket) + + test(=connection::terminal_close_delivery_tests::session_expiry_denial_closes_with_1008_on_real_socket) + + test(=connection::terminal_close_delivery_tests::server_holds_socket_for_client_close_reply) + + test(=connection::terminal_close_delivery_tests::server_finishes_promptly_when_client_answers_close) + + test(=connection::terminal_close_delivery_tests::stalled_cleanup_does_not_hold_socket_open) + + test(=connection::terminal_close_delivery_tests::bootstrap_drain_holds_socket_for_client_close_reply) + + test(=connection::terminal_close_delivery_tests::stalled_expiry_barrier_does_not_hold_socket_open) + + test(=connection::terminal_close_delivery_tests::blocked_frame_handler_does_not_hold_socket_open) + + test(=connection::terminal_close_delivery_tests::refused_bootstrap_holds_connection_slot_until_socket_ends) + + test(=connection::terminal_close_delivery_tests::stalled_bootstrap_holds_connection_slot) + + test(=connection::terminal_close_delivery_tests::peer_wait_replies_to_client_initiated_close) + test(=handlers::auth::tests::b2_pre_cancelled_connection_never_becomes_authenticated) + test(=handlers::auth::tests::fi_ban_check_error_emits_terminal_authorization_unavailable) + test(=handlers::auth::tests::fi_root_authorization_denied_rows_emit_identical_frames) diff --git a/crates/buzz-relay/Cargo.toml b/crates/buzz-relay/Cargo.toml index 6a4ab379fd..ef0808ef65 100644 --- a/crates/buzz-relay/Cargo.toml +++ b/crates/buzz-relay/Cargo.toml @@ -30,6 +30,9 @@ buzz-relay-mesh = { workspace = true } async-trait = "0.1" postcard = { workspace = true } axum = { workspace = true } +# Axum's WebSocket errors wrap tungstenite's; matching them tells a close +# handshake already under way apart from a failed terminal-frame write. +tungstenite = { version = "0.29", default-features = false } tokio = { workspace = true } # `io` (ReaderStream) is needed to stream git subprocess stdout straight into # the HTTP response body — see api::git::transport read-path streaming. The diff --git a/crates/buzz-relay/src/audio/handler.rs b/crates/buzz-relay/src/audio/handler.rs index 1a8b251ea2..3a4d2fcde1 100644 --- a/crates/buzz-relay/src/audio/handler.rs +++ b/crates/buzz-relay/src/audio/handler.rs @@ -294,6 +294,11 @@ pub(crate) async fn handle_audio_connection( let rx_for_run = Arc::clone(&rx_shared); let rx_for_drain = Arc::clone(&rx_shared); + // The drain path parks the read half here so the bounded wait for the + // client's Close runs after the registry entry is released. + let unread_shared = Arc::new(tokio::sync::Mutex::new(None)); + let unread_for_drain = Arc::clone(&unread_shared); + run_registered_community_connection( ®istry, Uuid::new_v4(), @@ -330,37 +335,21 @@ pub(crate) async fn handle_audio_connection( }, move || async move { let socket = socket_for_drain.lock().await.take(); - let mut pre_terminal_ctrl_rx = rx_for_drain.lock().await.take(); + let queued = rx_for_drain.lock().await.take(); if let Some(socket) = socket { - let (mut ws_send, _ws_recv) = socket.split(); - if let Some(ref mut rx) = pre_terminal_ctrl_rx { - while let Ok(msg) = rx.try_recv() { - let _ = tokio::time::timeout( - crate::connection::WS_TERMINAL_FLUSH_TIMEOUT, - futures_util::SinkExt::send(&mut ws_send, msg), - ) - .await; - } - } // A published reason (e.g. an expiry denial) closes with its // own frame, exactly as the writer would. let close = drain_reason.borrow().map(|reason| reason.close_message()); - if let Some(close) = close { - let _ = tokio::time::timeout( - crate::connection::WS_TERMINAL_FLUSH_TIMEOUT, - futures_util::SinkExt::send(&mut ws_send, close), - ) - .await; - } - let _ = tokio::time::timeout( - crate::connection::WS_TERMINAL_FLUSH_TIMEOUT, - futures_util::SinkExt::close(&mut ws_send), - ) - .await; + let ws_recv = crate::connection::drain_refused_socket(socket, queued, close).await; + *unread_for_drain.lock().await = Some(ws_recv); } }, ) .await; + let unread = unread_shared.lock().await.take(); + if let Some(ws_recv) = unread { + crate::connection::await_client_close(ws_recv).await; + } } /// Mid-admission cancellation fence for the audio route: if `$cancel` fired, @@ -411,6 +400,7 @@ async fn send_challenge_unless_cancelled( ) -> bool where S: futures_util::Sink + Unpin, + S::Error: std::error::Error + 'static, { tokio::select! { biased; @@ -438,6 +428,10 @@ pub(crate) fn audio_post_auth_register( control.set_proven_identity(pubkey_bytes, nip_fi_issuer); } +/// Runs the audio session, then holds the read half for the client's Close +/// reply, bounded by `WS_TERMINAL_FLUSH_TIMEOUT`. An exit whose cleanup +/// follows its terminal frames hands the socket off first through +/// [`hand_off_socket`], so that cleanup never holds the socket open. #[allow(clippy::too_many_arguments)] pub(crate) async fn handle_active_audio_connection( socket: WebSocket, @@ -451,13 +445,102 @@ pub(crate) async fn handle_active_audio_connection( // which arms them BEFORE the `is_community_active` bootstrap await. // `None` is used by test call sites that bypass the outer wrapper. pre_built: Option, +) { + let (ws_send, ws_recv) = socket.split(); + let mut ws_recv = AudioReadHalf::Held(ws_recv); + run_active_audio_connection( + ws_send, + &mut ws_recv, + state, + tenant, + channel_id, + control, + nip_fi_assertion, + connection_time, + pre_built, + ) + .await; + // Returning releases the caller's connection permit, so the socket must + // have ended first: wait on the read half here, or on the background + // close an exit handed it to. + match ws_recv { + AudioReadHalf::Held(ws_recv) => crate::connection::await_client_close(ws_recv).await, + AudioReadHalf::HandedOff(close) => { + if let Err(error) = close.await { + warn!(channel_id = %channel_id, "audio background close failed: {error}"); + } + } + } +} + +/// The audio session's read half, or the background close it was handed to. +enum AudioReadHalf { + Held(futures_util::stream::SplitStream), + HandedOff(tokio::task::JoinHandle<()>), +} + +/// Hands the socket to a task that waits, bounded, for the client's Close +/// reply and then drops it. Called once the terminal frames are sent and +/// before cleanup that may block on Postgres, Redis or a remote owner, so +/// that cleanup never holds the socket open. Pass the write half when the +/// caller still owns it: either half alone keeps the socket open. The task +/// handle is kept so the wrapper holds the connection permit until the +/// socket ends. +fn hand_off_socket( + ws_recv: &mut AudioReadHalf, + ws_send: Option>, +) { + if let AudioReadHalf::HandedOff(_) = ws_recv { + return; + } + // The task is spawned before the read half moves out of `ws_recv`, so + // `ws_recv` is never left without one of its two states. + let (recv_tx, recv_rx) = tokio::sync::oneshot::channel(); + let close = tokio::spawn(async move { + if let Ok(recv) = recv_rx.await { + crate::connection::await_client_close(recv).await; + } + drop(ws_send); + }); + if let AudioReadHalf::Held(recv) = std::mem::replace(ws_recv, AudioReadHalf::HandedOff(close)) { + let _ = recv_tx.send(recv); + } +} + +/// An admission error frame followed by the Close that ends the socket, so +/// the client gets a clean close handshake instead of waiting out the +/// bounded peer wait for a Close that never comes. `code` is +/// `close_code::POLICY` (1008) for a refusal, `close_code::ERROR` (1011) +/// for a server-side failure and `close_code::RESTART` (1012) when the relay +/// is draining and the client should reconnect. +fn error_exit_frames(error: serde_json::Value, code: u16) -> [WsMessage; 2] { + [WsMessage::Text(error.to_string().into()), exit_close(code)] +} + +fn exit_close(code: u16) -> WsMessage { + WsMessage::Close(Some(axum::extract::ws::CloseFrame { + code, + reason: axum::extract::ws::Utf8Bytes::default(), + })) +} + +#[allow(clippy::too_many_arguments)] +async fn run_active_audio_connection( + mut ws_send: futures_util::stream::SplitSink, + ws_recv: &mut AudioReadHalf, + state: Arc, + tenant: TenantContext, + channel_id: Uuid, + control: CommunityConnectionControl, + nip_fi_assertion: Option, + connection_time: chrono::DateTime, + pre_built: Option, ) { let cancel = control.cancellation_token(); let disconnect_reason = control.disconnect_reason(); // connection_time is threaded in from the HTTP handler (captured immediately // before on_upgrade) so the session partition is rooted at the true upgrade // instant, not the post-community-active-check instant. [FI-TRACE-LEASE-BOUND] - let (mut ws_send, mut ws_recv) = socket.split(); // P2 / Fix 3: Arm the NIP-FI session gate and expiry task. // @@ -550,6 +633,9 @@ pub(crate) async fn handle_active_audio_connection( return; }, result = tokio::time::timeout(AUTH_TIMEOUT, async { + let AudioReadHalf::Held(ws_recv) = ws_recv else { + return None; + }; while let Some(Ok(msg)) = ws_recv.next().await { if let WsMessage::Text(text) = msg { if text.len() > MAX_TEXT_FRAME_BYTES { @@ -569,8 +655,17 @@ pub(crate) async fn handle_active_audio_connection( let auth_msg = match auth_result { Ok(Some(a)) => a, - _ => { - debug!(channel_id = %channel_id, "audio auth timeout or disconnect"); + Ok(None) => { + debug!(channel_id = %channel_id, "audio client disconnected before auth"); + return; + } + Err(_) => { + debug!(channel_id = %channel_id, "audio auth timeout"); + crate::connection::send_exit_frames_bounded( + &mut ws_send, + [exit_close(axum::extract::ws::close_code::POLICY)], + ) + .await; return; } }; @@ -932,15 +1027,14 @@ pub(crate) async fn handle_active_audio_connection( shadow_attempt.refused(); crate::connection::send_exit_frames_bounded( &mut ws_send, - [WsMessage::Text( + error_exit_frames( serde_json::json!({ "type": "error", "code": "huddle_relay_draining", "message": "relay is draining; reconnect" - }) - .to_string() - .into(), - )], + }), + axum::extract::ws::close_code::RESTART, + ), ) .await; return; @@ -993,15 +1087,14 @@ pub(crate) async fn handle_active_audio_connection( ); crate::connection::send_exit_frames_bounded( &mut ws_send, - [WsMessage::Text( + error_exit_frames( serde_json::json!({ "type": "error", "code": "join_rejected", "message": "huddle join rejected" - }) - .to_string() - .into(), - )], + }), + axum::extract::ws::close_code::ERROR, + ), ) .await; return; @@ -1021,15 +1114,14 @@ pub(crate) async fn handle_active_audio_connection( ); crate::connection::send_exit_frames_bounded( &mut ws_send, - [WsMessage::Text( + error_exit_frames( serde_json::json!({ "type": "error", "code": "huddle_audio_unavailable", "message": "huddle audio unavailable in this deployment" - }) - .to_string() - .into(), - )], + }), + axum::extract::ws::close_code::POLICY, + ), ) .await; return; @@ -1052,19 +1144,23 @@ pub(crate) async fn handle_active_audio_connection( // get_or_create (the old room was already cleaned up). This DB check // catches that case. The room-level ended flag (checked inside add_peer) // handles the same-room case. + #[cfg(test)] + crate::nip_fi_test_hooks::before_room_archive_recheck(tenant.community()).await; match state.db.get_channel(tenant.community(), channel_id).await { Ok(ch) if ch.archived_at.is_some() => { shadow_attempt.refused(); debug!(channel_id = %channel_id, "channel archived before room join"); crate::connection::send_exit_frames_bounded( &mut ws_send, - [WsMessage::Text( - serde_json::json!({"type":"error","message":"huddle has ended"}) - .to_string() - .into(), - )], + error_exit_frames( + serde_json::json!({"type":"error","message":"huddle has ended"}), + axum::extract::ws::close_code::POLICY, + ), ) .await; + hand_off_socket(ws_recv, Some(ws_send)); + #[cfg(test)] + crate::nip_fi_test_hooks::before_archived_exit_cleanup(tenant.community()).await; // I1 residual: release lease with an awaited call, not a detached task. if let Some((lease, directory)) = staged_lease { if let Err(e) = directory.release(&lease).await { @@ -1079,6 +1175,12 @@ pub(crate) async fn handle_active_audio_connection( Err(e) => { shadow_attempt.refused(); warn!(channel_id = %channel_id, "pre-join channel check failed (fail-closed): {e}"); + crate::connection::send_exit_frames_bounded( + &mut ws_send, + [exit_close(axum::extract::ws::close_code::ERROR)], + ) + .await; + hand_off_socket(ws_recv, Some(ws_send)); // I1 residual: release lease with an awaited call, not a detached task. if let Some((lease, directory)) = staged_lease { if let Err(re) = directory.release(&lease).await { @@ -1108,7 +1210,7 @@ pub(crate) async fn handle_active_audio_connection( current = CURRENT_PROTOCOL_VERSION, "audio: client requested unsupported protocol version" ); - crate::connection::send_exit_frames_bounded(&mut ws_send, [WsMessage::Text( + crate::connection::send_exit_frames_bounded(&mut ws_send, error_exit_frames( serde_json::json!({ "type": "error", "code": "unsupported_version", @@ -1116,11 +1218,11 @@ pub(crate) async fn handle_active_audio_connection( "huddle audio protocol v{requested_version} not supported; relay max is v{CURRENT_PROTOCOL_VERSION}" ), "current_version": CURRENT_PROTOCOL_VERSION, - }) - .to_string() - .into(), - )]) + }), + axum::extract::ws::close_code::POLICY, + )) .await; + hand_off_socket(ws_recv, Some(ws_send)); if let Some((lease, directory)) = staged_lease { // I1 residual: release lease with an awaited call, not a detached task. if let Err(e) = directory.release(&lease).await { @@ -1182,11 +1284,13 @@ pub(crate) async fn handle_active_audio_connection( warn!(channel_id = %channel_id, pubkey = %pubkey_hex, "huddle owner rejected registration: {reason:?}"); crate::connection::send_exit_frames_bounded( &mut ws_send, - [WsMessage::Text( - remote_rejection_ws_error(&reason).to_string().into(), - )], + error_exit_frames( + remote_rejection_ws_error(&reason), + axum::extract::ws::close_code::POLICY, + ), ) .await; + hand_off_socket(ws_recv, Some(ws_send)); // I3 residual: await expiry task before resource teardown. cancel.cancel(); if let Some(t) = _nip_fi_admission_expiry.take() { @@ -1203,16 +1307,16 @@ pub(crate) async fn handle_active_audio_connection( warn!(channel_id = %channel_id, pubkey = %pubkey_hex, "huddle owner registration failed: {e}"); crate::connection::send_exit_frames_bounded( &mut ws_send, - [WsMessage::Text( + error_exit_frames( serde_json::json!({ "type": "error", "code": "huddle_owner_unreachable", "message": "could not reach the huddle owner" - }) - .to_string() - .into(), - )], + }), + axum::extract::ws::close_code::ERROR, + ), ) .await; + hand_off_socket(ws_recv, Some(ws_send)); // I3 residual: await expiry task before resource teardown. cancel.cancel(); if let Some(t) = _nip_fi_admission_expiry.take() { @@ -1296,7 +1400,8 @@ pub(crate) async fn handle_active_audio_connection( Err(crate::audio::room::AdmissionError::Full) => { shadow_attempt.refused(); warn!(channel_id = %channel_id, "audio room participant capacity reached"); - crate::connection::send_exit_frames_bounded(&mut ws_send, [WsMessage::Text(serde_json::json!({"type":"error","code":"room_full","message":"room participant capacity reached"}).to_string().into())]).await; + crate::connection::send_exit_frames_bounded(&mut ws_send, error_exit_frames(serde_json::json!({"type":"error","code":"room_full","message":"room participant capacity reached"}), axum::extract::ws::close_code::POLICY)).await; + hand_off_socket(ws_recv, Some(ws_send)); // IMPORTANT 3: cancel + await expiry task before guard release. cancel.cancel(); if let Some(t) = _nip_fi_admission_expiry.take() { @@ -1308,7 +1413,8 @@ pub(crate) async fn handle_active_audio_connection( Err(crate::audio::room::AdmissionError::Ended) => { shadow_attempt.refused(); debug!(channel_id = %channel_id, "room ended before admission"); - crate::connection::send_exit_frames_bounded(&mut ws_send, [WsMessage::Text(serde_json::json!({"type":"error","code":"room_ended","message":"huddle has ended"}).to_string().into())]).await; + crate::connection::send_exit_frames_bounded(&mut ws_send, error_exit_frames(serde_json::json!({"type":"error","code":"room_ended","message":"huddle has ended"}), axum::extract::ws::close_code::POLICY)).await; + hand_off_socket(ws_recv, Some(ws_send)); // IMPORTANT 3: cancel + await expiry task before guard release. cancel.cancel(); if let Some(t) = _nip_fi_admission_expiry.take() { @@ -1320,11 +1426,12 @@ pub(crate) async fn handle_active_audio_connection( Err(crate::audio::room::AdmissionError::VersionMismatch { pinned, requested }) => { shadow_attempt.refused(); info!(channel_id = %channel_id, pubkey = %pubkey_hex, pinned, requested, "audio: protocol version mismatch — upgrade required"); - crate::connection::send_exit_frames_bounded(&mut ws_send, [WsMessage::Text(serde_json::json!({ + crate::connection::send_exit_frames_bounded(&mut ws_send, error_exit_frames(serde_json::json!({ "type": "error", "code": "upgrade_required", "message": format!("this huddle is using audio protocol v{pinned}; your client requested v{requested}"), "pinned_version": pinned, "requested_version": requested, - }).to_string().into())]).await; + }), axum::extract::ws::close_code::POLICY)).await; + hand_off_socket(ws_recv, Some(ws_send)); // IMPORTANT 3: cancel + await expiry task before guard release. cancel.cancel(); if let Some(t) = _nip_fi_admission_expiry.take() { @@ -1626,6 +1733,7 @@ pub(crate) async fn handle_active_audio_connection( .map_or(WsMessage::Close(None), |reason| reason.close_message()); let _ = tokio::time::timeout_at(deadline, ws_send.send(close)).await; } + hand_off_socket(ws_recv, Some(ws_send)); if let (Some(session), Some(ref mut stream)) = ( guard.take_remote_session().as_ref(), guard.take_remote_stream().as_mut(), @@ -2021,20 +2129,29 @@ pub(crate) async fn handle_active_audio_connection( None }; - recv_loop( - ws_recv, - Arc::clone(&room), - peer_id, - requested_version, - ctrl_tx, - Arc::clone(&missed_pongs), - cancel.clone(), - remote_session.as_mut(), - ) - .await; + // Every exit that hands the socket off returns, so it is still held here; + // if it were not, the session would skip straight to teardown. + if let AudioReadHalf::Held(ws_recv) = ws_recv { + recv_loop( + ws_recv, + Arc::clone(&room), + peer_id, + requested_version, + ctrl_tx, + Arc::clone(&missed_pongs), + cancel.clone(), + remote_session.as_mut(), + ) + .await; + } cancel.cancel(); let _ = send_task.await; + // The writer has sent the Close and released the write half. Nothing + // below may hold the socket open, including the Postgres and Redis writes. + hand_off_socket(ws_recv, None); + #[cfg(test)] + crate::nip_fi_test_hooks::before_audio_teardown_cleanup(tenant.community()).await; let _ = heartbeat_task.await; let _ = forward_task.await; // The reader task owns the owner control stream; joining it here guarantees @@ -2226,7 +2343,7 @@ fn remote_rejection_ws_error(reason: &crate::audio::join::RegisterRejection) -> /// clarifies at this single call site. #[allow(clippy::too_many_arguments)] async fn recv_loop( - mut ws_recv: futures_util::stream::SplitStream, + ws_recv: &mut futures_util::stream::SplitStream, room: Arc, peer_id: Uuid, protocol_version: u8, @@ -2758,6 +2875,7 @@ pub(crate) async fn deny_audio_authorization( terminal_rx: &mut mpsc::Receiver, ) where S: futures_util::Sink + Unpin, + S::Error: std::error::Error + 'static, { control.disconnect_nip_fi(); crate::connection::send_exit_frames_bounded( @@ -2792,6 +2910,7 @@ async fn exit_authorization_refusal( off_mode: serde_json::Value, ) where S: futures_util::Sink + Unpin, + S::Error: std::error::Error + 'static, { if nip_fi && class == buzz_auth::DenialClass::AuthorizationDenied { deny_audio_authorization(ws_send, control, terminal_rx).await; @@ -5300,6 +5419,13 @@ mod tests { /// Reads `BUZZ_TEST_DATABASE_URL`; falls back to the local development URL. /// Returns `None` if the resolved database is not reachable. async fn audio_test_state_real_db() -> Option> { + audio_test_state_real_db_with(|_| {}).await + } + + /// `audio_test_state_real_db` with a final `configure` pass over the config. + async fn audio_test_state_real_db_with( + configure: impl FnOnce(&mut crate::config::Config), + ) -> Option> { use std::sync::Arc; let db_url = std::env::var("BUZZ_TEST_DATABASE_URL").unwrap_or_else(|_| { "postgres://buzz:buzz_dev@localhost:5432/buzz".to_string() // sadscan:disable np.postgres.1 -- local test-only credentials @@ -5312,6 +5438,7 @@ mod tests { config.require_relay_membership = false; config.database_url = db_url.clone(); config.redis_url = "redis://127.0.0.1:1".to_string(); + configure(&mut config); let pool = sqlx::PgPool::connect_lazy(&config.database_url).expect("lazy pg pool"); let db = buzz_db::Db::from_pool(pool.clone()); let redis_pool = deadpool_redis::Config::from_url(&config.redis_url) @@ -9992,6 +10119,369 @@ mod tests { let _ = server.await; } + /// Answers the server's NIP-42 challenge with an audio `auth` message. + async fn send_audio_auth( + client: &mut tokio_tungstenite::WebSocketStream< + tokio_tungstenite::MaybeTlsStream, + >, + key: &nostr::Keys, + relay_url: &str, + protocol_version: u8, + ) { + let challenge = + match tokio::time::timeout(std::time::Duration::from_secs(2), client.next()) + .await + .expect("challenge timeout") + .expect("challenge msg") + .expect("challenge ws msg") + { + tokio_tungstenite::tungstenite::Message::Text(t) => { + serde_json::from_str::(&t).expect("challenge JSON") + ["challenge"] + .as_str() + .expect("challenge field") + .to_string() + } + other => panic!("expected text challenge; got {other:?}"), + }; + let auth_event = nostr::EventBuilder::new(nostr::Kind::Authentication, "") + .tag(nostr::Tag::parse(["relay", relay_url]).unwrap()) + .tag(nostr::Tag::parse(["challenge", &challenge]).unwrap()) + .sign_with_keys(key) + .unwrap(); + let auth_msg = serde_json::json!({ + "type": "auth", + "event": auth_event, + "parent_channel_id": null, + "protocol_version": protocol_version, + }) + .to_string(); + client + .send(tokio_tungstenite::tungstenite::Message::Text( + auth_msg.into(), + )) + .await + .expect("send auth msg"); + } + + /// After an established session's writer sends its Close, the + /// teardown that follows writes to Postgres and Redis, which can stall + /// with no deadline. The socket must still end within the bounded wait + /// for the client's Close reply. This client never replies and the + /// teardown never finishes. + /// + /// Mutation: move `hand_off_socket(ws_recv, None)` to after the + /// teardown → the socket stays open while the teardown is held → + /// `time_to_eof` hits its watchdog and panics. + #[ignore = "requires Postgres — runs in postgres-ci nextest lane"] + #[tokio::test] + async fn stalled_teardown_does_not_hold_established_socket_open() { + use crate::connection::terminal_close_delivery_tests::time_to_eof; + use crate::connection::WS_TERMINAL_FLUSH_TIMEOUT; + use buzz_auth::VerifiedAssertion; + use chrono::{Duration, Utc}; + use std::sync::Arc; + + let state = audio_test_state_real_db().await.expect( + "PostgreSQL must be available — set BUZZ_TEST_DATABASE_URL or start local postgres", + ); + let pool = state.db.pool().clone(); + let (tenant, channel_id, key) = seed_audio_fixture(&pool).await; + let community = tenant.community(); + let assertion = VerifiedAssertion::for_test( + Some(key.public_key()), + vec![Utc::now() + Duration::hours(1)], + ); + let relay_url = format!("ws://{}", tenant.host()); + let conn_cancel = CancellationToken::new(); + + let (fanout_rx, fanout_release) = + crate::nip_fi_test_hooks::audio_participant_fanout_hook::arm(community); + let (teardown_rx, teardown_release) = + crate::nip_fi_test_hooks::audio_teardown_cleanup_hook::arm(community); + + let app = axum::Router::new().route( + "/", + axum::routing::get({ + let conn_cancel = conn_cancel.clone(); + move |ws: axum::extract::ws::WebSocketUpgrade| { + let (state, tenant, assertion) = + (Arc::clone(&state), tenant.clone(), assertion.clone()); + let control = + crate::state::CommunityConnectionControl::new(conn_cancel.clone()); + async move { + ws.on_upgrade(move |socket| { + handle_active_audio_connection( + socket, + state, + tenant, + channel_id, + control, + Some(assertion), + Utc::now(), + None, + ) + }) + } + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind test listener"); + let addr = listener.local_addr().expect("test listener addr"); + let server = tokio::spawn(async move { axum::serve(listener, app).await }); + + let (mut client, _) = tokio_tungstenite::connect_async(format!("ws://{addr}/")) + .await + .expect("connect client"); + send_audio_auth(&mut client, &key, &relay_url, 1).await; + + tokio::time::timeout(std::time::Duration::from_secs(10), fanout_rx) + .await + .expect("handler must commit the join within 10s") + .expect("fanout channel closed"); + fanout_release.notify_one(); + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + + // From here the client only reads raw bytes, so it never answers + // the server's Close. + let tokio_tungstenite::MaybeTlsStream::Plain(socket) = client.get_mut() else { + panic!("plain TCP client"); + }; + let started = std::time::Instant::now(); + conn_cancel.cancel(); + tokio::time::timeout(std::time::Duration::from_secs(5), teardown_rx) + .await + .expect("teardown must reach the stall point") + .expect("teardown hook dropped"); + // The hook runs just after the hand-off, so the socket must end + // within the bounded wait measured from here. + let handed_off = std::time::Instant::now(); + let elapsed = time_to_eof(socket, started).await; + assert!( + elapsed >= WS_TERMINAL_FLUSH_TIMEOUT, + "server must wait for the client's Close reply, ended after {elapsed:?}" + ); + let since_hand_off = handed_off.elapsed(); + assert!( + since_hand_off <= WS_TERMINAL_FLUSH_TIMEOUT + std::time::Duration::from_millis(500), + "socket must end within the bounded wait after the hand-off, took {since_hand_off:?}" + ); + + teardown_release.notify_one(); + server.abort(); + } + + /// The archived-channel exit sends its error and Close, then releases + /// the huddle lease in Redis, which can stall with no deadline. The + /// socket must end within the bounded wait for the client's Close + /// reply, not when that cleanup finishes. This client never replies + /// and the cleanup never finishes. + /// + /// Mutation: move that exit's `hand_off_socket` below the lease + /// release (or remove it) → the socket stays open while the cleanup + /// is held → `time_to_eof` hits its watchdog and panics. + #[ignore = "requires Postgres — runs in postgres-ci nextest lane"] + #[tokio::test] + async fn stalled_archived_exit_cleanup_does_not_hold_socket_open() { + use crate::connection::terminal_close_delivery_tests::time_to_eof; + use crate::connection::WS_TERMINAL_FLUSH_TIMEOUT; + use buzz_auth::VerifiedAssertion; + use chrono::{Duration, Utc}; + use std::sync::Arc; + + let state = audio_test_state_real_db().await.expect( + "PostgreSQL must be available — set BUZZ_TEST_DATABASE_URL or start local postgres", + ); + let pool = state.db.pool().clone(); + let (tenant, channel_id, key) = seed_audio_fixture(&pool).await; + let community = tenant.community(); + let assertion = VerifiedAssertion::for_test( + Some(key.public_key()), + vec![Utc::now() + Duration::hours(1)], + ); + let relay_url = format!("ws://{}", tenant.host()); + + let (recheck_rx, recheck_release) = + crate::nip_fi_test_hooks::room_archive_recheck_hook::arm(community); + let (cleanup_rx, cleanup_release) = + crate::nip_fi_test_hooks::archived_exit_cleanup_hook::arm(community); + + let app = axum::Router::new().route( + "/", + axum::routing::get(move |ws: axum::extract::ws::WebSocketUpgrade| { + let (state, tenant, assertion) = + (Arc::clone(&state), tenant.clone(), assertion.clone()); + let control = + crate::state::CommunityConnectionControl::new(CancellationToken::new()); + async move { + ws.on_upgrade(move |socket| { + handle_active_audio_connection( + socket, + state, + tenant, + channel_id, + control, + Some(assertion), + Utc::now(), + None, + ) + }) + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind test listener"); + let addr = listener.local_addr().expect("test listener addr"); + let server = tokio::spawn(async move { axum::serve(listener, app).await }); + + let (mut client, _) = tokio_tungstenite::connect_async(format!("ws://{addr}/")) + .await + .expect("connect client"); + send_audio_auth(&mut client, &key, &relay_url, 1).await; + + // Admission passed the membership check; archive the channel + // before the re-check so the join takes the archived exit. + tokio::time::timeout(std::time::Duration::from_secs(10), recheck_rx) + .await + .expect("admission must reach the archive re-check") + .expect("re-check hook dropped"); + sqlx::query("UPDATE channels SET archived_at = NOW() WHERE id = $1") + .bind(channel_id) + .execute(&pool) + .await + .expect("archive channel"); + + // From here the client only reads raw bytes, so it never answers + // the server's Close. + let tokio_tungstenite::MaybeTlsStream::Plain(socket) = client.get_mut() else { + panic!("plain TCP client"); + }; + let started = std::time::Instant::now(); + recheck_release.notify_one(); + tokio::time::timeout(std::time::Duration::from_secs(5), cleanup_rx) + .await + .expect("the archived exit must reach its cleanup") + .expect("cleanup hook dropped"); + let elapsed = time_to_eof(socket, started).await; + assert!( + elapsed <= WS_TERMINAL_FLUSH_TIMEOUT + std::time::Duration::from_millis(500), + "a stalled exit cleanup must not hold the socket open, ended after {elapsed:?}" + ); + + cleanup_release.notify_one(); + server.abort(); + } + + /// A fast refusal hands the socket to a background close and returns. + /// The connection permit must stay held until that socket ends, or a + /// client that never answers the Close lets open audio sockets exceed + /// `BUZZ_MAX_CONNECTIONS`. + /// + /// Mutation: stop the wrapper from joining the background close → the + /// permit is released while the socket is still open → the + /// `available_permits() == 0` assertion fails. + #[ignore = "requires Postgres — runs in postgres-ci nextest lane"] + #[tokio::test] + async fn background_close_holds_connection_permit_until_socket_ends() { + use crate::connection::terminal_close_delivery_tests::time_to_eof; + use buzz_auth::VerifiedAssertion; + use chrono::{Duration, Utc}; + use std::sync::Arc; + use tokio::io::AsyncReadExt as _; + + let state = audio_test_state_real_db().await.expect( + "PostgreSQL must be available — set BUZZ_TEST_DATABASE_URL or start local postgres", + ); + let pool = state.db.pool().clone(); + let (tenant, channel_id, key) = seed_audio_fixture(&pool).await; + let assertion = VerifiedAssertion::for_test( + Some(key.public_key()), + vec![Utc::now() + Duration::hours(1)], + ); + let relay_url = format!("ws://{}", tenant.host()); + let semaphore = Arc::new(tokio::sync::Semaphore::new(1)); + + let app = axum::Router::new().route( + "/", + axum::routing::get({ + let semaphore = Arc::clone(&semaphore); + move |ws: axum::extract::ws::WebSocketUpgrade| { + let (state, tenant, assertion) = + (Arc::clone(&state), tenant.clone(), assertion.clone()); + let permit = Arc::clone(&semaphore) + .try_acquire_owned() + .expect("one connection permit"); + async move { + ws.on_upgrade(move |socket| { + handle_audio_connection( + socket, + state, + tenant, + channel_id, + permit, + Some(assertion), + None, + Utc::now(), + ) + }) + } + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind test listener"); + let addr = listener.local_addr().expect("test listener addr"); + let server = tokio::spawn(async move { axum::serve(listener, app).await }); + + let (mut client, _) = tokio_tungstenite::connect_async(format!("ws://{addr}/")) + .await + .expect("connect client"); + send_audio_auth(&mut client, &key, &relay_url, u8::MAX).await; + + // From here the client only reads raw bytes, so it never answers + // the server's Close. + let tokio_tungstenite::MaybeTlsStream::Plain(socket) = client.get_mut() else { + panic!("plain TCP client"); + }; + let mut received = Vec::new(); + let mut buf = [0u8; 512]; + while !received + .windows(b"unsupported_version".len()) + .any(|w| w == b"unsupported_version") + { + let n = + tokio::time::timeout(std::time::Duration::from_secs(5), socket.read(&mut buf)) + .await + .expect("refusal must arrive") + .expect("read"); + assert_ne!(n, 0, "socket ended before the refusal"); + received.extend_from_slice(&buf[..n]); + } + let started = std::time::Instant::now(); + // The refusal has no staged lease, so its cleanup returns at once. + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + assert_eq!( + semaphore.available_permits(), + 0, + "permit must stay held while the refused socket is open" + ); + + time_to_eof(socket, started).await; + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while semaphore.available_permits() == 0 { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .expect("permit must be released once the socket ends"); + + server.abort(); + } + // ── F1: generation fencing witness ──────────────────────────────────────── // // `commit_participant_join` must include `generation` in the committed @@ -13756,6 +14246,39 @@ mod tests { assert_eq!(close.reason.as_str(), "access revoked"); } + /// A deployment without huddle audio refuses the join with its error + /// frame and then a Close 1008, so the client gets a clean handshake + /// instead of waiting out the peer wait. + /// Mutation: drop the Close from the refusal → no Close frame → RED. + #[tokio::test] + #[ignore = "requires Postgres — runs in postgres-ci nextest lane"] + async fn huddle_audio_unavailable_sends_error_then_policy_close() { + let state = audio_test_state_real_db_with(|c| c.huddle_audio_available = false) + .await + .expect( + "PostgreSQL must be available — set BUZZ_TEST_DATABASE_URL or start local postgres", + ); + let (tenant, channel_id, key) = seed_audio_fixture(state.db.pool()).await; + let (mut client, server) = + start_audio_auth(&state, tenant, channel_id, &key, None, None).await; + + let error = next_audio_text(&mut client).await; + assert!( + error.contains("\"huddle_audio_unavailable\""), + "expected the unavailable refusal; got {error}" + ); + match tokio::time::timeout(std::time::Duration::from_secs(5), client.next()) + .await + .expect("close must follow the refusal") + { + Some(Ok(tokio_tungstenite::tungstenite::Message::Close(Some(frame)))) => { + assert_eq!(u16::from(frame.code), 1008) + } + other => panic!("expected Close 1008 after the refusal; got {other:?}"), + } + server.abort(); + } + /// Adds `key` as a member of the audio fixture's channel. async fn add_channel_member( pool: &sqlx::PgPool, @@ -14662,4 +15185,99 @@ mod tests { future never resolves → test hangs → RED" ); } + + /// Real-socket proof that each audio exit path holds the socket for the + /// client's Close reply. The raw client never replies, so a server that + /// drops the socket right after its Close ends it before the bounded wait. + mod terminal_close_delivery { + use super::*; + use crate::connection::terminal_close_delivery_tests::{ + raw_client_after_server_close, time_to_eof, + }; + use crate::connection::WS_TERMINAL_FLUSH_TIMEOUT; + + fn expired_assertion() -> VerifiedAssertion { + let deadline = chrono::Utc::now() - chrono::Duration::seconds(1); + VerifiedAssertion::for_test(Some(nostr::Keys::generate().public_key()), vec![deadline]) + } + + fn tenant() -> TenantContext { + TenantContext::resolved( + buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), + "test.local".to_owned(), + ) + } + + async fn assert_holds_for_close_reply(app: Router) { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let address = listener.local_addr().expect("listener address"); + tokio::spawn(async move { axum::serve(listener, app).await.expect("serve") }); + + let started = std::time::Instant::now(); + let mut socket = raw_client_after_server_close(address).await; + let elapsed = time_to_eof(&mut socket, started).await; + assert!( + elapsed >= WS_TERMINAL_FLUSH_TIMEOUT, + "server must wait for the client's Close reply, ended after {elapsed:?}" + ); + } + + /// Expired session refused during bootstrap: the `on_not_run` drain. + #[tokio::test] + async fn bootstrap_drain_holds_socket_for_client_close_reply() { + let state = audio_test_state().await; + let app = Router::new().route( + "/", + get(move |ws: WebSocketUpgrade| { + let state = Arc::clone(&state); + let permit = Arc::new(tokio::sync::Semaphore::new(1)) + .try_acquire_owned() + .expect("permit"); + async move { + ws.on_upgrade(move |socket| { + handle_audio_connection( + socket, + state, + tenant(), + Uuid::new_v4(), + permit, + Some(expired_assertion()), + None, + chrono::Utc::now(), + ) + }) + } + }), + ); + assert_holds_for_close_reply(app).await; + } + + /// Expired session inside the active handler: an early admission exit. + #[tokio::test] + async fn active_handler_exit_holds_socket_for_client_close_reply() { + let state = audio_test_state().await; + let app = Router::new().route( + "/", + get(move |ws: WebSocketUpgrade| { + let state = Arc::clone(&state); + async move { + ws.on_upgrade(move |socket| { + let cancel = CancellationToken::new(); + handle_active_audio_connection( + socket, + state, + tenant(), + Uuid::new_v4(), + CommunityConnectionControl::new(cancel), + Some(expired_assertion()), + chrono::Utc::now(), + None, + ) + }) + } + }), + ); + assert_holds_for_close_reply(app).await; + } + } } diff --git a/crates/buzz-relay/src/connection.rs b/crates/buzz-relay/src/connection.rs index b341f9b214..ad58fd032f 100644 --- a/crates/buzz-relay/src/connection.rs +++ b/crates/buzz-relay/src/connection.rs @@ -29,8 +29,10 @@ use crate::state::{ /// Maximum time a new socket may hold a connection slot without completing NIP-42 auth. const AUTH_TIMEOUT: Duration = Duration::from_secs(5); -/// Maximum time the writer may spend flushing terminal frames after cancellation. -/// This stays well inside the process-wide 30-second hard drain. +/// Bounds each of the two stages of ending a socket: writing its terminal +/// frames after cancellation, then waiting for the client's Close reply +/// (`await_client_close`). Both stay well inside the process-wide 30-second +/// hard drain. pub(crate) const WS_TERMINAL_FLUSH_TIMEOUT: Duration = Duration::from_secs(1); /// Shared mutable subscription map for a single WebSocket connection: sub ID → @@ -362,6 +364,12 @@ pub(crate) async fn handle_connection( nip_fi_shadow: Option>, connection_time: chrono::DateTime, ) { + // Held through bootstrap, the active handler and any refusal drain, so a + // socket counts against `BUZZ_MAX_CONNECTIONS` for as long as it is open. + let Ok(_permit) = state.conn_semaphore.clone().try_acquire_owned() else { + warn!("Connection limit reached, rejecting {addr}"); + return; + }; let conn_id = Uuid::new_v4(); let cancel = CancellationToken::new(); @@ -417,6 +425,11 @@ pub(crate) async fn handle_connection( let rx_for_run = Arc::clone(&rx_shared); let rx_for_drain = Arc::clone(&rx_shared); + // The drain path parks the read half here so the bounded wait for the + // client's Close runs after the registry entry is released. + let unread_shared = Arc::new(Mutex::new(None)); + let unread_for_drain = Arc::clone(&unread_shared); + run_registered_community_connection( ®istry, conn_id, @@ -453,45 +466,22 @@ pub(crate) async fn handle_connection( .await }, move || async move { - // Drain the pre-terminal channel (NIP-FI denial frame, if any) and - // close the socket with a bounded timeout so a queued denial is - // delivered even when the active-connection path is never entered. + // Deliver a queued denial frame (if any) and the reason's Close + // even when the active-connection path is never entered. let socket = socket_for_drain.lock().await.take(); - let mut pre_terminal_ctrl_rx = rx_for_drain.lock().await.take(); + let queued = rx_for_drain.lock().await.take(); if let Some(socket) = socket { - let (mut ws_send, _ws_recv) = socket.split(); - // Terminal channel has capacity 1; if the expiry task fired, it holds - // exactly one denial frame here. Deliver it before closing. - if let Some(ref mut rx) = pre_terminal_ctrl_rx { - while let Ok(msg) = rx.try_recv() { - let _ = tokio::time::timeout( - WS_TERMINAL_FLUSH_TIMEOUT, - futures_util::SinkExt::send(&mut ws_send, msg), - ) - .await; - } - } - // A published reason (e.g. an expiry denial) closes with its - // own frame, exactly as the writer would. let close = drain_reason.borrow().map(|reason| reason.close_message()); - if let Some(close) = close { - let _ = tokio::time::timeout( - WS_TERMINAL_FLUSH_TIMEOUT, - futures_util::SinkExt::send(&mut ws_send, close), - ) - .await; - } - // Close the socket so the client sees a clean close rather than - // an abrupt TCP reset. - let _ = tokio::time::timeout( - WS_TERMINAL_FLUSH_TIMEOUT, - futures_util::SinkExt::close(&mut ws_send), - ) - .await; + let ws_recv = drain_refused_socket(socket, queued, close).await; + *unread_for_drain.lock().await = Some(ws_recv); } }, ) .await; + let unread = unread_shared.lock().await.take(); + if let Some(ws_recv) = unread { + await_client_close(ws_recv).await; + } } // `handle_active_connection` inherits the connection handler's natural parameter @@ -517,13 +507,6 @@ async fn handle_active_connection( // connection_time is threaded in from the HTTP handler (captured immediately // before on_upgrade) so the session partition is rooted at the true upgrade // instant, not the post-community-active-check instant. [FI-TRACE-LEASE-BOUND] - let permit = match state.conn_semaphore.clone().try_acquire_owned() { - Ok(p) => p, - Err(_) => { - warn!("Connection limit reached, rejecting {addr}"); - return; - } - }; let challenge = generate_challenge(); @@ -551,6 +534,8 @@ async fn handle_active_connection( let backpressure_count = Arc::new(AtomicU8::new(0)); let subscriptions = Arc::new(Mutex::new(HashMap::new())); + #[cfg(test)] + test_seam::hold_subscriptions(&subscriptions).await; // Compute the NIP-FI session deadline from the assertion. // @@ -729,7 +714,7 @@ async fn handle_active_connection( }; auth_cancel_conn.finish_pending_auth_on_cancel(outcome); }); - recv_loop( + let (ws_recv, in_flight) = recv_loop( ws_recv, Arc::clone(&conn), Arc::clone(&state), @@ -749,30 +734,43 @@ async fn handle_active_connection( let authenticated = auth_lifecycle.finish(close_outcome); let _ = send_task.await; - let _ = heartbeat_task.await; - let _ = auth_timeout_task.await; - if let Some(task) = nip_fi_expiry_task { - let _ = task.await; - } - let _ = auth_cancel_task.await; - crate::handlers::close::release_connection_subscriptions(&conn, &state).await; - state.conn_manager.deregister(conn.conn_id); - if let Some(auth_ctx) = authenticated { - if !state.conn_manager.has_admitted_connection( - conn.tenant.community(), - auth_ctx.pubkey.to_bytes().as_slice(), - ) { - let _ = state - .pubsub - .clear_presence(&conn.tenant, &auth_ctx.pubkey) - .await; + // The writer is done, so the server's Close (if any) is sent. Wait on the + // peer alongside the remaining joins and cleanup: the expiry barrier and + // Redis can block without a deadline, and neither may hold the socket + // open past the bounded peer wait. + let cleanup = async { + // A frame handler interrupted by cancellation was spawned, so it keeps + // running (and can release any effect permit the expiry barrier waits + // on). Finish it before application cleanup so nothing is half-retired. + if let Some(task) = in_flight { + if let Err(error) = task.await { + warn!(conn_id = %conn.conn_id, "interrupted frame handler failed: {error}"); + } } - } + let _ = heartbeat_task.await; + let _ = auth_timeout_task.await; + if let Some(task) = nip_fi_expiry_task { + let _ = task.await; + } + let _ = auth_cancel_task.await; + crate::handlers::close::release_connection_subscriptions(&conn, &state).await; + state.conn_manager.deregister(conn.conn_id); + if let Some(auth_ctx) = authenticated { + if !state.conn_manager.has_admitted_connection( + conn.tenant.community(), + auth_ctx.pubkey.to_bytes().as_slice(), + ) { + let _ = state + .pubsub + .clear_presence(&conn.tenant, &auth_ctx.pubkey) + .await; + } + } + }; + tokio::join!(cleanup, await_client_close(ws_recv)); metrics::gauge!("buzz_ws_connections_active").decrement(1.0); info!(conn_id = %conn_id, addr = %addr, "WebSocket connection closed"); - - drop(permit); } /// Send WebSocket messages in priority order: control frames before data frames. @@ -871,31 +869,23 @@ async fn flush_terminal_frames( first_ctrl: Option, ) where S: Sink + Unpin, + S::Error: std::error::Error + 'static, { let deadline = tokio::time::Instant::now() + WS_TERMINAL_FLUSH_TIMEOUT; // 1. Drain FI terminal channel first — denial frame must precede the close. while let Ok(terminal_msg) = terminal_ctrl_rx.try_recv() { - if !matches!( - tokio::time::timeout_at(deadline, sink.send(terminal_msg)).await, - Ok(Ok(())) - ) { + if !send_terminal_frame(sink, terminal_msg, deadline, "terminal").await { return; } } // 2. Drain ordinary control frames (ban reason, etc.). if let Some(ctrl_msg) = first_ctrl { - if !matches!( - tokio::time::timeout_at(deadline, sink.send(ctrl_msg)).await, - Ok(Ok(())) - ) { + if !send_terminal_frame(sink, ctrl_msg, deadline, "control").await { return; } } while let Ok(ctrl_msg) = ctrl_rx.try_recv() { - if !matches!( - tokio::time::timeout_at(deadline, sink.send(ctrl_msg)).await, - Ok(Ok(())) - ) { + if !send_terminal_frame(sink, ctrl_msg, deadline, "control").await { return; } } @@ -903,7 +893,46 @@ async fn flush_terminal_frames( let close = disconnect_reason .borrow() .map_or(WsMessage::Close(None), |reason| reason.close_message()); - let _ = tokio::time::timeout_at(deadline, sink.send(close)).await; + send_terminal_frame(sink, close, deadline, "close").await; +} + +/// One deadline-bounded terminal write. Failures stay best-effort but are +/// logged, so an undelivered denial or Close is diagnosable. +async fn send_terminal_frame( + sink: &mut S, + frame: WsMessage, + deadline: tokio::time::Instant, + frame_kind: &'static str, +) -> bool +where + S: Sink + Unpin, + S::Error: std::error::Error + 'static, +{ + match tokio::time::timeout_at(deadline, sink.send(frame)).await { + Ok(Ok(())) => true, + Ok(Err(error)) => { + if !closed_by_peer(&error) { + debug!(frame_kind, "terminal frame send failed: {error}"); + } + false + } + Err(_) => { + debug!(frame_kind, "terminal frame send timed out"); + false + } + } +} + +/// Whether a write failed only because the close handshake is already under +/// way, as when the client closed first. Not a delivery failure, so not logged. +fn closed_by_peer(error: &(dyn std::error::Error + 'static)) -> bool { + use tungstenite::error::{Error, ProtocolError}; + std::iter::successors(Some(error), |error| error.source()).any(|error| { + matches!( + error.downcast_ref::(), + Some(Error::AlreadyClosed | Error::Protocol(ProtocolError::SendAfterClosing)) + ) + }) } /// Writes a WebSocket exit path's frames before a writer task owns the socket. @@ -916,16 +945,55 @@ pub(crate) async fn send_exit_frames_bounded( frames: impl IntoIterator, ) where S: Sink + Unpin, + S::Error: std::error::Error + 'static, { let deadline = tokio::time::Instant::now() + WS_TERMINAL_FLUSH_TIMEOUT; + send_frames_by(sink, frames, deadline).await; +} + +/// Writes `frames` in order until one fails or `deadline` passes. Returns +/// whether every frame was written. +async fn send_frames_by( + sink: &mut S, + frames: impl IntoIterator, + deadline: tokio::time::Instant, +) -> bool +where + S: Sink + Unpin, + S::Error: std::error::Error + 'static, +{ for frame in frames { - if !matches!( - tokio::time::timeout_at(deadline, sink.send(frame)).await, - Ok(Ok(())) - ) { - return; + if !send_terminal_frame(sink, frame, deadline, "exit").await { + return false; } } + true +} + +/// Ends a socket refused before its session handler ran: writes the frames +/// still `queued` for it, then `close`, then closes the sink, all under one +/// `WS_TERMINAL_FLUSH_TIMEOUT` deadline. Returns the read half, which the +/// caller passes to `await_client_close` once it has released anything the +/// wait must not hold. Shared by the relay and audio routes. +pub(crate) async fn drain_refused_socket( + socket: WebSocket, + queued: Option>, + close: Option, +) -> futures_util::stream::SplitStream { + let (mut ws_send, ws_recv) = socket.split(); + let deadline = tokio::time::Instant::now() + WS_TERMINAL_FLUSH_TIMEOUT; + let queued = queued + .into_iter() + .flat_map(|mut rx| std::iter::from_fn(move || rx.try_recv().ok())); + if send_frames_by(&mut ws_send, queued.chain(close), deadline).await { + match tokio::time::timeout_at(deadline, ws_send.close()).await { + Ok(Ok(())) => {} + Ok(Err(error)) if closed_by_peer(&error) => {} + Ok(Err(error)) => debug!("refused socket close failed: {error}"), + Err(_) => debug!("refused socket close timed out"), + } + } + ws_recv } async fn send_loop_inner( @@ -938,6 +1006,7 @@ async fn send_loop_inner( disconnect_reason: watch::Receiver>, ) where S: Sink + Unpin, + S::Error: std::error::Error + 'static, { loop { // Priority: drain all pending control frames before data. @@ -1102,12 +1171,20 @@ async fn heartbeat_loop( } } +/// Reads and handles client frames until the connection is cancelled or the +/// stream ends. Returns the read half, for the bounded wait on the client's +/// Close, and the handle of a frame handler that cancellation interrupted, if +/// any: the caller must join that handle before subscription cleanup so the +/// interrupted handler never leaves a subscription half-retired. async fn recv_loop( mut ws_recv: futures_util::stream::SplitStream, conn: Arc, state: Arc, missed_pongs: Arc, cancel: CancellationToken, +) -> ( + futures_util::stream::SplitStream, + Option>, ) { loop { tokio::select! { @@ -1130,7 +1207,10 @@ async fn recv_loop( break; } trace!(len = text.len(), "frame received"); - handle_text_message(text.to_string(), Arc::clone(&conn), Arc::clone(&state)).await; + let frame = handle_text_message(text.to_string(), Arc::clone(&conn), Arc::clone(&state)); + if let Some(task) = run_frame(frame, &cancel).await { + return (ws_recv, Some(task)); + } } Some(Ok(WsMessage::Binary(bytes))) => { let max_frame_bytes = state.config.max_frame_bytes; @@ -1152,7 +1232,10 @@ async fn recv_loop( // (notably certain Nostr libraries) send text payloads in binary frames. // NIP-01 is text-only, but accepting binary is a common relay extension. if let Ok(text) = String::from_utf8(bytes.to_vec()) { - handle_text_message(text, Arc::clone(&conn), Arc::clone(&state)).await; + let frame = handle_text_message(text, Arc::clone(&conn), Arc::clone(&state)); + if let Some(task) = run_frame(frame, &cancel).await { + return (ws_recv, Some(task)); + } } } Some(Ok(WsMessage::Pong(_))) => { @@ -1181,6 +1264,42 @@ async fn recv_loop( _ = cancel.cancelled() => break, } } + (ws_recv, None) +} + +/// Runs one frame handler inline, keeping frames serial while the connection +/// is live. If the connection is cancelled first, the handler is spawned to +/// finish rather than dropped (dropping could half-retire a subscription), and +/// its handle is returned so the reader can give back its read half at once. +async fn run_frame( + frame: impl std::future::Future + Send + 'static, + cancel: &CancellationToken, +) -> Option> { + let mut frame = Box::pin(frame); + tokio::select! { + biased; + () = &mut frame => None, + () = cancel.cancelled() => Some(tokio::spawn(frame)), + } +} + +/// Holds the read half open until the close handshake finishes and the stream +/// ends, bounded by `WS_TERMINAL_FLUSH_TIMEOUT`. Dropping TCP before the peer +/// finishes the close handshake (including in-flight peer frames) can make the +/// client lose the unread NOTICE and Close and report an abnormal `1006`. +/// +/// Reading on past the client's Close matters: when the client closed first, +/// tungstenite only queues our reply Close and writes it on the next read, so +/// stopping at the client's Close would drop TCP with the reply unsent. Once +/// the handshake completes, the next read ends the stream at once. +pub(crate) async fn await_client_close(mut ws_recv: futures_util::stream::SplitStream) { + let drained = tokio::time::timeout(WS_TERMINAL_FLUSH_TIMEOUT, async { + while let Some(Ok(_)) = ws_recv.next().await {} + }) + .await; + if drained.is_err() { + debug!("client did not answer Close before the terminal deadline"); + } } async fn handle_text_message(text: String, conn: Arc, state: Arc) { @@ -1303,6 +1422,28 @@ async fn handle_text_message(text: String, conn: Arc, state: Ar } } +#[cfg(test)] +pub(crate) mod test_seam { + use super::*; + + pub(crate) type Hold = + tokio::sync::oneshot::Sender>>; + + tokio::task_local! { + /// Hands the test this connection's subscription lock, already held, + /// before the reader can take it. + pub(crate) static SUBSCRIPTIONS_HOLD: StdMutex>; + } + + pub(super) async fn hold_subscriptions(subscriptions: &ConnectionSubscriptions) { + if let Ok(Some(hold)) = + SUBSCRIPTIONS_HOLD.try_with(|hold| hold.lock().expect("hold").take()) + { + let _ = hold.send(Arc::clone(subscriptions).lock_owned().await); + } + } +} + #[cfg(test)] pub(crate) mod tests { use super::*; @@ -3198,3 +3339,665 @@ pub(crate) mod tests { let _ = server.await; } } + +/// Close-code delivery of NIP-FI authorization denials over a real WebSocket. +/// The mock-sink writer tests prove frame order; these prove what a client +/// actually reads off the wire. +#[cfg(test)] +pub(crate) mod terminal_close_delivery_tests { + use super::*; + use axum::{extract::ws::WebSocketUpgrade, routing::get, Router}; + use buzz_auth::VerifiedAssertion; + use nostr::{EventBuilder, Keys, RelayUrl}; + use tokio::net::TcpListener; + use tokio_tungstenite::tungstenite::{protocol::frame::coding::CloseCode, Message}; + + type Client = tokio_tungstenite::WebSocketStream< + tokio_tungstenite::MaybeTlsStream, + >; + + /// Serves `handle_active_connection` for one assertion. `prefire_expiry` + /// arms an already-expired session gate, as `handle_connection` does. + async fn serve( + assertion: VerifiedAssertion, + prefire_expiry: bool, + ) -> (std::net::SocketAddr, RelayUrl, Arc, Uuid) { + serve_with_seams(assertion, prefire_expiry, Seams::default()).await + } + + /// How long after upgrade a session armed with `Seams::effect_hold` + /// expires. Long enough that a loaded runner still delivers a frame the + /// client sends at once before the deadline. + const EFFECT_HOLD_EXPIRY: chrono::Duration = chrono::Duration::seconds(2); + + /// Optional test hooks for `serve_with_seams`. + #[derive(Default)] + struct Seams { + /// Runs the handler inside this pause so the test can stall the + /// connection's final subscription cleanup. + cleanup_pause: Option>, + /// The session instead expires `EFFECT_HOLD_EXPIRY` after upgrade, and + /// the test receives a session effect permit taken before the + /// deadline, so the expiry task stalls at its quiescence barrier. + effect_hold: Option>, + /// Receives the connection's subscription lock, taken before the + /// reader starts, so a frame handler that needs it blocks until the + /// test drops it. + subscriptions_hold: Option, + } + + /// Like `serve`, with the test hooks in `seams`. + async fn serve_with_seams( + assertion: VerifiedAssertion, + prefire_expiry: bool, + Seams { + cleanup_pause, + effect_hold, + subscriptions_hold, + }: Seams, + ) -> (std::net::SocketAddr, RelayUrl, Arc, Uuid) { + let effect_hold = Arc::new(std::sync::Mutex::new(effect_hold)); + let subscriptions_hold = Arc::new(std::sync::Mutex::new(subscriptions_hold)); + let state = crate::state::tests::test_state().await; + let served_state = Arc::clone(&state); + let conn_id = Uuid::new_v4(); + let tenant = TenantContext::resolved( + buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), + "test.local".to_owned(), + ); + let relay_url = + crate::api::bridge::nip42_expected_relay_url(&state.config.relay_url, &tenant) + .parse() + .expect("expected NIP-42 relay URL"); + let app = Router::new().route( + "/", + get(move |ws: WebSocketUpgrade| { + let (state, tenant, assertion, cleanup_pause) = ( + Arc::clone(&state), + tenant.clone(), + assertion.clone(), + cleanup_pause.clone(), + ); + let effect_hold = effect_hold.lock().expect("effect hold").take(); + let subscriptions_hold = subscriptions_hold.lock().expect("subs hold").take(); + async move { + ws.on_upgrade(move |socket| async move { + let cancel = CancellationToken::new(); + let control = CommunityConnectionControl::new(cancel.clone()); + let offset = if effect_hold.is_some() { + EFFECT_HOLD_EXPIRY + } else { + chrono::Duration::seconds(-1) + }; + let pre_built = prefire_expiry.then(|| { + let deadline = chrono::Utc::now() + offset; + let (tx, rx) = mpsc::channel::(1); + let gate = + crate::nip_fi_gate::SessionAdmissionGate::new(deadline, cancel); + let expiry = crate::nip_fi_session::spawn_nip_fi_expiry_task( + deadline, + Arc::clone(&gate), + control.clone(), + tx.clone(), + crate::nip_fi_session::NipFiWsRoute::Root, + ); + (gate, tx, rx, Some(expiry)) + }); + if let (Some(hold), Some((gate, ..))) = (effect_hold, &pre_built) { + let permit = gate.acquire_effect().await.expect("effect permit"); + let _ = hold.send(permit); + } + let handler = handle_active_connection( + socket, + state, + "127.0.0.1:1234".parse().expect("client address"), + tenant, + conn_id, + control, + Some(assertion), + chrono::Utc::now(), + pre_built, + ); + let handler = super::test_seam::SUBSCRIPTIONS_HOLD + .scope(std::sync::Mutex::new(subscriptions_hold), handler); + match cleanup_pause { + Some(pause) => { + crate::handlers::close::test_seam::CLEANUP_PAUSE + .scope(pause, handler) + .await; + } + None => handler.await, + } + }) + } + }), + ); + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let address = listener.local_addr().expect("listener address"); + tokio::spawn(async move { axum::serve(listener, app).await.expect("serve") }); + (address, relay_url, served_state, conn_id) + } + + async fn connect(address: std::net::SocketAddr) -> Client { + tokio_tungstenite::connect_async(format!("ws://{address}/")) + .await + .expect("connect") + .0 + } + + fn expired_assertion() -> VerifiedAssertion { + let deadline = chrono::Utc::now() - chrono::Duration::seconds(1); + VerifiedAssertion::for_test(Some(Keys::generate().public_key()), vec![deadline]) + } + + /// Reads until the stream ends. Returns the text frames and the close + /// frame, or the error that ended the stream without one. + async fn read_to_close(client: &mut Client) -> (Vec, Result<(u16, String), String>) { + let mut texts = Vec::new(); + loop { + let frame = tokio::time::timeout(std::time::Duration::from_secs(5), client.next()) + .await + .expect("socket must terminate"); + match frame { + Some(Ok(Message::Text(text))) => texts.push(text.to_string()), + Some(Ok(Message::Close(Some(close)))) => { + return (texts, Ok((u16::from(close.code), close.reason.to_string()))) + } + Some(Ok(Message::Close(None))) => return (texts, Err("close without code".into())), + Some(Ok(_)) => {} + Some(Err(error)) => return (texts, Err(error.to_string())), + None => return (texts, Err("stream ended without close".into())), + } + } + } + + fn assert_denied_with_1008(texts: &[String], close: Result<(u16, String), String>) { + assert!( + texts + .iter() + .any(|t| t.contains("restricted: authorization denied")), + "denial NOTICE must arrive, got {texts:?}" + ); + assert_eq!( + close, + Ok(( + u16::from(CloseCode::Policy), + "authorization denied".to_owned() + )), + "denial must end with close 1008" + ); + } + + #[tokio::test] + async fn key_pairing_mismatch_closes_with_1008_on_real_socket() { + let asserted = Keys::generate(); + let assertion = VerifiedAssertion::for_test( + Some(asserted.public_key()), + vec![chrono::Utc::now() + chrono::Duration::hours(1)], + ); + let (address, relay_url, _, _) = serve(assertion, false).await; + let mut client = connect(address).await; + + let Some(Ok(Message::Text(challenge))) = client.next().await else { + panic!("expected AUTH challenge") + }; + let challenge: serde_json::Value = + serde_json::from_str(&challenge).expect("challenge json"); + let auth = EventBuilder::auth(challenge[1].as_str().expect("challenge"), relay_url) + .sign_with_keys(&Keys::generate()) + .expect("sign AUTH"); + client + .send(Message::Text( + serde_json::json!(["AUTH", auth]).to_string().into(), + )) + .await + .expect("send AUTH"); + + let (texts, close) = read_to_close(&mut client).await; + assert_denied_with_1008(&texts, close); + } + + #[tokio::test] + async fn session_expiry_denial_closes_with_1008_on_real_socket() { + let (address, _, _, _) = serve(expired_assertion(), true).await; + let mut client = connect(address).await; + + let (texts, close) = read_to_close(&mut client).await; + assert_denied_with_1008(&texts, close); + } + + /// Upgrades a raw TCP client and reads whole frames (server frames are + /// unmasked) through the server's Close, so a 0x88 byte inside a payload + /// or length cannot end the wait early. + pub(crate) async fn raw_client_after_server_close( + address: std::net::SocketAddr, + ) -> tokio::net::TcpStream { + raw_client_sending_then_after_server_close(address, &[]).await + } + + /// Like `raw_client_after_server_close`, but writes `frames` (masked + /// client frames) once the upgrade completes. + async fn raw_client_sending_then_after_server_close( + address: std::net::SocketAddr, + frames: &[u8], + ) -> tokio::net::TcpStream { + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + + let mut socket = tokio::net::TcpStream::connect(address) + .await + .expect("connect"); + socket + .write_all( + b"GET / HTTP/1.1\r\nHost: test.local\r\nUpgrade: websocket\r\n\ + Connection: Upgrade\r\nSec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\ + Sec-WebSocket-Version: 13\r\n\r\n", + ) + .await + .expect("send upgrade"); + + let mut read_exact = async |len: usize| { + let mut bytes = vec![0u8; len]; + tokio::time::timeout( + std::time::Duration::from_secs(5), + socket.read_exact(&mut bytes), + ) + .await + .expect("server frames must arrive") + .expect("socket ended before the server sent Close"); + bytes + }; + let mut response = Vec::new(); + while !response.ends_with(b"\r\n\r\n") { + response.extend(read_exact(1).await); + } + socket.write_all(frames).await.expect("send frames"); + let mut read_exact = async |len: usize| { + let mut bytes = vec![0u8; len]; + tokio::time::timeout( + std::time::Duration::from_secs(5), + socket.read_exact(&mut bytes), + ) + .await + .expect("server frames must arrive") + .expect("socket ended before the server sent Close"); + bytes + }; + loop { + let header = read_exact(2).await; + let payload_len = match header[1] & 0x7f { + 126 => u64::from(u16::from_be_bytes(read_exact(2).await.try_into().unwrap())), + 127 => u64::from_be_bytes(read_exact(8).await.try_into().unwrap()), + len => u64::from(len), + }; + read_exact(usize::try_from(payload_len).expect("frame length")).await; + if header[0] & 0x0f == 0x8 { + break; + } + } + socket + } + + /// Generous watchdog for a socket the server must end on its own. + const EOF_WATCHDOG: Duration = Duration::from_secs(10); + + /// Reads until EOF and returns the time elapsed since `started`. + pub(crate) async fn time_to_eof( + socket: &mut tokio::net::TcpStream, + started: std::time::Instant, + ) -> Duration { + use tokio::io::AsyncReadExt as _; + + let mut buf = [0u8; 64]; + tokio::time::timeout(EOF_WATCHDOG, async { + while socket.read(&mut buf).await.expect("read") != 0 {} + }) + .await + .expect("server must end the socket on its own"); + started.elapsed() + } + + /// Dropping TCP before the peer finishes the close handshake (including + /// in-flight peer frames) can make a client lose the NOTICE and Close + /// (seen live as `1006`). This raw client never replies: the server must + /// hold the socket for its full bounded wait, which starts only after + /// connect, so timing from before connect makes the lower bound + /// independent of scheduler delay. + #[tokio::test] + async fn server_holds_socket_for_client_close_reply() { + let (address, _, _, _) = serve(expired_assertion(), true).await; + let started = std::time::Instant::now(); + let mut socket = raw_client_after_server_close(address).await; + + let elapsed = time_to_eof(&mut socket, started).await; + assert!( + elapsed >= WS_TERMINAL_FLUSH_TIMEOUT, + "server must wait for the client's Close reply, ended after {elapsed:?}" + ); + } + + /// The client withholds its Close reply until the connection has + /// deregistered, so cleanup must not wait on the peer: a server that + /// cleans up after the wait, or ignores the reply, needs at + /// least `WS_TERMINAL_FLUSH_TIMEOUT` from connect. + #[tokio::test] + async fn server_finishes_promptly_when_client_answers_close() { + use tokio::io::AsyncWriteExt as _; + + let (address, _, state, conn_id) = serve(expired_assertion(), true).await; + let started = std::time::Instant::now(); + let mut socket = raw_client_after_server_close(address).await; + + tokio::time::timeout(EOF_WATCHDOG, async { + while state.conn_manager.community_for_conn(conn_id).is_some() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("connection must deregister before the client replies"); + // Masked Close, code 1000: payload [0x03, 0xe8] XOR a zero mask. + socket + .write_all(&[0x88, 0x82, 0, 0, 0, 0, 0x03, 0xe8]) + .await + .expect("send Close reply"); + let elapsed = time_to_eof(&mut socket, started).await; + assert!( + elapsed < WS_TERMINAL_FLUSH_TIMEOUT, + "server must finish on the client's Close, not the timeout; took {elapsed:?}" + ); + } + + /// Final cleanup can block on Redis with no deadline (the pool has no + /// wait timeout). The peer wait runs alongside it, so a stalled cleanup + /// still lets the socket end once the bounded wait expires. + #[tokio::test] + async fn stalled_cleanup_does_not_hold_socket_open() { + let pause = Arc::new(crate::handlers::close::test_seam::Pause::default()); + let (address, _, _, _) = serve_with_seams( + expired_assertion(), + true, + Seams { + cleanup_pause: Some(Arc::clone(&pause)), + ..Seams::default() + }, + ) + .await; + let started = std::time::Instant::now(); + let mut socket = raw_client_after_server_close(address).await; + + tokio::time::timeout(EOF_WATCHDOG, pause.reached.notified()) + .await + .expect("cleanup never reached the pause point"); + let elapsed = time_to_eof(&mut socket, started).await; + assert!(elapsed >= WS_TERMINAL_FLUSH_TIMEOUT); + } + + /// An established session's expiry task stalls at its quiescence barrier + /// while a handler holds an effect permit (e.g. an ephemeral EVENT blocked + /// on Redis). The writer has already sent the Close, so the socket must + /// still end within the bounded peer wait, not when the permit drops. + #[tokio::test] + async fn stalled_expiry_barrier_does_not_hold_socket_open() { + let (hold_tx, hold_rx) = tokio::sync::oneshot::channel(); + let assertion = VerifiedAssertion::for_test( + Some(Keys::generate().public_key()), + vec![chrono::Utc::now() + chrono::Duration::hours(1)], + ); + let (address, _, _, _) = serve_with_seams( + assertion, + true, + Seams { + effect_hold: Some(hold_tx), + ..Seams::default() + }, + ) + .await; + let mut socket = raw_client_after_server_close(address).await; + let _effect = hold_rx + .await + .expect("effect permit taken before the deadline"); + + let elapsed = time_to_eof(&mut socket, std::time::Instant::now()).await; + assert!( + elapsed < WS_TERMINAL_FLUSH_TIMEOUT + Duration::from_millis(500), + "a stalled expiry barrier must not hold the socket open, ended after {elapsed:?}" + ); + } + + /// A Nostr `CLOSE` is handled inline by the reader and can wait on the + /// subscription lock indefinitely (e.g. behind a REQ stuck in Redis). + /// Once the session expires the reader must give back its read half so + /// the socket ends within the bounded peer wait, and the interrupted + /// handler must still finish once the lock frees, before cleanup. + #[tokio::test] + async fn blocked_frame_handler_does_not_hold_socket_open() { + let (effect_tx, effect_rx) = tokio::sync::oneshot::channel(); + let (subs_tx, subs_rx) = tokio::sync::oneshot::channel(); + let assertion = VerifiedAssertion::for_test( + Some(Keys::generate().public_key()), + vec![chrono::Utc::now() + chrono::Duration::hours(1)], + ); + let pause = Arc::new(crate::handlers::close::test_seam::Pause::default()); + let (address, _, state, conn_id) = serve_with_seams( + assertion, + true, + Seams { + cleanup_pause: Some(Arc::clone(&pause)), + effect_hold: Some(effect_tx), + subscriptions_hold: Some(subs_tx), + }, + ) + .await; + // Masked text frame, zero mask: ["CLOSE","s"]. + let close = br#"["CLOSE","s"]"#; + let mut frame = vec![ + 0x81, + 0x80 | u8::try_from(close.len()).expect("len"), + 0, + 0, + 0, + 0, + ]; + frame.extend_from_slice(close); + let connected = tokio::spawn(async move { + raw_client_sending_then_after_server_close(address, &frame).await + }); + drop(effect_rx.await.expect("effect permit")); + let mut subscriptions = subs_rx.await.expect("subscription lock"); + subscriptions.insert("s".to_owned(), 1); + let mut socket = connected.await.expect("raw client"); + + let elapsed = time_to_eof(&mut socket, std::time::Instant::now()).await; + assert!( + elapsed < WS_TERMINAL_FLUSH_TIMEOUT + Duration::from_millis(500), + "a blocked frame handler must not hold the socket open, ended after {elapsed:?}" + ); + assert!( + state.conn_manager.community_for_conn(conn_id).is_some(), + "cleanup must wait for the interrupted handler" + ); + + let lock = Arc::clone(tokio::sync::OwnedMutexGuard::mutex(&subscriptions)); + assert!( + tokio::time::timeout(Duration::from_millis(200), pause.reached.notified()) + .await + .is_err(), + "cleanup must not start while the handed-off CLOSE is still blocked" + ); + drop(subscriptions); + tokio::time::timeout(EOF_WATCHDOG, pause.reached.notified()) + .await + .expect("cleanup never reached subscription release"); + assert!( + !lock.lock().await.contains_key("s"), + "the handed-off CLOSE must finish retiring before cleanup starts" + ); + pause.resume.notify_one(); + tokio::time::timeout(EOF_WATCHDOG, async { + while state.conn_manager.community_for_conn(conn_id).is_some() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("cleanup must complete once the handler finishes"); + assert!( + lock.lock().await.is_empty(), + "the subscription must be retired" + ); + } + + /// When the client closes first, tungstenite only queues our reply Close + /// and writes it on a later read. The peer wait must keep reading past + /// the client's Close so the reply goes out before TCP is dropped, as on + /// an audio exit that hands the socket off without sending a Close. + #[tokio::test] + async fn peer_wait_replies_to_client_initiated_close() { + let app = Router::new().route( + "/", + get(|ws: WebSocketUpgrade| async move { + ws.on_upgrade(|socket| async move { + let (_ws_send, ws_recv) = socket.split(); + await_client_close(ws_recv).await; + }) + }), + ); + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let address = listener.local_addr().expect("listener address"); + tokio::spawn(async move { axum::serve(listener, app).await.expect("serve") }); + let mut client = connect(address).await; + + client.close(None).await.expect("send Close"); + let reply = tokio::time::timeout(EOF_WATCHDOG, client.next()) + .await + .expect("server must answer or end the socket"); + assert!( + matches!(reply, Some(Ok(Message::Close(_)))), + "client-initiated Close must get its reply Close, got {reply:?}" + ); + } + + /// Serves the real outer `handle_connection` with `state`. + fn serve_outer( + state: Arc, + assertion: Option, + ) -> impl std::future::Future { + let tenant = TenantContext::resolved( + buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()), + "test.local".to_owned(), + ); + let app = Router::new().route( + "/", + get(move |ws: WebSocketUpgrade| { + let (state, tenant, assertion) = + (Arc::clone(&state), tenant.clone(), assertion.clone()); + async move { + ws.on_upgrade(move |socket| { + handle_connection( + socket, + state, + "127.0.0.1:1234".parse().expect("client address"), + tenant, + assertion, + None, + chrono::Utc::now(), + ) + }) + } + }), + ); + async move { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let address = listener.local_addr().expect("listener address"); + tokio::spawn(async move { axum::serve(listener, app).await.expect("serve") }); + address + } + } + + /// Leaves exactly one free slot in the connection budget. + async fn one_slot(state: &AppState) -> tokio::sync::OwnedSemaphorePermit { + let others = u32::try_from(state.conn_semaphore.available_permits() - 1).expect("permits"); + Arc::clone(&state.conn_semaphore) + .acquire_many_owned(others) + .await + .expect("reserve other slots") + } + + /// A socket refused during bootstrap (here an expired session) gets the + /// denial and Close from the `on_not_run` drain. The raw client never + /// replies, so the server must hold the socket for its bounded wait + /// instead of dropping it with client input possibly unread. + #[tokio::test] + async fn bootstrap_drain_holds_socket_for_client_close_reply() { + let state = crate::state::tests::test_state().await; + let address = serve_outer(state, Some(expired_assertion())).await; + + let started = std::time::Instant::now(); + let mut socket = raw_client_after_server_close(address).await; + let elapsed = time_to_eof(&mut socket, started).await; + assert!( + elapsed >= WS_TERMINAL_FLUSH_TIMEOUT, + "bootstrap drain must wait for the client's Close reply, ended after {elapsed:?}" + ); + } + + /// A socket refused during bootstrap keeps its connection slot until the + /// socket ends, including the bounded wait for a Close reply that never + /// comes, so refusals cannot open sockets beyond `BUZZ_MAX_CONNECTIONS`. + #[tokio::test] + async fn refused_bootstrap_holds_connection_slot_until_socket_ends() { + let state = crate::state::tests::test_state().await; + let _others = one_slot(&state).await; + let address = serve_outer(Arc::clone(&state), Some(expired_assertion())).await; + + let mut socket = raw_client_after_server_close(address).await; + assert_eq!( + state.conn_semaphore.available_permits(), + 0, + "the refused socket must hold its slot while open" + ); + time_to_eof(&mut socket, std::time::Instant::now()).await; + tokio::time::timeout(Duration::from_secs(1), async { + while state.conn_semaphore.available_permits() == 0 { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("the slot must be released once the socket ends"); + } + + /// The community check runs inside the connection budget: while it + /// stalls on the database, the socket already holds its slot even though + /// the active handler (which sends the AUTH challenge) has not started. + #[tokio::test] + async fn stalled_bootstrap_holds_connection_slot() { + let fake_database = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind database"); + let pool = sqlx::postgres::PgPoolOptions::new() + .acquire_timeout(Duration::from_secs(60)) + .connect_lazy(&format!( + "postgres://buzz@{}/buzz", + fake_database.local_addr().expect("database address") + )) + .expect("lazy pool"); + let state = crate::state::tests::test_state_with_database_pool(pool).await; + let _others = one_slot(&state).await; + let address = serve_outer(Arc::clone(&state), None).await; + + let mut client = connect(address).await; + let _stalled_check = tokio::time::timeout(Duration::from_secs(5), fake_database.accept()) + .await + .expect("the community check must reach the database") + .expect("accept database connection"); + assert_eq!( + state.conn_semaphore.available_permits(), + 0, + "a socket in bootstrap must hold its slot" + ); + assert!( + tokio::time::timeout(Duration::from_millis(200), client.next()) + .await + .is_err(), + "the active handler must not have started" + ); + } +} diff --git a/crates/buzz-relay/src/handlers/close.rs b/crates/buzz-relay/src/handlers/close.rs index d900bd4241..33b8b6020f 100644 --- a/crates/buzz-relay/src/handlers/close.rs +++ b/crates/buzz-relay/src/handlers/close.rs @@ -89,6 +89,8 @@ pub(crate) async fn close_if_owner( /// detached REQ task can register after this runs. pub(crate) async fn release_connection_subscriptions(conn: &ConnectionState, state: &AppState) { debug_assert!(conn.cancel.is_cancelled()); + #[cfg(test)] + test_seam::pause_at(&test_seam::CLEANUP_PAUSE).await; let mut subs = conn.subscriptions.lock().await; subs.clear(); for removed in state.sub_registry.remove_connection(conn.conn_id) { @@ -116,7 +118,8 @@ pub(crate) async fn release_scope_topics( /// Test-only pause points, scoped to one task so parallel tests don't /// interfere: `PAUSE` inside `close_if_owner` between teardown and the /// terminal frame, `RELEASE_PAUSE` inside `retire_locked` between the -/// registry removal and the topic release. +/// registry removal and the topic release, `CLEANUP_PAUSE` at the start of +/// `release_connection_subscriptions`. #[cfg(test)] pub(crate) mod test_seam { use std::sync::Arc; @@ -131,6 +134,7 @@ pub(crate) mod test_seam { tokio::task_local! { pub(crate) static PAUSE: Arc; pub(crate) static RELEASE_PAUSE: Arc; + pub(crate) static CLEANUP_PAUSE: Arc; } pub(super) async fn pause_at(point: &'static tokio::task::LocalKey>) { diff --git a/crates/buzz-relay/src/main.rs b/crates/buzz-relay/src/main.rs index effda6cc0a..a154f70757 100644 --- a/crates/buzz-relay/src/main.rs +++ b/crates/buzz-relay/src/main.rs @@ -2087,23 +2087,31 @@ async fn run_periodic_until_cancelled( /// grace, this backstops the whole drain and force-exits the process if /// exceeded. It bounds everything after the grace, not the grace itself. /// -/// A single WebSocket can therefore stay open, from SIGTERM, for up to: +/// Each socket's teardown after the grace runs in stages, each with its own +/// budget: /// -/// ```text -/// 5s grace + up to 20s jitter + up to 5s close-frame ack = 30s -/// (fixed) (MAX_DRAIN_JITTER_MS) (RESTART_CLOSE_ACK_TIMEOUT) -/// ``` +/// - up to 20s jitter before the relay sends its restart Close +/// ([`buzz_relay::config::MAX_DRAIN_JITTER_MS`]); +/// - up to 5s for the writer to acknowledge sending that Close +/// (`RESTART_CLOSE_ACK_TIMEOUT` in `state.rs`). If the acknowledgement +/// times out, the connection is cancelled and its teardown can spend a +/// separate 1s flushing the terminal frame (`WS_TERMINAL_FLUSH_TIMEOUT` in +/// `connection.rs`); +/// - up to 1s waiting for the client's Close reply (the same constant). This +/// wait starts only once the reader has returned and the writer has joined. /// -/// The 5s grace runs before the 30s hard-drain clock starts, so the jitter -/// (capped at [`buzz_relay::config::MAX_DRAIN_JITTER_MS`] = 20s) plus the -/// per-connection close-frame ack wait (`RESTART_CLOSE_ACK_TIMEOUT` = 5s in -/// `state.rs`) sum to 25s and stay inside the 30s hard drain. Total worst -/// case from SIGTERM to forced exit is 5s + 30s = 35s. Both fit inside the -/// chart's `terminationGracePeriodSeconds: 60` (`deploy/charts/buzz/values.yaml`), -/// which leaves headroom but assumes no `preStop` hook adds further delay. -/// With jitter off (`BUZZ_DRAIN_JITTER_MS=0`, the default) sockets close -/// all-at-once right after the grace, so the per-socket delay collapses to -/// roughly the 5s grace plus the ack wait. +/// These are stage budgets, not a guaranteed per-socket total: after a +/// successful acknowledgement nothing cancels the reader, so a silent peer +/// does not go straight into the peer wait. The 30s hard drain is what bounds +/// every socket. The configured process shutdown budget, from SIGTERM to +/// forced exit, is 5s + 30s = 35s, which fits within the chart's +/// `terminationGracePeriodSeconds: 60` (`deploy/charts/buzz/values.yaml`). +/// That leaves headroom but assumes no `preStop` hook adds further delay. +/// With jitter off (`BUZZ_DRAIN_JITTER_MS=0`, the default), `drain_all()` +/// queues every restart Close right after the grace and cancels the +/// connections immediately, without waiting for any acknowledgement. The +/// terminal-flush and peer-wait budgets then apply only to sockets whose +/// teardown reaches those stages. const GRACEFUL_DRAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); async fn serve( diff --git a/crates/buzz-relay/src/nip_fi_test_hooks.rs b/crates/buzz-relay/src/nip_fi_test_hooks.rs index 55d8e5b0e8..1769164ef7 100644 --- a/crates/buzz-relay/src/nip_fi_test_hooks.rs +++ b/crates/buzz-relay/src/nip_fi_test_hooks.rs @@ -210,6 +210,19 @@ make_hook!(audio_refusal_frames_hook, before_audio_refusal_frames); make_hook!(join_refusal_rollback_hook, before_join_refusal_rollback); make_hook!(audio_participant_fanout_hook, after_participant_fanout); make_hook!(audio_add_peer_hook, after_add_peer); +// `before_audio_teardown_cleanup`: fires in `run_active_audio_connection` once +// the writer has sent its Close and the socket is handed off, before the +// teardown that writes to Postgres and Redis. Holding it stands in for a +// stalled cleanup that must not hold the socket open. +make_hook!(audio_teardown_cleanup_hook, before_audio_teardown_cleanup); +// `before_room_archive_recheck`: fires in `run_active_audio_connection` after the +// membership check, just before the archive re-check, so a test can archive +// the channel between the two and drive the archived-channel exit. +make_hook!(room_archive_recheck_hook, before_room_archive_recheck); +// `before_archived_exit_cleanup`: fires on the archived-channel exit once the +// socket is handed off, before the lease release. Holding it stands in for a +// stalled Redis release that must not hold the socket open. +make_hook!(archived_exit_cleanup_hook, before_archived_exit_cleanup); // `before_lease_permit`: fires in `handle_active_audio_connection` immediately // before the effect permit that covers join-owner resolution. Cancelling here // proves the resolver (and its lease CAS) never runs without a permit. diff --git a/scripts/run-tests.sh b/scripts/run-tests.sh index bd7961124d..3850d288c5 100755 --- a/scripts/run-tests.sh +++ b/scripts/run-tests.sh @@ -343,6 +343,17 @@ run_unit_tests() { connection::tests::expiry_notice_queued_on_ctrl_before_cancel connection::tests::f3_root_outer_wrapper_delivers_denial_on_bootstrap_cancellation connection::tests::f3_root_pre_built_expired_gate_terminates_connection + connection::terminal_close_delivery_tests::key_pairing_mismatch_closes_with_1008_on_real_socket + connection::terminal_close_delivery_tests::session_expiry_denial_closes_with_1008_on_real_socket + connection::terminal_close_delivery_tests::server_holds_socket_for_client_close_reply + connection::terminal_close_delivery_tests::server_finishes_promptly_when_client_answers_close + connection::terminal_close_delivery_tests::stalled_cleanup_does_not_hold_socket_open + connection::terminal_close_delivery_tests::bootstrap_drain_holds_socket_for_client_close_reply + connection::terminal_close_delivery_tests::stalled_expiry_barrier_does_not_hold_socket_open + connection::terminal_close_delivery_tests::blocked_frame_handler_does_not_hold_socket_open + connection::terminal_close_delivery_tests::refused_bootstrap_holds_connection_slot_until_socket_ends + connection::terminal_close_delivery_tests::stalled_bootstrap_holds_connection_slot + connection::terminal_close_delivery_tests::peer_wait_replies_to_client_initiated_close handlers::auth::tests::b2_pre_cancelled_connection_never_becomes_authenticated handlers::auth::tests::fi_ban_check_error_emits_terminal_authorization_unavailable handlers::auth::tests::fi_root_authorization_denied_rows_emit_identical_frames