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