mirror of
https://github.com/block/buzz.git
synced 2026-11-08 02:11:36 +00:00
19874bb7bf
A sync REQ refused before its handler runs (handler saturation or WS admission) sent CLOSED and left the previous holder of the ID live, so fan-out could continue under a closed ID. Both paths now go through artifact_sync::refuse, which retires the old holder under the lifecycle lock before CLOSED. Co-authored-by: Larry <627498bd4bd1f281a16431e3c6cce3b5c25b6692798c78672298aefbf2f8f8b5@buzz.block.builderlab.xyz> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Logan Johnson <loganj@squareup.com>
4113 lines
166 KiB
Rust
4113 lines
166 KiB
Rust
//! WebSocket connection lifecycle: semaphore → challenge → recv/send/heartbeat loops → cleanup.
|
|
|
|
use std::collections::HashMap;
|
|
use std::net::SocketAddr;
|
|
use std::sync::atomic::{AtomicU8, Ordering};
|
|
use std::sync::{Arc, Mutex as StdMutex, PoisonError};
|
|
use std::time::{Duration, Instant};
|
|
|
|
use axum::extract::ws::{Message as WsMessage, WebSocket};
|
|
use futures_util::{Sink, SinkExt, StreamExt};
|
|
use tokio::sync::{mpsc, watch, Mutex};
|
|
use tokio_util::sync::CancellationToken;
|
|
use tracing::Instrument as _;
|
|
use tracing::{debug, info, trace, warn};
|
|
use uuid::Uuid;
|
|
|
|
use buzz_auth::{generate_challenge, AuthContext};
|
|
use buzz_core::tenant::TenantContext;
|
|
|
|
use crate::handlers;
|
|
use crate::metrics::AuthOutcome;
|
|
use crate::protocol::{ClientMessage, RelayMessage};
|
|
use crate::rejection::{enforce_ws_admission, request_rejection_message, RejectionTarget};
|
|
use crate::state::{
|
|
run_registered_community_connection, AppState, CommunityConnectionControl,
|
|
CommunityDisconnectReason,
|
|
};
|
|
|
|
/// Maximum time a new socket may hold a connection slot without completing NIP-42 auth.
|
|
const AUTH_TIMEOUT: Duration = Duration::from_secs(5);
|
|
|
|
/// 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 →
|
|
/// owner token of the request that currently holds it (see `close::next_owner`).
|
|
pub(crate) type ConnectionSubscriptions = Arc<Mutex<HashMap<String, u64>>>;
|
|
|
|
/// Pre-built NIP-FI session components created before the `is_community_active`
|
|
/// bootstrap await. Passed from the HTTP-layer wrapper into the active handler so
|
|
/// the session deadline is enforced from the true upgrade instant.
|
|
/// [FI-TRACE-LEASE-BOUND, Fix 3]
|
|
type PreBuiltNipFiBundle = (
|
|
Arc<crate::nip_fi_gate::SessionAdmissionGate>,
|
|
mpsc::Sender<WsMessage>,
|
|
mpsc::Receiver<WsMessage>,
|
|
Option<tokio::task::JoinHandle<()>>,
|
|
);
|
|
|
|
/// Request for the writer to flush a restart close and report the result.
|
|
pub(crate) struct RestartClose {
|
|
pub(crate) flushed: tokio::sync::oneshot::Sender<bool>,
|
|
}
|
|
|
|
/// Maximum outbound data frames buffered into the websocket sink before one flush.
|
|
const MAX_WS_SEND_BATCH: usize = 64;
|
|
|
|
/// NIP-42 authentication state for a single connection.
|
|
#[derive(Debug, Clone)]
|
|
pub enum AuthState {
|
|
/// Challenge has been sent; awaiting a signed AUTH event from the client.
|
|
Pending {
|
|
/// The random challenge string sent to the client.
|
|
challenge: String,
|
|
/// When the challenge was delivered and this attempt began.
|
|
started_at: Instant,
|
|
},
|
|
/// Client has successfully authenticated.
|
|
Authenticated(AuthContext),
|
|
/// Authentication attempt was rejected.
|
|
Failed,
|
|
}
|
|
|
|
/// Per-connection state split by access pattern:
|
|
/// - `auth_state`: synchronous mutex (short, non-awaiting transitions; drop-safe cleanup)
|
|
/// - `subscriptions`: Mutex (write-heavy during REQ/CLOSE)
|
|
/// - `send_tx`, `ctrl_tx`, `cancel`: outside any lock (Clone+Send, no coordination needed)
|
|
pub struct ConnectionState {
|
|
/// Unique identifier for this connection.
|
|
pub conn_id: Uuid,
|
|
/// The community this connection is bound to, resolved from the connection
|
|
/// host at row zero (before any frame is read) and never overridable by
|
|
/// client-supplied input. Every handler reads tenant scope from here.
|
|
pub tenant: TenantContext,
|
|
/// Remote socket address of the client.
|
|
pub remote_addr: SocketAddr,
|
|
/// Current NIP-42 authentication state.
|
|
pub auth_state: StdMutex<AuthState>,
|
|
/// Active subscriptions keyed by subscription ID.
|
|
pub subscriptions: ConnectionSubscriptions,
|
|
/// Sender for outbound data messages (EVENT, NOTICE, OK, etc.).
|
|
pub send_tx: mpsc::Sender<WsMessage>,
|
|
/// Sender for outbound control frames (Pong, Close).
|
|
/// Separate channel with priority drain — if this channel fills too,
|
|
/// the connection is closed (writer is completely stalled).
|
|
pub ctrl_tx: mpsc::Sender<WsMessage>,
|
|
/// Dedicated one-slot sender for the terminal NIP-FI denial frame.
|
|
///
|
|
/// Because only one terminal event fires per connection lifetime (either key
|
|
/// pairing mismatch or session expiry, never both), this channel is always
|
|
/// available when the denial is enqueued — it cannot be saturated by ordinary
|
|
/// control traffic. The send_loop drains it in its cancel branch ahead of
|
|
/// `Close`, guaranteeing the denial frame is delivered even when `ctrl_tx`
|
|
/// (capacity 8) is full. [FI-INV-05, FI-TRACE-LEASE-BOUND]
|
|
pub terminal_ctrl_tx: mpsc::Sender<WsMessage>,
|
|
/// Token used to signal graceful shutdown of this connection's tasks.
|
|
pub cancel: CancellationToken,
|
|
/// Consecutive buffer-full events. Cancel only after `grace_limit`.
|
|
/// Shared with `ConnectionManager::ConnEntry` so both direct sends and
|
|
/// fan-out broadcasts track the same counter.
|
|
pub backpressure_count: Arc<AtomicU8>,
|
|
/// Configurable slow-client grace limit (from `Config::slow_client_grace_limit`).
|
|
pub grace_limit: u8,
|
|
/// Lifecycle control shared with the community registry; the auth
|
|
/// handler's post-registration deny check routes its denial through it.
|
|
/// [FI-TRACE-DENY-SET]
|
|
pub(crate) community_control: crate::state::CommunityConnectionControl,
|
|
|
|
/// The NIP-FI assertion presented at upgrade, when enforcement is enabled.
|
|
///
|
|
/// `None` means the relay is in `Off` mode — no assertion is required.
|
|
/// When `Some`, the NIP-42 key pairing check uses this to enforce that
|
|
/// `assertion.asserted_key() == nip42_pubkey` unconditionally (S3 invariant:
|
|
/// no flag reads — S2 deleted `require_attested_key`). [FI-INV-05]
|
|
pub nip_fi_assertion: Option<buzz_auth::VerifiedAssertion>,
|
|
|
|
/// The UTC deadline after which this connection's NIP-FI lease expires.
|
|
///
|
|
/// `None` means no assertion-based lifetime is enforced (mode is `Off`).
|
|
/// When `Some`, the session-expiry task fires at this instant and sends
|
|
/// `restricted: authorization denied` + cancels. Equality is expired.
|
|
/// [FI-TRACE-LEASE-BOUND]
|
|
pub session_deadline: Option<chrono::DateTime<chrono::Utc>>,
|
|
|
|
/// The NIP-FI session admission gate. Every WS connection has exactly one
|
|
/// gate — this is the [one-gate-per-connection] invariant.
|
|
///
|
|
/// In enforce mode (assertion presented at upgrade), the gate has a
|
|
/// deadline and the expiry task calls `gate.expire()` at that deadline.
|
|
/// In off-mode (no assertion), the gate has no deadline and never
|
|
/// self-expires — `acquire_effect()` always succeeds unless the outer
|
|
/// cancel token fires.
|
|
///
|
|
/// Handlers that perform irreversible side effects (AUTH state commit,
|
|
/// EVENT persistence, REQ subscription registration, COUNT query) must
|
|
/// call `gate.acquire_effect()` at the irreversible seam. The gate's
|
|
/// quiescence barrier ensures connection teardown (subscription removal,
|
|
/// peer cleanup) cannot start until all pre-expiry effects finish their
|
|
/// bounded commits. [FI-TRACE-LEASE-BOUND, one-gate-per-connection]
|
|
pub(crate) nip_fi_gate: std::sync::Arc<crate::nip_fi_gate::SessionAdmissionGate>,
|
|
}
|
|
|
|
impl ConnectionState {
|
|
fn lock_auth_state(&self) -> std::sync::MutexGuard<'_, AuthState> {
|
|
self.auth_state
|
|
.lock()
|
|
.unwrap_or_else(PoisonError::into_inner)
|
|
}
|
|
|
|
/// Snapshot the current authentication state without holding its lock over an await.
|
|
pub(crate) fn auth_state_snapshot(&self) -> AuthState {
|
|
self.lock_auth_state().clone()
|
|
}
|
|
|
|
fn transition_pending_auth(&self, next: AuthState, outcome: AuthOutcome) -> bool {
|
|
let mut auth = self.lock_auth_state();
|
|
let AuthState::Pending { started_at, .. } = &*auth else {
|
|
return false;
|
|
};
|
|
let duration = started_at.elapsed();
|
|
let became_authenticated = matches!(next, AuthState::Authenticated(_));
|
|
*auth = next;
|
|
crate::metrics::record_auth_outcome(outcome, duration);
|
|
if became_authenticated {
|
|
// Keep the gauge update under the same state lock as the
|
|
// Pending -> Authenticated transition. Cleanup must never observe
|
|
// Authenticated before its increment and decrement first.
|
|
metrics::gauge!("buzz_ws_authenticated_connections_active").increment(1.0);
|
|
}
|
|
true
|
|
}
|
|
|
|
/// Atomically finish the initial challenge as authenticated.
|
|
pub(crate) fn authenticate(&self, auth_context: AuthContext) -> bool {
|
|
self.transition_pending_auth(AuthState::Authenticated(auth_context), AuthOutcome::Success)
|
|
}
|
|
|
|
/// Atomically finish the initial challenge with a bounded denial.
|
|
pub(crate) fn reject_auth(&self, outcome: AuthOutcome) -> bool {
|
|
debug_assert!(!matches!(outcome, AuthOutcome::Success));
|
|
self.transition_pending_auth(AuthState::Failed, outcome)
|
|
}
|
|
|
|
/// Finish a pending challenge on timeout and preserve the historical rule
|
|
/// that an already-failed connection is closed when its timeout expires.
|
|
fn expire_auth(&self) -> bool {
|
|
let mut auth = self.lock_auth_state();
|
|
match &*auth {
|
|
AuthState::Pending { started_at, .. } => {
|
|
let duration = started_at.elapsed();
|
|
*auth = AuthState::Failed;
|
|
crate::metrics::record_auth_outcome(AuthOutcome::Timeout, duration);
|
|
true
|
|
}
|
|
AuthState::Failed => true,
|
|
AuthState::Authenticated(_) => false,
|
|
}
|
|
}
|
|
|
|
/// Finalize authentication accounting when a connection closes.
|
|
///
|
|
/// Replacing the state with `Failed` makes cleanup idempotent: an
|
|
/// authenticated gauge can be decremented at most once, and a pending
|
|
/// attempt can receive at most one disconnect/shutdown terminal.
|
|
fn finish_auth_on_close(&self, outcome: AuthOutcome) -> Option<AuthContext> {
|
|
debug_assert!(matches!(
|
|
outcome,
|
|
AuthOutcome::Disconnect | AuthOutcome::Shutdown
|
|
));
|
|
let mut auth = self.lock_auth_state();
|
|
match std::mem::replace(&mut *auth, AuthState::Failed) {
|
|
AuthState::Pending { started_at, .. } => {
|
|
crate::metrics::record_auth_outcome(outcome, started_at.elapsed());
|
|
None
|
|
}
|
|
AuthState::Authenticated(auth_context) => {
|
|
metrics::gauge!("buzz_ws_authenticated_connections_active").decrement(1.0);
|
|
Some(auth_context)
|
|
}
|
|
AuthState::Failed => None,
|
|
}
|
|
}
|
|
|
|
/// Let the cancellation watcher claim only a still-pending attempt.
|
|
/// Authenticated cleanup remains owned by `AuthLifecycleGuard`, which must
|
|
/// retain the authentication context for presence cleanup after task joins.
|
|
fn finish_pending_auth_on_cancel(&self, outcome: AuthOutcome) -> bool {
|
|
debug_assert!(matches!(
|
|
outcome,
|
|
AuthOutcome::Disconnect | AuthOutcome::Shutdown
|
|
));
|
|
let mut auth = self.lock_auth_state();
|
|
let AuthState::Pending { started_at, .. } = &*auth else {
|
|
return false;
|
|
};
|
|
let duration = started_at.elapsed();
|
|
*auth = AuthState::Failed;
|
|
crate::metrics::record_auth_outcome(outcome, duration);
|
|
true
|
|
}
|
|
|
|
/// Sends a data message to this connection's outbound channel.
|
|
///
|
|
/// On a full buffer, increments the backpressure counter. The first
|
|
/// `grace_limit` occurrences log a warning; sustained backpressure
|
|
/// cancels the connection to prevent unbounded memory growth.
|
|
pub fn send(&self, msg: String) -> bool {
|
|
match self.send_tx.try_send(WsMessage::Text(msg.into())) {
|
|
Ok(_) => {
|
|
// Successful send resets the grace counter.
|
|
self.backpressure_count.store(0, Ordering::Relaxed);
|
|
true
|
|
}
|
|
Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
|
|
let count = self.backpressure_count.fetch_add(1, Ordering::Relaxed) + 1;
|
|
if count >= self.grace_limit {
|
|
warn!(conn_id = %self.conn_id, count, "sustained backpressure — closing slow client");
|
|
metrics::counter!("buzz_ws_backpressure_disconnects_total").increment(1);
|
|
self.cancel.cancel();
|
|
} else {
|
|
warn!(conn_id = %self.conn_id, count, grace = self.grace_limit, "send buffer full — grace {count}/{}", self.grace_limit);
|
|
}
|
|
false
|
|
}
|
|
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
|
|
debug!(conn_id = %self.conn_id, "send channel closed");
|
|
false
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Owns authentication accounting for exactly the lifetime of the production
|
|
/// connection future. Explicit teardown selects the precise close outcome;
|
|
/// aborts and panics fall back to `disconnect` and cancel child tasks.
|
|
struct AuthLifecycleGuard {
|
|
conn: Arc<ConnectionState>,
|
|
finished: bool,
|
|
}
|
|
|
|
impl AuthLifecycleGuard {
|
|
fn new(conn: Arc<ConnectionState>) -> Self {
|
|
Self {
|
|
conn,
|
|
finished: false,
|
|
}
|
|
}
|
|
|
|
fn finish(&mut self, outcome: AuthOutcome) -> Option<AuthContext> {
|
|
self.finished = true;
|
|
self.conn.finish_auth_on_close(outcome)
|
|
}
|
|
}
|
|
|
|
impl Drop for AuthLifecycleGuard {
|
|
fn drop(&mut self) {
|
|
if !self.finished {
|
|
self.conn.cancel.cancel();
|
|
self.conn.finish_auth_on_close(AuthOutcome::Disconnect);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Compute the NIP-FI session deadline from a verified assertion and the
|
|
/// configured `max_connection_lifetime`.
|
|
///
|
|
/// Per spec [FI-TRACE-LEASE-BOUND]:
|
|
/// ```text
|
|
/// session_deadline = min(
|
|
/// assertion.upstream_authority_deadline(), // min(exp, iat+max_age, key-snapshot-hard)
|
|
/// connection_time + max_connection_lifetime // partitions, never shortens
|
|
/// )
|
|
/// ```
|
|
///
|
|
/// `upstream_authority_deadline()` already includes the key-snapshot hard
|
|
/// deadline (one of the three authority_deadlines terms), so this two-term min
|
|
/// covers all four normative terms. Equality at any deadline is expired.
|
|
///
|
|
/// `connection_time` must be captured at or immediately before the WebSocket
|
|
/// upgrade — not after the NIP-42 exchange — so the partition is rooted at the
|
|
/// true connection establishment instant and the session cannot outlive
|
|
/// `connection_time + max_connection_lifetime` by the authentication interval.
|
|
pub(crate) fn compute_session_deadline(
|
|
assertion: &buzz_auth::VerifiedAssertion,
|
|
connection_time: chrono::DateTime<chrono::Utc>,
|
|
max_connection_lifetime: Option<std::time::Duration>,
|
|
) -> chrono::DateTime<chrono::Utc> {
|
|
let upstream = assertion.upstream_authority_deadline();
|
|
match max_connection_lifetime {
|
|
Some(lifetime) => {
|
|
let partition = match chrono::Duration::from_std(lifetime) {
|
|
Ok(d) => connection_time + d,
|
|
// lifetime so large it overflows chrono — treat as effectively
|
|
// infinite, so the upstream deadline wins.
|
|
Err(_) => chrono::DateTime::<chrono::Utc>::MAX_UTC,
|
|
};
|
|
upstream.min(partition)
|
|
}
|
|
None => upstream,
|
|
}
|
|
}
|
|
|
|
/// Acquires a connection semaphore permit, sends the NIP-42 AUTH challenge,
|
|
/// then drives the send, heartbeat, and receive loops until the connection closes.
|
|
pub(crate) async fn handle_connection(
|
|
socket: WebSocket,
|
|
state: Arc<AppState>,
|
|
addr: SocketAddr,
|
|
tenant: TenantContext,
|
|
nip_fi_assertion: Option<buzz_auth::VerifiedAssertion>,
|
|
nip_fi_shadow: Option<Arc<crate::nip_fi_shadow_session::ShadowSession>>,
|
|
connection_time: chrono::DateTime<chrono::Utc>,
|
|
) {
|
|
// 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();
|
|
|
|
// Fix 3: Arm the NIP-FI gate and expiry task BEFORE the `is_community_active`
|
|
// bootstrap await so the session deadline is enforced even when the DB check
|
|
// is delayed. The deadline is computed from `connection_time` and
|
|
// `nip_fi_assertion` — both are available here, before bootstrap.
|
|
// [FI-TRACE-LEASE-BOUND, NIP-FI §"terminated no later than"]
|
|
let pre_session_deadline = nip_fi_assertion.as_ref().map(|a| {
|
|
compute_session_deadline(
|
|
a,
|
|
connection_time,
|
|
state.config.nip_fi.max_connection_lifetime(),
|
|
)
|
|
});
|
|
let pre_gate = if let Some(deadline) = pre_session_deadline {
|
|
crate::nip_fi_gate::SessionAdmissionGate::new(deadline, cancel.clone())
|
|
} else {
|
|
crate::nip_fi_gate::SessionAdmissionGate::off_mode(cancel.clone())
|
|
};
|
|
let (pre_terminal_ctrl_tx, pre_terminal_ctrl_rx) = mpsc::channel::<WsMessage>(1);
|
|
let control = CommunityConnectionControl::new(cancel);
|
|
control.attach_nip_fi_shadow(nip_fi_shadow);
|
|
let drain_reason = control.disconnect_reason();
|
|
let pre_expiry_task = pre_session_deadline.map(|deadline| {
|
|
crate::nip_fi_session::spawn_nip_fi_expiry_task(
|
|
deadline,
|
|
Arc::clone(&pre_gate),
|
|
control.clone(),
|
|
pre_terminal_ctrl_tx.clone(),
|
|
crate::nip_fi_session::NipFiWsRoute::Root,
|
|
)
|
|
});
|
|
|
|
let community_id = tenant.community();
|
|
let registry = Arc::clone(&state.community_connections);
|
|
let check_state = Arc::clone(&state);
|
|
let run_state = Arc::clone(&state);
|
|
|
|
// Fix 3 (F3 drain): when run is not called (cancellation or inactive community),
|
|
// drain the pre-terminal channel and close the socket so a queued NIP-FI denial
|
|
// is delivered before the connection drops. Both the run closure and the drain
|
|
// closure need the socket and receiver — use a shared Arc<Mutex<Option<...>>> so
|
|
// exactly one path takes each value. [FI-TRACE-BOOTSTRAP-DENIAL-DRAIN]
|
|
let socket_shared = Arc::new(Mutex::new(Some(socket)));
|
|
let socket_for_run = Arc::clone(&socket_shared);
|
|
let socket_for_drain = Arc::clone(&socket_shared);
|
|
|
|
// Similarly share the pre-terminal receiver: the run path passes it into the
|
|
// active handler (which drains it via the writer); the drain path drains it
|
|
// directly. The Arc ensures each path can move-capture the shared state.
|
|
let rx_shared = Arc::new(Mutex::new(Some(pre_terminal_ctrl_rx)));
|
|
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,
|
|
community_id,
|
|
control,
|
|
move || async move { check_state.db.is_community_active(community_id).await },
|
|
move |control| async move {
|
|
let socket = socket_for_run
|
|
.lock()
|
|
.await
|
|
.take()
|
|
.expect("socket taken by drain before run — logic error");
|
|
let pre_terminal_ctrl_rx = rx_for_run
|
|
.lock()
|
|
.await
|
|
.take()
|
|
.expect("rx taken by drain before run — logic error");
|
|
handle_active_connection(
|
|
socket,
|
|
run_state,
|
|
addr,
|
|
tenant,
|
|
conn_id,
|
|
control,
|
|
nip_fi_assertion,
|
|
connection_time,
|
|
Some((
|
|
pre_gate,
|
|
pre_terminal_ctrl_tx,
|
|
pre_terminal_ctrl_rx,
|
|
pre_expiry_task,
|
|
)),
|
|
)
|
|
.await
|
|
},
|
|
move || async move {
|
|
// 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 queued = rx_for_drain.lock().await.take();
|
|
if let Some(socket) = socket {
|
|
let close = drain_reason.borrow().map(|reason| reason.close_message());
|
|
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
|
|
// surface (socket, state, addr, tenant, conn_id, control, assertion, connection_time).
|
|
// Collapsing into a struct would just move the fields without reducing coupling.
|
|
#[allow(clippy::too_many_arguments)]
|
|
async fn handle_active_connection(
|
|
socket: WebSocket,
|
|
state: Arc<AppState>,
|
|
addr: SocketAddr,
|
|
tenant: TenantContext,
|
|
conn_id: Uuid,
|
|
control: CommunityConnectionControl,
|
|
nip_fi_assertion: Option<buzz_auth::VerifiedAssertion>,
|
|
connection_time: chrono::DateTime<chrono::Utc>,
|
|
// Fix 3: pre-built gate/expiry/channels from `handle_connection`, armed
|
|
// BEFORE the `is_community_active` bootstrap await. `None` is not currently
|
|
// produced by any caller but retained for future test extensibility.
|
|
pre_built: Option<PreBuiltNipFiBundle>,
|
|
) {
|
|
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 challenge = generate_challenge();
|
|
|
|
let (tx, rx) = mpsc::channel::<WsMessage>(state.config.send_buffer_size);
|
|
// Control channel for Pong/Close — small capacity, guaranteed delivery
|
|
// even when the data buffer is full.
|
|
let (ctrl_tx, ctrl_rx) = mpsc::channel::<WsMessage>(8);
|
|
|
|
// Dedicated one-slot channel for the terminal NIP-FI denial frame.
|
|
// Cannot be saturated by ordinary traffic — only one terminal event fires.
|
|
// Fix 3: Use the pre-built terminal channel from `handle_connection` (armed
|
|
// before bootstrap) if available; otherwise create fresh here.
|
|
let (terminal_ctrl_tx, terminal_ctrl_rx, pre_nip_fi_gate, pre_nip_fi_expiry_task) =
|
|
if let Some((gate, tx, rx, expiry)) = pre_built {
|
|
(tx, rx, Some(gate), expiry)
|
|
} else {
|
|
let (tx, rx) = mpsc::channel::<WsMessage>(1);
|
|
(tx, rx, None, None)
|
|
};
|
|
|
|
// Dedicated restart-close channel carries a flush acknowledgement. Keeping
|
|
// ordinary control frames unchanged avoids coupling heartbeat/ban traffic
|
|
// to graceful-shutdown delivery tracking.
|
|
let (restart_tx, restart_rx) = mpsc::channel::<RestartClose>(1);
|
|
|
|
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.
|
|
//
|
|
// Per spec (Request and session bounds, [FI-TRACE-LEASE-BOUND]):
|
|
// session_deadline = min(
|
|
// assertion.upstream_authority_deadline(), // = min(exp, iat+max_age, key-snapshot hard deadline)
|
|
// connection_time + max_connection_lifetime // partitions, never shortens per spec
|
|
// )
|
|
//
|
|
// Equality at any deadline is expired. `upstream_authority_deadline()` already
|
|
// includes the key-snapshot hard deadline (one of the three authority_deadlines
|
|
// terms), so this min covers all normative terms.
|
|
let session_deadline = nip_fi_assertion.as_ref().map(|a| {
|
|
compute_session_deadline(
|
|
a,
|
|
connection_time,
|
|
state.config.nip_fi.max_connection_lifetime(),
|
|
)
|
|
});
|
|
|
|
// Create the NIP-FI session admission gate. Every WS connection gets
|
|
// exactly one gate — the [one-gate-per-connection] invariant.
|
|
//
|
|
// The gate is the lifetime authority for this connection: handlers acquire
|
|
// an effect permit before any DB write, and the expiry task closes the gate
|
|
// at the session deadline so no new effects can start after expiry.
|
|
//
|
|
// Enforce mode (assertion + deadline): gate has a deadline; the expiry
|
|
// task calls gate.expire() at the deadline.
|
|
// Off-mode (no assertion): gate has no deadline and never self-expires;
|
|
// acquire_effect() always succeeds unless the outer cancel token fires.
|
|
// [FI-TRACE-LEASE-BOUND, one-gate-per-connection]
|
|
//
|
|
// Fix 3: when pre_built, use the already-armed pre_nip_fi_gate (created
|
|
// before bootstrap in `handle_connection`) so the deadline is enforced
|
|
// across the full session lifecycle.
|
|
let nip_fi_gate = pre_nip_fi_gate.unwrap_or_else(|| {
|
|
if let Some(deadline) = session_deadline {
|
|
crate::nip_fi_gate::SessionAdmissionGate::new(deadline, cancel.clone())
|
|
} else {
|
|
crate::nip_fi_gate::SessionAdmissionGate::off_mode(cancel.clone())
|
|
}
|
|
});
|
|
|
|
let conn = Arc::new(ConnectionState {
|
|
conn_id,
|
|
tenant,
|
|
remote_addr: addr,
|
|
auth_state: StdMutex::new(AuthState::Pending {
|
|
challenge: challenge.clone(),
|
|
started_at: Instant::now(),
|
|
}),
|
|
subscriptions: Arc::clone(&subscriptions),
|
|
send_tx: tx.clone(),
|
|
ctrl_tx: ctrl_tx.clone(),
|
|
terminal_ctrl_tx,
|
|
cancel: cancel.clone(),
|
|
backpressure_count: Arc::clone(&backpressure_count),
|
|
grace_limit: state.config.slow_client_grace_limit,
|
|
nip_fi_assertion,
|
|
session_deadline,
|
|
nip_fi_gate: nip_fi_gate.clone(),
|
|
community_control: control.clone(),
|
|
});
|
|
|
|
info!(conn_id = %conn_id, addr = %addr, "WebSocket connection established");
|
|
metrics::counter!(
|
|
"buzz_ws_connections_total",
|
|
"community" => conn.tenant.host().to_owned()
|
|
)
|
|
.increment(1);
|
|
|
|
let challenge_msg = RelayMessage::auth_challenge(&challenge);
|
|
if tx
|
|
.send(WsMessage::Text(challenge_msg.into()))
|
|
.await
|
|
.is_err()
|
|
{
|
|
warn!(conn_id = %conn_id, "Failed to send AUTH challenge — client disconnected immediately");
|
|
return;
|
|
}
|
|
|
|
// Gauge incremented AFTER challenge send succeeds — early disconnects
|
|
// don't leak. Decremented in the cleanup path below.
|
|
metrics::gauge!("buzz_ws_connections_active").increment(1.0);
|
|
crate::metrics::record_auth_attempt_started();
|
|
let mut auth_lifecycle = AuthLifecycleGuard::new(Arc::clone(&conn));
|
|
|
|
// Register after challenge succeeds — avoids leaked entries on early disconnect.
|
|
state.conn_manager.register(
|
|
conn_id,
|
|
tx.clone(),
|
|
ctrl_tx.clone(),
|
|
conn.terminal_ctrl_tx.clone(),
|
|
Some(restart_tx),
|
|
cancel.clone(),
|
|
conn.tenant.community(),
|
|
Arc::clone(&backpressure_count),
|
|
subscriptions,
|
|
state.config.slow_client_grace_limit,
|
|
control.clone(),
|
|
);
|
|
|
|
let (ws_send, ws_recv) = socket.split();
|
|
|
|
let send_cancel = cancel.child_token();
|
|
let send_task = tokio::spawn(send_loop(
|
|
ws_send,
|
|
rx,
|
|
ctrl_rx,
|
|
terminal_ctrl_rx,
|
|
restart_rx,
|
|
send_cancel,
|
|
disconnect_reason,
|
|
));
|
|
|
|
let missed_pongs = Arc::new(AtomicU8::new(0));
|
|
let heartbeat_cancel = cancel.clone();
|
|
let heartbeat_task = tokio::spawn(heartbeat_loop(
|
|
ctrl_tx,
|
|
Arc::clone(&missed_pongs),
|
|
heartbeat_cancel,
|
|
));
|
|
|
|
let auth_timeout_conn = Arc::clone(&conn);
|
|
let auth_timeout_cancel = cancel.clone();
|
|
let auth_timeout_task = tokio::spawn(async move {
|
|
tokio::select! {
|
|
_ = tokio::time::sleep(AUTH_TIMEOUT) => {
|
|
if auth_timeout_conn.expire_auth() {
|
|
warn!(
|
|
conn_id = %auth_timeout_conn.conn_id,
|
|
timeout_secs = AUTH_TIMEOUT.as_secs(),
|
|
"NIP-42 auth timeout — closing connection"
|
|
);
|
|
metrics::counter!("buzz_ws_auth_timeouts_total").increment(1);
|
|
auth_timeout_cancel.cancel();
|
|
}
|
|
}
|
|
_ = auth_timeout_cancel.cancelled() => {}
|
|
}
|
|
});
|
|
|
|
// NIP-FI session-lifetime enforcement task.
|
|
//
|
|
// Uses gate.expire() so the quiescence barrier (write lock) ensures
|
|
// connection teardown cannot start until all pre-expiry effects have
|
|
// finished. [FI-TRACE-LEASE-BOUND]
|
|
//
|
|
// Fix 3: if the expiry task was already spawned pre-bootstrap
|
|
// (`pre_nip_fi_expiry_task`), use it directly; otherwise spawn fresh.
|
|
let nip_fi_expiry_task = pre_nip_fi_expiry_task.or_else(|| {
|
|
conn.session_deadline.map(|deadline| {
|
|
crate::nip_fi_session::spawn_nip_fi_expiry_task(
|
|
deadline,
|
|
Arc::clone(&nip_fi_gate),
|
|
conn.community_control.clone(),
|
|
conn.terminal_ctrl_tx.clone(),
|
|
crate::nip_fi_session::NipFiWsRoute::Root,
|
|
)
|
|
})
|
|
});
|
|
|
|
// Cancellation races database-backed AUTH work. This watcher claims the
|
|
// pending lifecycle under the same lock as success/denial transitions, so
|
|
// whichever terminal happens first wins and a late handler cannot overwrite it.
|
|
let auth_cancel_conn = Arc::clone(&conn);
|
|
let auth_cancel_state = Arc::clone(&state);
|
|
let auth_cancel_token = cancel.clone();
|
|
let auth_cancel_task = tokio::spawn(async move {
|
|
auth_cancel_token.cancelled().await;
|
|
let outcome = if auth_cancel_state.shutting_down.load(Ordering::Acquire) {
|
|
AuthOutcome::Shutdown
|
|
} else {
|
|
AuthOutcome::Disconnect
|
|
};
|
|
auth_cancel_conn.finish_pending_auth_on_cancel(outcome);
|
|
});
|
|
let (ws_recv, in_flight) = recv_loop(
|
|
ws_recv,
|
|
Arc::clone(&conn),
|
|
Arc::clone(&state),
|
|
Arc::clone(&missed_pongs),
|
|
cancel.clone(),
|
|
)
|
|
.await;
|
|
|
|
cancel.cancel();
|
|
let close_outcome = if state.shutting_down.load(Ordering::Acquire) {
|
|
AuthOutcome::Shutdown
|
|
} else {
|
|
AuthOutcome::Disconnect
|
|
};
|
|
// Terminalize before joining writer/heartbeat tasks. A blocked socket sink
|
|
// must not keep the authenticated gauge high during shutdown.
|
|
let authenticated = auth_lifecycle.finish(close_outcome);
|
|
|
|
let _ = send_task.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");
|
|
}
|
|
|
|
/// Send WebSocket messages in priority order: control frames before data frames.
|
|
///
|
|
/// Control frames (Pong, Close) are drained first on every iteration,
|
|
/// giving them priority over data frames. If the underlying socket writer
|
|
/// is stalled, control frames queue in the small ctrl_rx buffer; callers
|
|
/// treat a full control channel as terminal (Bug 7 fix).
|
|
async fn send_loop(
|
|
ws_send: futures_util::stream::SplitSink<WebSocket, WsMessage>,
|
|
data_rx: mpsc::Receiver<WsMessage>,
|
|
ctrl_rx: mpsc::Receiver<WsMessage>,
|
|
terminal_ctrl_rx: mpsc::Receiver<WsMessage>,
|
|
restart_rx: mpsc::Receiver<RestartClose>,
|
|
cancel: CancellationToken,
|
|
disconnect_reason: watch::Receiver<Option<CommunityDisconnectReason>>,
|
|
) {
|
|
send_loop_inner(
|
|
ws_send,
|
|
data_rx,
|
|
ctrl_rx,
|
|
terminal_ctrl_rx,
|
|
restart_rx,
|
|
cancel,
|
|
disconnect_reason,
|
|
)
|
|
.await;
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
|
enum WriterStep {
|
|
Completed,
|
|
Cancelled,
|
|
Failed,
|
|
}
|
|
|
|
async fn send_or_cancel<S>(
|
|
sink: &mut S,
|
|
message: WsMessage,
|
|
cancel: &CancellationToken,
|
|
) -> WriterStep
|
|
where
|
|
S: Sink<WsMessage> + Unpin,
|
|
{
|
|
tokio::select! {
|
|
biased;
|
|
_ = cancel.cancelled() => WriterStep::Cancelled,
|
|
result = sink.send(message) => {
|
|
if result.is_ok() { WriterStep::Completed } else { WriterStep::Failed }
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn feed_or_cancel<S>(
|
|
sink: &mut S,
|
|
message: WsMessage,
|
|
cancel: &CancellationToken,
|
|
) -> WriterStep
|
|
where
|
|
S: Sink<WsMessage> + Unpin,
|
|
{
|
|
tokio::select! {
|
|
biased;
|
|
_ = cancel.cancelled() => WriterStep::Cancelled,
|
|
result = sink.feed(message) => {
|
|
if result.is_ok() { WriterStep::Completed } else { WriterStep::Failed }
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn flush_or_cancel<S>(sink: &mut S, cancel: &CancellationToken) -> WriterStep
|
|
where
|
|
S: Sink<WsMessage> + Unpin,
|
|
{
|
|
tokio::select! {
|
|
biased;
|
|
_ = cancel.cancelled() => WriterStep::Cancelled,
|
|
result = sink.flush() => {
|
|
if result.is_ok() { WriterStep::Completed } else { WriterStep::Failed }
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Best-effort terminal delivery with one shared deadline. A socket that never
|
|
/// becomes writable cannot retain its connection task or semaphore permit.
|
|
///
|
|
/// Drain order: FI terminal frames first (highest priority — denial must reach
|
|
/// the client before the close frame), then ordinary control frames, then the
|
|
/// Close frame. All sends are bounded by the shared `deadline` so a never-ready
|
|
/// sink cannot block indefinitely. [FI-TRACE-TERMINAL-BOUNDED, Fix 8]
|
|
async fn flush_terminal_frames<S>(
|
|
sink: &mut S,
|
|
terminal_ctrl_rx: &mut mpsc::Receiver<WsMessage>,
|
|
ctrl_rx: &mut mpsc::Receiver<WsMessage>,
|
|
disconnect_reason: &watch::Receiver<Option<CommunityDisconnectReason>>,
|
|
first_ctrl: Option<WsMessage>,
|
|
) where
|
|
S: Sink<WsMessage> + 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 !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 !send_terminal_frame(sink, ctrl_msg, deadline, "control").await {
|
|
return;
|
|
}
|
|
}
|
|
while let Ok(ctrl_msg) = ctrl_rx.try_recv() {
|
|
if !send_terminal_frame(sink, ctrl_msg, deadline, "control").await {
|
|
return;
|
|
}
|
|
}
|
|
// 3. Send Close.
|
|
let close = disconnect_reason
|
|
.borrow()
|
|
.map_or(WsMessage::Close(None), |reason| reason.close_message());
|
|
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<S>(
|
|
sink: &mut S,
|
|
frame: WsMessage,
|
|
deadline: tokio::time::Instant,
|
|
frame_kind: &'static str,
|
|
) -> bool
|
|
where
|
|
S: Sink<WsMessage> + 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::<Error>(),
|
|
Some(Error::AlreadyClosed | Error::Protocol(ProtocolError::SendAfterClosing))
|
|
)
|
|
})
|
|
}
|
|
|
|
/// Writes a WebSocket exit path's frames before a writer task owns the socket.
|
|
///
|
|
/// Every write shares one `WS_TERMINAL_FLUSH_TIMEOUT` deadline and the first
|
|
/// failed or timed-out write ends the drain, so a never-ready sink cannot hold
|
|
/// the exit, or the cleanup that follows it, open. [FI-TRACE-TERMINAL-BOUNDED]
|
|
pub(crate) async fn send_exit_frames_bounded<S>(
|
|
sink: &mut S,
|
|
frames: impl IntoIterator<Item = WsMessage>,
|
|
) where
|
|
S: Sink<WsMessage> + 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<S>(
|
|
sink: &mut S,
|
|
frames: impl IntoIterator<Item = WsMessage>,
|
|
deadline: tokio::time::Instant,
|
|
) -> bool
|
|
where
|
|
S: Sink<WsMessage> + Unpin,
|
|
S::Error: std::error::Error + 'static,
|
|
{
|
|
for frame in frames {
|
|
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<mpsc::Receiver<WsMessage>>,
|
|
close: Option<WsMessage>,
|
|
) -> futures_util::stream::SplitStream<WebSocket> {
|
|
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<S>(
|
|
mut ws_send: S,
|
|
mut data_rx: mpsc::Receiver<WsMessage>,
|
|
mut ctrl_rx: mpsc::Receiver<WsMessage>,
|
|
mut terminal_ctrl_rx: mpsc::Receiver<WsMessage>,
|
|
mut restart_rx: mpsc::Receiver<RestartClose>,
|
|
cancel: CancellationToken,
|
|
disconnect_reason: watch::Receiver<Option<CommunityDisconnectReason>>,
|
|
) where
|
|
S: Sink<WsMessage> + Unpin,
|
|
S::Error: std::error::Error + 'static,
|
|
{
|
|
loop {
|
|
// Priority: drain all pending control frames before data.
|
|
while let Ok(ctrl_msg) = ctrl_rx.try_recv() {
|
|
match send_or_cancel(&mut ws_send, ctrl_msg.clone(), &cancel).await {
|
|
WriterStep::Completed => {}
|
|
WriterStep::Cancelled => {
|
|
flush_terminal_frames(
|
|
&mut ws_send,
|
|
&mut terminal_ctrl_rx,
|
|
&mut ctrl_rx,
|
|
&disconnect_reason,
|
|
Some(ctrl_msg),
|
|
)
|
|
.await;
|
|
return;
|
|
}
|
|
WriterStep::Failed => return,
|
|
}
|
|
}
|
|
|
|
tokio::select! {
|
|
// Biased: restart > cancel > ordinary control > data. A restart
|
|
// command owns shutdown delivery and must flush its 1012 before
|
|
// cancellation can fall back to an unacknowledged close.
|
|
biased;
|
|
Some(restart) = restart_rx.recv() => {
|
|
let sent = matches!(
|
|
tokio::time::timeout(
|
|
WS_TERMINAL_FLUSH_TIMEOUT,
|
|
ws_send.send(WsMessage::Close(Some(axum::extract::ws::CloseFrame {
|
|
code: axum::extract::ws::close_code::RESTART,
|
|
reason: axum::extract::ws::Utf8Bytes::from_static("relay restarting"),
|
|
}))),
|
|
).await,
|
|
Ok(Ok(()))
|
|
);
|
|
let _ = restart.flushed.send(sent);
|
|
break;
|
|
}
|
|
_ = cancel.cancelled() => {
|
|
// Route FI terminal + ordinary control + Close through the shared
|
|
// bounded terminal path. flush_terminal_frames drains terminal_ctrl_rx
|
|
// first (denial before close), then ctrl_rx, then Close — all bounded
|
|
// by WS_TERMINAL_FLUSH_TIMEOUT. [FI-TRACE-TERMINAL-BOUNDED, Fix 8]
|
|
flush_terminal_frames(&mut ws_send, &mut terminal_ctrl_rx, &mut ctrl_rx, &disconnect_reason, None).await;
|
|
break;
|
|
}
|
|
Some(ctrl_msg) = ctrl_rx.recv() => {
|
|
match send_or_cancel(&mut ws_send, ctrl_msg.clone(), &cancel).await {
|
|
WriterStep::Completed => {}
|
|
WriterStep::Cancelled => {
|
|
flush_terminal_frames(
|
|
&mut ws_send,
|
|
&mut terminal_ctrl_rx,
|
|
&mut ctrl_rx,
|
|
&disconnect_reason,
|
|
Some(ctrl_msg),
|
|
)
|
|
.await;
|
|
break;
|
|
}
|
|
WriterStep::Failed => break,
|
|
}
|
|
}
|
|
Some(msg) = data_rx.recv() => {
|
|
let mut batched = 1usize;
|
|
match feed_or_cancel(&mut ws_send, msg, &cancel).await {
|
|
WriterStep::Completed => {}
|
|
WriterStep::Cancelled => {
|
|
flush_terminal_frames(
|
|
&mut ws_send,
|
|
&mut terminal_ctrl_rx,
|
|
&mut ctrl_rx,
|
|
&disconnect_reason,
|
|
None,
|
|
)
|
|
.await;
|
|
break;
|
|
}
|
|
WriterStep::Failed => break,
|
|
}
|
|
|
|
while batched < MAX_WS_SEND_BATCH {
|
|
match data_rx.try_recv() {
|
|
Ok(next) => {
|
|
match feed_or_cancel(&mut ws_send, next, &cancel).await {
|
|
WriterStep::Completed => {}
|
|
WriterStep::Cancelled => {
|
|
flush_terminal_frames(
|
|
&mut ws_send,
|
|
&mut terminal_ctrl_rx,
|
|
&mut ctrl_rx,
|
|
&disconnect_reason,
|
|
None,
|
|
)
|
|
.await;
|
|
return;
|
|
}
|
|
WriterStep::Failed => return,
|
|
}
|
|
batched += 1;
|
|
}
|
|
Err(tokio::sync::mpsc::error::TryRecvError::Empty) => break,
|
|
Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => break,
|
|
}
|
|
}
|
|
|
|
match flush_or_cancel(&mut ws_send, &cancel).await {
|
|
WriterStep::Completed => {}
|
|
WriterStep::Cancelled => {
|
|
flush_terminal_frames(
|
|
&mut ws_send,
|
|
&mut terminal_ctrl_rx,
|
|
&mut ctrl_rx,
|
|
&disconnect_reason,
|
|
None,
|
|
)
|
|
.await;
|
|
break;
|
|
}
|
|
WriterStep::Failed => break,
|
|
}
|
|
metrics::histogram!("buzz_ws_send_batch_size").record(batched as f64);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// 3 missed pongs → disconnect.
|
|
///
|
|
/// Sends Ping through the control channel so it isn't blocked by a full
|
|
/// data buffer. Uses `try_send` to keep the select loop responsive to
|
|
/// cancellation — a full control channel means the writer is stalled.
|
|
async fn heartbeat_loop(
|
|
ctrl_tx: mpsc::Sender<WsMessage>,
|
|
missed_pongs: Arc<AtomicU8>,
|
|
cancel: CancellationToken,
|
|
) {
|
|
let mut interval = tokio::time::interval(Duration::from_secs(30));
|
|
loop {
|
|
tokio::select! {
|
|
_ = interval.tick() => {
|
|
// fetch_add returns the *previous* value before incrementing:
|
|
// prev=0 → now 1 (first miss)
|
|
// prev=1 → now 2 (second miss)
|
|
// prev=2 → now 3 (third miss → disconnect)
|
|
let missed = missed_pongs.fetch_add(1, Ordering::Relaxed);
|
|
if missed >= 2 {
|
|
warn!("3 missed pongs — closing connection");
|
|
cancel.cancel();
|
|
break;
|
|
}
|
|
if ctrl_tx.try_send(WsMessage::Ping(axum::body::Bytes::new())).is_err() {
|
|
warn!("control channel full — cannot send Ping, closing");
|
|
cancel.cancel();
|
|
break;
|
|
}
|
|
}
|
|
_ = cancel.cancelled() => break,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// 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<WebSocket>,
|
|
conn: Arc<ConnectionState>,
|
|
state: Arc<AppState>,
|
|
missed_pongs: Arc<AtomicU8>,
|
|
cancel: CancellationToken,
|
|
) -> (
|
|
futures_util::stream::SplitStream<WebSocket>,
|
|
Option<tokio::task::JoinHandle<()>>,
|
|
) {
|
|
loop {
|
|
tokio::select! {
|
|
msg = ws_recv.next() => {
|
|
match msg {
|
|
Some(Ok(WsMessage::Text(text))) => {
|
|
let max_frame_bytes = state.config.max_frame_bytes;
|
|
if text.len() > max_frame_bytes {
|
|
warn!(
|
|
conn_id = %conn.conn_id,
|
|
bytes = text.len(),
|
|
max_frame_bytes,
|
|
"frame too large — disconnecting"
|
|
);
|
|
conn.send(format!(
|
|
r#"["NOTICE","error: frame too large ({} bytes, limit {})"]"#,
|
|
text.len(),
|
|
max_frame_bytes
|
|
));
|
|
break;
|
|
}
|
|
trace!(len = text.len(), "frame received");
|
|
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;
|
|
if bytes.len() > max_frame_bytes {
|
|
warn!(
|
|
conn_id = %conn.conn_id,
|
|
bytes = bytes.len(),
|
|
max_frame_bytes,
|
|
"binary frame too large — disconnecting"
|
|
);
|
|
conn.send(format!(
|
|
r#"["NOTICE","error: binary frame too large ({} bytes, limit {})"]"#,
|
|
bytes.len(),
|
|
max_frame_bytes
|
|
));
|
|
break;
|
|
}
|
|
// Binary frames: attempt UTF-8 decode and treat as text. Some clients
|
|
// (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()) {
|
|
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(_))) => {
|
|
missed_pongs.store(0, Ordering::Relaxed);
|
|
}
|
|
Some(Ok(WsMessage::Ping(data))) => {
|
|
// Send Pong through the control channel — priority
|
|
// delivery even when the data buffer is full (Bug 7 fix).
|
|
if conn.ctrl_tx.try_send(WsMessage::Pong(data)).is_err() {
|
|
// Control channel full means the socket writer is
|
|
// completely stalled — treat as terminal.
|
|
warn!(conn_id = %conn.conn_id, "control channel full — cannot send Pong, closing");
|
|
break;
|
|
}
|
|
}
|
|
Some(Ok(WsMessage::Close(_))) | None => {
|
|
debug!("WebSocket closed by client");
|
|
break;
|
|
}
|
|
Some(Err(e)) => {
|
|
debug!("WebSocket error: {e}");
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
_ = 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<Output = ()> + Send + 'static,
|
|
cancel: &CancellationToken,
|
|
) -> Option<tokio::task::JoinHandle<()>> {
|
|
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<WebSocket>) {
|
|
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<ConnectionState>, state: Arc<AppState>) {
|
|
// B2: Frame admission fence. If the connection's NIP-FI session has already
|
|
// expired (cancel fired by the expiry task), drop this frame before any
|
|
// handler dispatch. This closes the window where a buffered EVENT/REQ/AUTH
|
|
// is selected from the recv queue after expiry fires the cancel token.
|
|
// The check at the top of handle_text_message covers all message types
|
|
// uniformly — no individual handler needs its own fence.
|
|
if conn.cancel.is_cancelled() {
|
|
return;
|
|
}
|
|
|
|
let msg = match ClientMessage::parse(&text) {
|
|
Ok(m) => m,
|
|
Err(e) => {
|
|
conn.send(RelayMessage::notice(&format!("invalid message: {e}")));
|
|
return;
|
|
}
|
|
};
|
|
|
|
if !enforce_ws_admission(&msg, &conn, &state).await {
|
|
return;
|
|
}
|
|
|
|
match msg {
|
|
ClientMessage::Auth(event) => {
|
|
// AUTH remains inline so only one frame can race the connection's
|
|
// pending lifecycle, but cancellation can preempt dependency waits.
|
|
let span = tracing::info_span!("ws.auth", conn_id = %conn.conn_id);
|
|
tokio::select! {
|
|
biased;
|
|
_ = conn.cancel.cancelled() => {}
|
|
_ = handlers::auth::handle_auth(event, Arc::clone(&conn), Arc::clone(&state))
|
|
.instrument(span) => {}
|
|
}
|
|
}
|
|
ClientMessage::Event(event) => {
|
|
let conn = Arc::clone(&conn);
|
|
let state = Arc::clone(&state);
|
|
let permit = match state.handler_semaphore.clone().try_acquire_owned() {
|
|
Ok(p) => p,
|
|
Err(_) => {
|
|
// Correlate to the event id: a bare NOTICE here strands the
|
|
// client's pending publish exactly as an over-quota one did.
|
|
conn.send(request_rejection_message(
|
|
RejectionTarget::Event(event.id),
|
|
"rate-limited: too many concurrent requests",
|
|
));
|
|
return;
|
|
}
|
|
};
|
|
// Capture the parent span BEFORE the spawn so it is propagated into
|
|
// the spawned future. A bare `tokio::spawn` drops tracing context.
|
|
let span = tracing::info_span!(
|
|
"ws.event",
|
|
conn_id = %conn.conn_id,
|
|
event_id = tracing::field::Empty,
|
|
kind = tracing::field::Empty,
|
|
);
|
|
tokio::spawn(
|
|
async move {
|
|
handlers::event::handle_event(event, conn, state).await;
|
|
drop(permit);
|
|
}
|
|
.instrument(span),
|
|
);
|
|
}
|
|
ClientMessage::Req {
|
|
sub_id,
|
|
filters,
|
|
before_ids,
|
|
} => {
|
|
let conn = Arc::clone(&conn);
|
|
let state = Arc::clone(&state);
|
|
let permit = match state.handler_semaphore.clone().try_acquire_owned() {
|
|
Ok(p) => p,
|
|
Err(_) => {
|
|
conn.send(request_rejection_message(
|
|
RejectionTarget::Subscription(&sub_id),
|
|
"rate-limited: too many concurrent requests",
|
|
));
|
|
return;
|
|
}
|
|
};
|
|
let span = tracing::info_span!("ws.req", conn_id = %conn.conn_id, sub_id = %sub_id);
|
|
tokio::spawn(
|
|
async move {
|
|
handlers::req::handle_req(sub_id, filters, before_ids, conn, state).await;
|
|
drop(permit);
|
|
}
|
|
.instrument(span),
|
|
);
|
|
}
|
|
ClientMessage::ArtifactSync { sub_id, filter } => {
|
|
let conn = Arc::clone(&conn);
|
|
let state = Arc::clone(&state);
|
|
let permit = match state.handler_semaphore.clone().try_acquire_owned() {
|
|
Ok(p) => p,
|
|
Err(_) => {
|
|
// Retires whatever held the ID before `CLOSED`.
|
|
handlers::artifact_sync::refuse(
|
|
&sub_id,
|
|
"rate-limited: too many concurrent requests",
|
|
&conn,
|
|
&state,
|
|
)
|
|
.await;
|
|
return;
|
|
}
|
|
};
|
|
let span =
|
|
tracing::info_span!("ws.artifact_sync", conn_id = %conn.conn_id, sub_id = %sub_id);
|
|
tokio::spawn(
|
|
async move {
|
|
handlers::artifact_sync::handle_sync(sub_id, filter, conn, state, permit).await;
|
|
}
|
|
.instrument(span),
|
|
);
|
|
}
|
|
ClientMessage::Count { sub_id, filters } => {
|
|
let conn = Arc::clone(&conn);
|
|
let state = Arc::clone(&state);
|
|
let permit = match state.handler_semaphore.clone().try_acquire_owned() {
|
|
Ok(p) => p,
|
|
Err(_) => {
|
|
conn.send(request_rejection_message(
|
|
RejectionTarget::Subscription(&sub_id),
|
|
"rate-limited: too many concurrent requests",
|
|
));
|
|
return;
|
|
}
|
|
};
|
|
let span = tracing::info_span!("ws.count", conn_id = %conn.conn_id, sub_id = %sub_id);
|
|
tokio::spawn(
|
|
async move {
|
|
handlers::count::handle_count(sub_id, filters, conn, state).await;
|
|
drop(permit);
|
|
}
|
|
.instrument(span),
|
|
);
|
|
}
|
|
ClientMessage::Close(sub_id) => {
|
|
handlers::close::handle_close(sub_id, Arc::clone(&conn), Arc::clone(&state)).await;
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(crate) mod test_seam {
|
|
use super::*;
|
|
|
|
pub(crate) type Hold =
|
|
tokio::sync::oneshot::Sender<tokio::sync::OwnedMutexGuard<HashMap<String, u64>>>;
|
|
|
|
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<Option<Hold>>;
|
|
}
|
|
|
|
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::*;
|
|
use axum::{extract::ws::WebSocketUpgrade, routing::get, Router};
|
|
use metrics_util::debugging::DebugValue;
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
use buzz_auth::AuthMethod;
|
|
use nostr::{EventBuilder, Keys, Kind, RelayUrl};
|
|
use tokio::net::TcpListener;
|
|
use tokio::sync::Notify;
|
|
use tokio_tungstenite::{connect_async, tungstenite::Message};
|
|
|
|
/// A test connection together with the receiving end of each outbound queue.
|
|
pub(crate) struct TestConn {
|
|
/// The connection under test.
|
|
pub(crate) conn: Arc<ConnectionState>,
|
|
/// Receives data frames (EVENT, NOTICE, OK, ...).
|
|
pub(crate) send_rx: mpsc::Receiver<WsMessage>,
|
|
/// Receives control frames; capacity 8, as in production.
|
|
pub(crate) ctrl_rx: mpsc::Receiver<WsMessage>,
|
|
/// Receives the terminal NIP-FI denial frame; capacity 1, as in production.
|
|
pub(crate) terminal_rx: mpsc::Receiver<WsMessage>,
|
|
}
|
|
|
|
/// Builds a connection in `auth` carrying the NIP-FI assertion presented at
|
|
/// upgrade, if any. The admission gate is off-mode and no session deadline
|
|
/// is set.
|
|
///
|
|
/// Lives here, next to `ConnectionState`, so the crate has one place that
|
|
/// knows how to build one. Shared with `crate::rejection`'s and
|
|
/// `crate::nip_fi_session`'s tests.
|
|
pub(crate) fn test_conn(
|
|
auth: AuthState,
|
|
nip_fi_assertion: Option<buzz_auth::VerifiedAssertion>,
|
|
) -> TestConn {
|
|
test_conn_in(
|
|
auth,
|
|
nip_fi_assertion,
|
|
buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()),
|
|
)
|
|
}
|
|
|
|
/// [`test_conn`] whose tenant is `community`.
|
|
pub(crate) fn test_conn_in(
|
|
auth: AuthState,
|
|
nip_fi_assertion: Option<buzz_auth::VerifiedAssertion>,
|
|
community: buzz_core::tenant::CommunityId,
|
|
) -> TestConn {
|
|
let (send_tx, send_rx) = mpsc::channel(4);
|
|
let (ctrl_tx, ctrl_rx) = mpsc::channel(8);
|
|
let (terminal_ctrl_tx, terminal_rx) = mpsc::channel(1);
|
|
let cancel = CancellationToken::new();
|
|
let conn = ConnectionState {
|
|
conn_id: Uuid::new_v4(),
|
|
tenant: TenantContext::resolved(community, "test.local".to_string()),
|
|
remote_addr: "127.0.0.1:1234".parse().expect("socket addr"),
|
|
auth_state: StdMutex::new(auth),
|
|
subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
|
|
send_tx,
|
|
ctrl_tx,
|
|
terminal_ctrl_tx,
|
|
cancel: cancel.clone(),
|
|
backpressure_count: Arc::new(AtomicU8::new(0)),
|
|
grace_limit: 3,
|
|
nip_fi_assertion,
|
|
session_deadline: None,
|
|
nip_fi_gate: crate::nip_fi_gate::SessionAdmissionGate::off_mode(cancel.clone()),
|
|
community_control: crate::state::CommunityConnectionControl::new(cancel.clone()),
|
|
};
|
|
TestConn {
|
|
conn: Arc::new(conn),
|
|
send_rx,
|
|
ctrl_rx,
|
|
terminal_rx,
|
|
}
|
|
}
|
|
|
|
/// A connection without a NIP-FI assertion whose data frames a test can
|
|
/// read back.
|
|
pub(crate) fn test_conn_with_auth(
|
|
auth: AuthState,
|
|
) -> (Arc<ConnectionState>, mpsc::Receiver<WsMessage>) {
|
|
let TestConn { conn, send_rx, .. } = test_conn(auth, None);
|
|
(conn, send_rx)
|
|
}
|
|
|
|
/// An authenticated connection — the only state admission quotas apply to.
|
|
pub(crate) fn authenticated_state() -> AuthState {
|
|
AuthState::Authenticated(auth_context())
|
|
}
|
|
|
|
fn auth_context() -> AuthContext {
|
|
AuthContext {
|
|
pubkey: Keys::generate().public_key(),
|
|
scopes: Vec::new(),
|
|
channel_ids: None,
|
|
auth_method: AuthMethod::Nip42,
|
|
agent_owner_pubkey: None,
|
|
}
|
|
}
|
|
|
|
/// A connection awaiting its NIP-42 AUTH response.
|
|
pub(crate) fn pending_state() -> AuthState {
|
|
AuthState::Pending {
|
|
challenge: "test-challenge".to_owned(),
|
|
started_at: Instant::now(),
|
|
}
|
|
}
|
|
|
|
type MetricSnapshot = Vec<(
|
|
metrics_util::CompositeKey,
|
|
Option<metrics::Unit>,
|
|
Option<metrics::SharedString>,
|
|
DebugValue,
|
|
)>;
|
|
|
|
fn counter_value(snapshot: &MetricSnapshot, name: &str, outcome: Option<&str>) -> u64 {
|
|
snapshot
|
|
.iter()
|
|
.find_map(|(key, _, _, value)| {
|
|
if key.key().name() != name {
|
|
return None;
|
|
}
|
|
let labels = key.key().labels().collect::<Vec<_>>();
|
|
if outcome.is_some_and(|expected| {
|
|
!labels
|
|
.iter()
|
|
.any(|label| label.key() == "outcome" && label.value() == expected)
|
|
}) {
|
|
return None;
|
|
}
|
|
let DebugValue::Counter(value) = value else {
|
|
panic!("{name} must be a counter");
|
|
};
|
|
Some(*value)
|
|
})
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
fn labeled_counter_value(
|
|
snapshot: &MetricSnapshot,
|
|
name: &str,
|
|
label_key: &str,
|
|
label_value: &str,
|
|
) -> u64 {
|
|
snapshot
|
|
.iter()
|
|
.find_map(|(key, _, _, value)| {
|
|
(key.key().name() == name
|
|
&& key
|
|
.key()
|
|
.labels()
|
|
.any(|label| label.key() == label_key && label.value() == label_value))
|
|
.then(|| match value {
|
|
DebugValue::Counter(value) => *value,
|
|
_ => panic!("{name} must be a counter"),
|
|
})
|
|
})
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
fn labeled_gauge_value(snapshot: &MetricSnapshot, name: &str, labels: &[(&str, &str)]) -> f64 {
|
|
snapshot
|
|
.iter()
|
|
.find_map(|(key, _, _, value)| {
|
|
let matches = key.key().name() == name
|
|
&& labels.iter().all(|(expected_key, expected_value)| {
|
|
key.key().labels().any(|label| {
|
|
label.key() == *expected_key && label.value() == *expected_value
|
|
})
|
|
});
|
|
matches.then(|| match value {
|
|
DebugValue::Gauge(value) => value.into_inner(),
|
|
_ => panic!("{name} must be a gauge"),
|
|
})
|
|
})
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
fn authenticated_gauge(snapshot: &MetricSnapshot) -> f64 {
|
|
snapshot
|
|
.iter()
|
|
.find_map(|(key, _, _, value)| {
|
|
if key.key().name() != "buzz_ws_authenticated_connections_active" {
|
|
return None;
|
|
}
|
|
let DebugValue::Gauge(value) = value else {
|
|
panic!("authenticated connections must be a gauge");
|
|
};
|
|
Some(value.into_inner())
|
|
})
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
pub(crate) fn read_frame(rx: &mut mpsc::Receiver<WsMessage>) -> serde_json::Value {
|
|
match rx.try_recv().expect("a frame was sent") {
|
|
WsMessage::Text(text) => serde_json::from_str(&text).expect("valid JSON frame"),
|
|
other => panic!("unexpected websocket message: {other:?}"),
|
|
}
|
|
}
|
|
|
|
/// Exercise the real state-transition methods for every terminal. The
|
|
/// attempt counter must reconcile with exactly one terminal per completed
|
|
/// attempt, and repeated/racing terminal calls must not drive the active
|
|
/// authenticated gauge below zero.
|
|
#[tokio::test(flavor = "current_thread")]
|
|
async fn auth_lifecycle_reconciles_every_terminal_and_never_leaks_gauge() {
|
|
let recorder = metrics_util::debugging::DebuggingRecorder::new();
|
|
let snapshotter = recorder.snapshotter();
|
|
let _recorder_guard = metrics::set_default_local_recorder(&recorder);
|
|
|
|
crate::metrics::record_auth_attempt_started();
|
|
let (success, _rx) = test_conn_with_auth(pending_state());
|
|
assert!(success.authenticate(auth_context()));
|
|
assert!(
|
|
!success.authenticate(auth_context()),
|
|
"a terminal attempt cannot authenticate twice"
|
|
);
|
|
assert!(
|
|
!success.finish_pending_auth_on_cancel(AuthOutcome::Disconnect),
|
|
"the cancellation watcher must leave authenticated context for lifecycle cleanup"
|
|
);
|
|
assert!(success
|
|
.finish_auth_on_close(AuthOutcome::Disconnect)
|
|
.is_some());
|
|
assert!(
|
|
success
|
|
.finish_auth_on_close(AuthOutcome::Shutdown)
|
|
.is_none(),
|
|
"cleanup must be idempotent"
|
|
);
|
|
|
|
for outcome in [
|
|
AuthOutcome::Invalid,
|
|
AuthOutcome::Banned,
|
|
AuthOutcome::BanCheckError,
|
|
AuthOutcome::AllowlistCheckError,
|
|
AuthOutcome::AllowlistDenied,
|
|
AuthOutcome::RelayMembershipCheckError,
|
|
AuthOutcome::NotRelayMember,
|
|
AuthOutcome::PairingMismatch,
|
|
] {
|
|
crate::metrics::record_auth_attempt_started();
|
|
let (denied, _rx) = test_conn_with_auth(pending_state());
|
|
assert!(denied.reject_auth(outcome));
|
|
assert!(!denied.reject_auth(outcome), "a denial cannot record twice");
|
|
assert!(denied
|
|
.finish_auth_on_close(AuthOutcome::Disconnect)
|
|
.is_none());
|
|
}
|
|
|
|
crate::metrics::record_auth_attempt_started();
|
|
let (timed_out, _rx) = test_conn_with_auth(pending_state());
|
|
assert!(timed_out.expire_auth());
|
|
assert!(timed_out
|
|
.finish_auth_on_close(AuthOutcome::Disconnect)
|
|
.is_none());
|
|
|
|
for outcome in [AuthOutcome::Disconnect, AuthOutcome::Shutdown] {
|
|
crate::metrics::record_auth_attempt_started();
|
|
let (closed, _rx) = test_conn_with_auth(pending_state());
|
|
assert!(closed.finish_auth_on_close(outcome).is_none());
|
|
assert!(closed.finish_auth_on_close(outcome).is_none());
|
|
}
|
|
|
|
let snapshot = snapshotter.snapshot().into_vec();
|
|
let attempts = counter_value(&snapshot, "buzz_auth_attempts_total", None);
|
|
let outcomes = AuthOutcome::ALL
|
|
.iter()
|
|
.map(|outcome| {
|
|
let value = counter_value(
|
|
&snapshot,
|
|
"buzz_auth_outcomes_total",
|
|
Some(outcome.as_str()),
|
|
);
|
|
assert_eq!(
|
|
value,
|
|
1,
|
|
"{} terminal must be recorded exactly once",
|
|
outcome.as_str()
|
|
);
|
|
value
|
|
})
|
|
.sum::<u64>();
|
|
|
|
assert_eq!(attempts, AuthOutcome::ALL.len() as u64);
|
|
assert_eq!(attempts, outcomes);
|
|
assert_eq!(authenticated_gauge(&snapshot), 0.0);
|
|
}
|
|
|
|
/// Post-terminal AUTH floods traverse the production frame dispatcher but
|
|
/// cannot mint rollout-gating challenge lifecycles or outcomes.
|
|
#[tokio::test(flavor = "current_thread")]
|
|
async fn auth_flood_after_terminal_only_increments_protocol_noise_metric() {
|
|
const FRAMES_PER_STATE: u64 = 64;
|
|
|
|
let recorder = metrics_util::debugging::DebuggingRecorder::new();
|
|
let snapshotter = recorder.snapshotter();
|
|
let _recorder_guard = metrics::set_default_local_recorder(&recorder);
|
|
let state = crate::state::tests::test_state().await;
|
|
let mut authoritative_attempts = 0;
|
|
let mut authoritative_outcomes = 0;
|
|
|
|
for (auth_state, expected_state) in [
|
|
(
|
|
authenticated_state(),
|
|
crate::metrics::AuthPostTerminalState::Authenticated,
|
|
),
|
|
(
|
|
AuthState::Failed,
|
|
crate::metrics::AuthPostTerminalState::Failed,
|
|
),
|
|
] {
|
|
let (conn, mut rx) = test_conn_with_auth(auth_state);
|
|
for sequence in 0..FRAMES_PER_STATE {
|
|
let event = EventBuilder::new(Kind::Authentication, format!("noise-{sequence}"))
|
|
.sign_with_keys(&Keys::generate())
|
|
.expect("sign AUTH noise event");
|
|
let raw = serde_json::json!(["AUTH", event]).to_string();
|
|
handle_text_message(raw, Arc::clone(&conn), Arc::clone(&state)).await;
|
|
let frame = read_frame(&mut rx);
|
|
assert_eq!(frame[2], false);
|
|
}
|
|
assert!(!conn.cancel.is_cancelled());
|
|
let snapshot = snapshotter.snapshot().into_vec();
|
|
authoritative_attempts += counter_value(&snapshot, "buzz_auth_attempts_total", None);
|
|
authoritative_outcomes += counter_value(&snapshot, "buzz_auth_outcomes_total", None);
|
|
assert_eq!(
|
|
labeled_counter_value(
|
|
&snapshot,
|
|
"buzz_auth_post_terminal_frames_total",
|
|
"state",
|
|
expected_state.as_str(),
|
|
),
|
|
FRAMES_PER_STATE
|
|
);
|
|
}
|
|
|
|
let snapshot = snapshotter.snapshot().into_vec();
|
|
authoritative_attempts += counter_value(&snapshot, "buzz_auth_attempts_total", None);
|
|
authoritative_outcomes += counter_value(&snapshot, "buzz_auth_outcomes_total", None);
|
|
assert_eq!(
|
|
authoritative_attempts, 0,
|
|
"post-terminal protocol noise cannot create authoritative attempts"
|
|
);
|
|
assert_eq!(
|
|
authoritative_outcomes, 0,
|
|
"post-terminal protocol noise cannot create authoritative outcomes"
|
|
);
|
|
}
|
|
|
|
/// Aborting the production lifecycle owner must synchronously terminalize
|
|
/// pending accounting and release an authenticated gauge exactly once.
|
|
#[tokio::test(flavor = "current_thread")]
|
|
async fn aborting_lifecycle_owner_is_drop_safe() {
|
|
let recorder = metrics_util::debugging::DebuggingRecorder::new();
|
|
let snapshotter = recorder.snapshotter();
|
|
let _recorder_guard = metrics::set_default_local_recorder(&recorder);
|
|
|
|
crate::metrics::record_auth_attempt_started();
|
|
let (authenticated, _rx) = test_conn_with_auth(pending_state());
|
|
assert!(authenticated.authenticate(auth_context()));
|
|
let authenticated_started = Arc::new(Notify::new());
|
|
let task = {
|
|
let conn = Arc::clone(&authenticated);
|
|
let started = Arc::clone(&authenticated_started);
|
|
tokio::spawn(async move {
|
|
let _guard = AuthLifecycleGuard::new(conn);
|
|
started.notify_one();
|
|
std::future::pending::<()>().await;
|
|
})
|
|
};
|
|
authenticated_started.notified().await;
|
|
task.abort();
|
|
assert!(task.await.expect_err("task was aborted").is_cancelled());
|
|
|
|
crate::metrics::record_auth_attempt_started();
|
|
let (pending, _rx) = test_conn_with_auth(pending_state());
|
|
let pending_started = Arc::new(Notify::new());
|
|
let task = {
|
|
let conn = Arc::clone(&pending);
|
|
let started = Arc::clone(&pending_started);
|
|
tokio::spawn(async move {
|
|
let _guard = AuthLifecycleGuard::new(conn);
|
|
started.notify_one();
|
|
std::future::pending::<()>().await;
|
|
})
|
|
};
|
|
pending_started.notified().await;
|
|
task.abort();
|
|
assert!(task.await.expect_err("task was aborted").is_cancelled());
|
|
|
|
let snapshot = snapshotter.snapshot().into_vec();
|
|
assert_eq!(
|
|
counter_value(&snapshot, "buzz_auth_attempts_total", None),
|
|
2
|
|
);
|
|
assert_eq!(
|
|
counter_value(
|
|
&snapshot,
|
|
"buzz_auth_outcomes_total",
|
|
Some(AuthOutcome::Success.as_str()),
|
|
),
|
|
1
|
|
);
|
|
assert_eq!(
|
|
counter_value(
|
|
&snapshot,
|
|
"buzz_auth_outcomes_total",
|
|
Some(AuthOutcome::Disconnect.as_str()),
|
|
),
|
|
1
|
|
);
|
|
assert_eq!(authenticated_gauge(&snapshot), 0.0);
|
|
assert!(matches!(pending.auth_state_snapshot(), AuthState::Failed));
|
|
}
|
|
|
|
/// Drive the real upgraded WebSocket lifecycle until AUTH is waiting for
|
|
/// the sole database connection. Production shutdown must claim the
|
|
/// pending attempt before that dependency is released, and the late DB
|
|
/// result must not turn the terminal into success.
|
|
#[tokio::test(flavor = "current_thread")]
|
|
async fn shutdown_terminalizes_auth_stalled_on_database_before_late_success() {
|
|
let recorder = metrics_util::debugging::DebuggingRecorder::new();
|
|
let snapshotter = recorder.snapshotter();
|
|
let _recorder_guard = metrics::set_default_local_recorder(&recorder);
|
|
|
|
// Occupy the pool's only connection attempt with a fake PostgreSQL
|
|
// endpoint that accepts TCP and never completes the startup handshake.
|
|
// The production AUTH query then waits for the pool permit without
|
|
// requiring a developer database or relying on timing alone.
|
|
let fake_database = TcpListener::bind("127.0.0.1:0")
|
|
.await
|
|
.expect("bind fake PostgreSQL endpoint");
|
|
let database_url = format!(
|
|
"postgres://buzz@{}/buzz",
|
|
fake_database.local_addr().expect("fake database address")
|
|
);
|
|
let pool = sqlx::postgres::PgPoolOptions::new()
|
|
.max_connections(1)
|
|
.acquire_timeout(Duration::from_secs(30))
|
|
.connect_lazy(&database_url)
|
|
.expect("create lifecycle test pool");
|
|
let blocker_pool = pool.clone();
|
|
let blocker = tokio::spawn(async move { blocker_pool.acquire().await });
|
|
let (blocked_database_stream, _) = fake_database
|
|
.accept()
|
|
.await
|
|
.expect("pool reached fake PostgreSQL endpoint");
|
|
let state = crate::state::tests::test_state_with_database_pool(pool.clone()).await;
|
|
let tenant = TenantContext::resolved(
|
|
buzz_core::tenant::CommunityId::from_uuid(Uuid::nil()),
|
|
"test.local".to_owned(),
|
|
);
|
|
let expected_relay_url: RelayUrl =
|
|
crate::api::bridge::nip42_expected_relay_url(&state.config.relay_url, &tenant)
|
|
.parse()
|
|
.expect("expected NIP-42 relay URL");
|
|
|
|
let connection_finished = Arc::new(Notify::new());
|
|
let route_state = Arc::clone(&state);
|
|
let route_tenant = tenant.clone();
|
|
let route_finished = Arc::clone(&connection_finished);
|
|
let app = Router::new().route(
|
|
"/",
|
|
get(move |ws: WebSocketUpgrade| {
|
|
let state = Arc::clone(&route_state);
|
|
let tenant = route_tenant.clone();
|
|
let finished = Arc::clone(&route_finished);
|
|
async move {
|
|
ws.on_upgrade(move |socket| async move {
|
|
let cancel = CancellationToken::new();
|
|
let control = CommunityConnectionControl::new(cancel);
|
|
handle_active_connection(
|
|
socket,
|
|
state,
|
|
"127.0.0.1:1234".parse().expect("client address"),
|
|
tenant,
|
|
Uuid::new_v4(),
|
|
control,
|
|
None,
|
|
chrono::Utc::now(),
|
|
None,
|
|
)
|
|
.await;
|
|
finished.notify_one();
|
|
})
|
|
}
|
|
}),
|
|
);
|
|
let listener = TcpListener::bind("127.0.0.1:0")
|
|
.await
|
|
.expect("bind lifecycle WebSocket listener");
|
|
let address = listener.local_addr().expect("listener address");
|
|
let server = tokio::spawn(async move {
|
|
axum::serve(listener, app)
|
|
.await
|
|
.expect("serve lifecycle WebSocket");
|
|
});
|
|
|
|
let (mut client, _) = connect_async(format!("ws://{address}/"))
|
|
.await
|
|
.expect("connect lifecycle WebSocket");
|
|
let challenge_frame = client
|
|
.next()
|
|
.await
|
|
.expect("challenge frame")
|
|
.expect("read challenge");
|
|
let Message::Text(challenge_text) = challenge_frame else {
|
|
panic!("expected text challenge")
|
|
};
|
|
let challenge_json: serde_json::Value =
|
|
serde_json::from_str(&challenge_text).expect("parse challenge");
|
|
assert_eq!(challenge_json[0], "AUTH");
|
|
let challenge = challenge_json[1].as_str().expect("challenge string");
|
|
let auth_event = EventBuilder::auth(challenge, expected_relay_url)
|
|
.sign_with_keys(&Keys::generate())
|
|
.expect("sign NIP-42 AUTH");
|
|
client
|
|
.send(Message::Text(
|
|
serde_json::json!(["AUTH", auth_event]).to_string().into(),
|
|
))
|
|
.await
|
|
.expect("send NIP-42 AUTH");
|
|
|
|
let attempts_before_shutdown = tokio::time::timeout(Duration::from_secs(2), async {
|
|
let mut attempts = 0;
|
|
loop {
|
|
let snapshot = snapshotter.snapshot().into_vec();
|
|
attempts += counter_value(&snapshot, "buzz_auth_attempts_total", None);
|
|
if labeled_gauge_value(
|
|
&snapshot,
|
|
"buzz_db_pool_waiters",
|
|
&[("pool_role", "writer"), ("operation", "authorization")],
|
|
) >= 1.0
|
|
{
|
|
break attempts;
|
|
}
|
|
tokio::task::yield_now().await;
|
|
}
|
|
})
|
|
.await
|
|
.expect("AUTH reached the blocked production DB acquisition");
|
|
assert_eq!(
|
|
attempts_before_shutdown, 1,
|
|
"the production-issued challenge must own the attempt start"
|
|
);
|
|
|
|
state.begin_shutdown();
|
|
assert_eq!(state.conn_manager.drain_all(), 1);
|
|
tokio::time::timeout(Duration::from_secs(2), connection_finished.notified())
|
|
.await
|
|
.expect("production connection lifecycle finishes after shutdown");
|
|
|
|
// Release the dependency only after production shutdown has claimed
|
|
// the pending lifecycle, then prove no late handler result can win.
|
|
drop(blocked_database_stream);
|
|
blocker.abort();
|
|
let _ = blocker.await;
|
|
pool.close().await;
|
|
|
|
let snapshot = snapshotter.snapshot().into_vec();
|
|
assert_eq!(
|
|
counter_value(
|
|
&snapshot,
|
|
"buzz_auth_outcomes_total",
|
|
Some(AuthOutcome::Shutdown.as_str()),
|
|
),
|
|
1
|
|
);
|
|
assert_eq!(
|
|
counter_value(
|
|
&snapshot,
|
|
"buzz_auth_outcomes_total",
|
|
Some(AuthOutcome::Success.as_str()),
|
|
),
|
|
0,
|
|
"releasing the dependency after shutdown must not record late success"
|
|
);
|
|
assert_eq!(authenticated_gauge(&snapshot), 0.0);
|
|
|
|
drop(client);
|
|
server.abort();
|
|
let _ = server.await;
|
|
}
|
|
|
|
/// Drives the real `handle_text_message` with every handler permit held, so
|
|
/// the EVENT saturation branch is reached through production dispatch rather
|
|
/// than by calling its helpers directly.
|
|
///
|
|
/// This must go through `handle_text_message`: a test that renders the
|
|
/// rejection frame itself stays green when the call site inside the match
|
|
/// arm is reverted to a bare `NOTICE`.
|
|
#[tokio::test]
|
|
async fn saturated_handler_rejects_an_event_on_the_ok_channel() {
|
|
let state = crate::state::tests::test_state().await;
|
|
// An unauthenticated connection skips the admission quotas, so the
|
|
// semaphore is the only gate the frame can trip.
|
|
let (conn, mut rx) = test_conn_with_auth(AuthState::Failed);
|
|
|
|
let permits = state.handler_semaphore.available_permits();
|
|
let _held = Arc::clone(&state.handler_semaphore)
|
|
.acquire_many_owned(permits as u32)
|
|
.await
|
|
.expect("hold every handler permit");
|
|
|
|
let event = EventBuilder::new(Kind::TextNote, "hello")
|
|
.sign_with_keys(&Keys::generate())
|
|
.expect("sign event");
|
|
let event_id = event.id.to_hex();
|
|
let raw = serde_json::json!(["EVENT", event]).to_string();
|
|
|
|
handle_text_message(raw, Arc::clone(&conn), Arc::clone(&state)).await;
|
|
|
|
let frame = read_frame(&mut rx);
|
|
assert_eq!(
|
|
frame[0], "OK",
|
|
"an EVENT turned away for handler saturation must be rejected on the \
|
|
OK channel — a NOTICE carries no event id, so the client's pending \
|
|
publish cannot be settled and the send only times out"
|
|
);
|
|
assert_eq!(frame[1], event_id);
|
|
assert_eq!(frame[2], false);
|
|
assert_eq!(frame[3], "rate-limited: too many concurrent requests");
|
|
}
|
|
|
|
/// The REQ arm of the same branch still settles on CLOSED.
|
|
#[tokio::test]
|
|
async fn saturated_handler_rejects_a_req_on_the_closed_channel() {
|
|
let state = crate::state::tests::test_state().await;
|
|
let (conn, mut rx) = test_conn_with_auth(AuthState::Failed);
|
|
|
|
let permits = state.handler_semaphore.available_permits();
|
|
let _held = Arc::clone(&state.handler_semaphore)
|
|
.acquire_many_owned(permits as u32)
|
|
.await
|
|
.expect("hold every handler permit");
|
|
|
|
let raw = serde_json::json!(["REQ", "history-abc", {"kinds": [1]}]).to_string();
|
|
handle_text_message(raw, Arc::clone(&conn), Arc::clone(&state)).await;
|
|
|
|
let frame = read_frame(&mut rx);
|
|
assert_eq!(frame[0], "CLOSED");
|
|
assert_eq!(frame[1], "history-abc");
|
|
}
|
|
|
|
/// COUNT refusals follow NIP-45 and close the named query.
|
|
#[tokio::test]
|
|
async fn saturated_handler_rejects_a_count_on_the_closed_channel() {
|
|
let state = crate::state::tests::test_state().await;
|
|
let (conn, mut rx) = test_conn_with_auth(AuthState::Failed);
|
|
|
|
let permits = state.handler_semaphore.available_permits();
|
|
let _held = Arc::clone(&state.handler_semaphore)
|
|
.acquire_many_owned(permits as u32)
|
|
.await
|
|
.expect("hold every handler permit");
|
|
|
|
let raw = serde_json::json!(["COUNT", "count-abc", {"kinds": [1]}]).to_string();
|
|
handle_text_message(raw, Arc::clone(&conn), Arc::clone(&state)).await;
|
|
|
|
let frame = read_frame(&mut rx);
|
|
assert_eq!(frame[0], "CLOSED");
|
|
assert_eq!(frame[1], "count-abc");
|
|
assert_eq!(frame[2], "rate-limited: too many concurrent requests");
|
|
}
|
|
|
|
/// A sync REQ refused before its handler runs still replaces the old
|
|
/// holder of its ID (NIP-01): the `CLOSED` must not leave a live
|
|
/// subscription that keeps fanning out under a closed ID. Covers handler
|
|
/// saturation and WS admission, both through production dispatch.
|
|
/// Mutation: send either refusal without `retire_locked` → RED.
|
|
#[tokio::test]
|
|
async fn pre_handler_sync_refusal_retires_the_old_holder() {
|
|
use buzz_pubsub::EventTopic;
|
|
for (auth, saturate, reason) in [
|
|
(
|
|
AuthState::Failed,
|
|
true,
|
|
"rate-limited: too many concurrent requests",
|
|
),
|
|
(
|
|
authenticated_state(),
|
|
false,
|
|
"rate-limited: shared admission unavailable",
|
|
),
|
|
] {
|
|
// Redis is unreachable here, so authenticated admission is
|
|
// `Unavailable`; an unauthenticated connection skips admission.
|
|
let state = crate::state::tests::test_state().await;
|
|
let (conn, mut rx) = test_conn_with_auth(auth);
|
|
let channel = uuid::Uuid::new_v4();
|
|
let topic = EventTopic::Channel(channel);
|
|
handlers::req::claim_live_subscription(
|
|
"s",
|
|
&[nostr::Filter::new().kind(Kind::TextNote)],
|
|
Some(&[channel]),
|
|
&conn,
|
|
&state,
|
|
)
|
|
.await
|
|
.expect("claim the old holder");
|
|
let _held = if saturate {
|
|
let permits = state.handler_semaphore.available_permits();
|
|
Some(
|
|
Arc::clone(&state.handler_semaphore)
|
|
.acquire_many_owned(permits as u32)
|
|
.await
|
|
.expect("hold every handler permit"),
|
|
)
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let raw = serde_json::json!(
|
|
["REQ", "s", {"artifact": "current", "#h": [channel.to_string()]}]
|
|
)
|
|
.to_string();
|
|
handle_text_message(raw, Arc::clone(&conn), Arc::clone(&state)).await;
|
|
|
|
assert!(conn.subscriptions.lock().await.is_empty(), "{reason}");
|
|
assert!(
|
|
state.sub_registry.get_filters(conn.conn_id, "s").is_none(),
|
|
"{reason}"
|
|
);
|
|
assert!(
|
|
!state.artifact_syncs.contains(conn.conn_id, "s"),
|
|
"{reason}"
|
|
);
|
|
assert_eq!(
|
|
state.pubsub.topic_refcount(&conn.tenant, topic).await,
|
|
0,
|
|
"{reason}"
|
|
);
|
|
let frame = read_frame(&mut rx);
|
|
assert_eq!(frame, serde_json::json!(["CLOSED", "s", reason]));
|
|
assert!(rx.try_recv().is_err(), "{reason}: only CLOSED is sent");
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
struct MockSinkState {
|
|
messages: Vec<WsMessage>,
|
|
flush_count: usize,
|
|
fail_after_flushes: Option<usize>,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
struct MockSink {
|
|
state: Arc<Mutex<MockSinkState>>,
|
|
}
|
|
|
|
impl MockSink {
|
|
fn new(fail_after_flushes: Option<usize>) -> (Self, Arc<Mutex<MockSinkState>>) {
|
|
let state = Arc::new(Mutex::new(MockSinkState {
|
|
fail_after_flushes,
|
|
..MockSinkState::default()
|
|
}));
|
|
(
|
|
Self {
|
|
state: Arc::clone(&state),
|
|
},
|
|
state,
|
|
)
|
|
}
|
|
}
|
|
|
|
impl Sink<WsMessage> for MockSink {
|
|
type Error = std::io::Error;
|
|
|
|
fn poll_ready(
|
|
self: std::pin::Pin<&mut Self>,
|
|
_cx: &mut std::task::Context<'_>,
|
|
) -> std::task::Poll<Result<(), Self::Error>> {
|
|
std::task::Poll::Ready(Ok(()))
|
|
}
|
|
|
|
fn start_send(self: std::pin::Pin<&mut Self>, item: WsMessage) -> Result<(), Self::Error> {
|
|
self.state
|
|
.lock()
|
|
.expect("mock sink poisoned")
|
|
.messages
|
|
.push(item);
|
|
Ok(())
|
|
}
|
|
|
|
fn poll_flush(
|
|
self: std::pin::Pin<&mut Self>,
|
|
_cx: &mut std::task::Context<'_>,
|
|
) -> std::task::Poll<Result<(), Self::Error>> {
|
|
let mut state = self.state.lock().expect("mock sink poisoned");
|
|
state.flush_count += 1;
|
|
if state
|
|
.fail_after_flushes
|
|
.is_some_and(|limit| state.flush_count >= limit)
|
|
{
|
|
return std::task::Poll::Ready(Err(std::io::Error::new(
|
|
std::io::ErrorKind::BrokenPipe,
|
|
"mock flush failure",
|
|
)));
|
|
}
|
|
std::task::Poll::Ready(Ok(()))
|
|
}
|
|
|
|
fn poll_close(
|
|
self: std::pin::Pin<&mut Self>,
|
|
cx: &mut std::task::Context<'_>,
|
|
) -> std::task::Poll<Result<(), Self::Error>> {
|
|
self.poll_flush(cx)
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct NeverReadySink {
|
|
ready_polled: Arc<Notify>,
|
|
}
|
|
|
|
impl Sink<WsMessage> for NeverReadySink {
|
|
type Error = std::io::Error;
|
|
|
|
fn poll_ready(
|
|
self: std::pin::Pin<&mut Self>,
|
|
_cx: &mut std::task::Context<'_>,
|
|
) -> std::task::Poll<Result<(), Self::Error>> {
|
|
self.ready_polled.notify_one();
|
|
std::task::Poll::Pending
|
|
}
|
|
|
|
fn start_send(self: std::pin::Pin<&mut Self>, _item: WsMessage) -> Result<(), Self::Error> {
|
|
panic!("a never-ready sink must not accept a frame")
|
|
}
|
|
|
|
fn poll_flush(
|
|
self: std::pin::Pin<&mut Self>,
|
|
_cx: &mut std::task::Context<'_>,
|
|
) -> std::task::Poll<Result<(), Self::Error>> {
|
|
std::task::Poll::Pending
|
|
}
|
|
|
|
fn poll_close(
|
|
self: std::pin::Pin<&mut Self>,
|
|
_cx: &mut std::task::Context<'_>,
|
|
) -> std::task::Poll<Result<(), Self::Error>> {
|
|
std::task::Poll::Pending
|
|
}
|
|
}
|
|
|
|
fn ordinary_disconnect_reason() -> watch::Receiver<Option<CommunityDisconnectReason>> {
|
|
let (_tx, rx) = watch::channel(None);
|
|
rx
|
|
}
|
|
|
|
/// Every frame the production root `send_loop_inner` writes for a
|
|
/// cancelled `conn` whose terminal receiver is `terminal_rx`: the terminal
|
|
/// drain followed by the close chosen from the connection's disconnect
|
|
/// reason. Shared with `handlers::auth`'s wire-equality tests.
|
|
pub(crate) async fn root_wire_frames(
|
|
conn: &ConnectionState,
|
|
terminal_rx: mpsc::Receiver<WsMessage>,
|
|
) -> Vec<WsMessage> {
|
|
assert!(conn.cancel.is_cancelled(), "denial must cancel the socket");
|
|
let (sink, state) = MockSink::new(None);
|
|
send_loop_inner(
|
|
sink,
|
|
mpsc::channel(1).1,
|
|
mpsc::channel(1).1,
|
|
terminal_rx,
|
|
mpsc::channel(1).1,
|
|
conn.cancel.clone(),
|
|
conn.community_control.disconnect_reason(),
|
|
)
|
|
.await;
|
|
let mut recorded = state.lock().expect("mock sink poisoned");
|
|
std::mem::take(&mut recorded.messages)
|
|
}
|
|
|
|
fn deleted_community_disconnect_reason() -> watch::Receiver<Option<CommunityDisconnectReason>> {
|
|
let (tx, rx) = watch::channel(None);
|
|
tx.send_replace(Some(CommunityDisconnectReason::CommunityDeleted));
|
|
rx
|
|
}
|
|
|
|
fn archived_community_disconnect_reason() -> watch::Receiver<Option<CommunityDisconnectReason>>
|
|
{
|
|
let (tx, rx) = watch::channel(None);
|
|
tx.send_replace(Some(CommunityDisconnectReason::CommunityArchived));
|
|
rx
|
|
}
|
|
|
|
fn text_payloads(messages: &[WsMessage]) -> Vec<String> {
|
|
messages
|
|
.iter()
|
|
.map(|msg| match msg {
|
|
WsMessage::Text(text) => text.to_string(),
|
|
other => panic!("unexpected websocket message in test: {other:?}"),
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
/// Cancellation must break a writer blocked in `poll_ready`, and terminal
|
|
/// close delivery gets one bounded best-effort window before teardown wins.
|
|
#[tokio::test(start_paused = true)]
|
|
async fn cancelled_never_ready_sink_cannot_retain_writer_task() {
|
|
let (data_tx, data_rx) = mpsc::channel(1);
|
|
let (_ctrl_tx, ctrl_rx) = mpsc::channel(1);
|
|
let (_terminal_ctrl_tx, terminal_ctrl_rx) = mpsc::channel::<WsMessage>(1);
|
|
let (_restart_tx, restart_rx) = mpsc::channel(1);
|
|
let cancel = CancellationToken::new();
|
|
let ready_polled = Arc::new(Notify::new());
|
|
data_tx
|
|
.send(WsMessage::Text("blocked".into()))
|
|
.await
|
|
.expect("queue blocked frame");
|
|
|
|
let writer = tokio::spawn(send_loop_inner(
|
|
NeverReadySink {
|
|
ready_polled: Arc::clone(&ready_polled),
|
|
},
|
|
data_rx,
|
|
ctrl_rx,
|
|
terminal_ctrl_rx,
|
|
restart_rx,
|
|
cancel.clone(),
|
|
ordinary_disconnect_reason(),
|
|
));
|
|
|
|
ready_polled.notified().await;
|
|
cancel.cancel();
|
|
tokio::task::yield_now().await;
|
|
tokio::time::advance(WS_TERMINAL_FLUSH_TIMEOUT + Duration::from_millis(1)).await;
|
|
writer
|
|
.await
|
|
.expect("writer exits after bounded terminal flush");
|
|
}
|
|
|
|
/// Cancellation during ordinary traffic with a queued FI denial must not
|
|
/// block indefinitely — even when the sink is never-ready, the bounded
|
|
/// terminal drain in `flush_terminal_frames` enforces the timeout and the
|
|
/// writer task exits.
|
|
///
|
|
/// This tests the merge-integrity fix: the cancel arm in `send_loop_inner`
|
|
/// routes through `flush_terminal_frames` (not unbounded `ws_send.send`),
|
|
/// so a queued FI denial cannot retain the writer task past the deadline.
|
|
///
|
|
/// ## Mutation oracle
|
|
///
|
|
/// Replace the `flush_terminal_frames` call in the cancel arm with a bare
|
|
/// unbounded `ws_send.send(terminal_msg).await` loop → the never-ready sink
|
|
/// blocks indefinitely → `tokio::time::advance` does not help →
|
|
/// `writer.await` times out → test panics.
|
|
#[tokio::test(start_paused = true)]
|
|
async fn cancelled_never_ready_sink_with_queued_fi_denial_exits_within_timeout() {
|
|
let (data_tx, data_rx) = mpsc::channel(1);
|
|
let (_ctrl_tx, ctrl_rx) = mpsc::channel(1);
|
|
let (terminal_ctrl_tx, terminal_ctrl_rx) = mpsc::channel::<WsMessage>(1);
|
|
let (_restart_tx, restart_rx) = mpsc::channel(1);
|
|
let cancel = CancellationToken::new();
|
|
let ready_polled = Arc::new(Notify::new());
|
|
data_tx
|
|
.send(WsMessage::Text("in-flight".into()))
|
|
.await
|
|
.expect("queue in-flight data frame");
|
|
|
|
// Queue a FI denial on the terminal channel — simulates an expiry
|
|
// task firing and queuing a denial just before cancellation.
|
|
let denial = crate::nip_fi_session::denial_frame(
|
|
crate::nip_fi_session::NipFiWsRoute::Root,
|
|
buzz_auth::DenialClass::AuthorizationDenied,
|
|
);
|
|
terminal_ctrl_tx
|
|
.try_send(denial)
|
|
.expect("queue FI denial on terminal_ctrl_tx");
|
|
|
|
let writer = tokio::spawn(send_loop_inner(
|
|
NeverReadySink {
|
|
ready_polled: Arc::clone(&ready_polled),
|
|
},
|
|
data_rx,
|
|
ctrl_rx,
|
|
terminal_ctrl_rx,
|
|
restart_rx,
|
|
cancel.clone(),
|
|
ordinary_disconnect_reason(),
|
|
));
|
|
|
|
ready_polled.notified().await;
|
|
cancel.cancel();
|
|
tokio::task::yield_now().await;
|
|
// Advance past the terminal flush timeout — the writer must exit.
|
|
tokio::time::advance(WS_TERMINAL_FLUSH_TIMEOUT + Duration::from_millis(1)).await;
|
|
writer.await.expect(
|
|
"merge-integrity: writer must exit within bounded terminal flush \
|
|
even when a FI denial is queued and the sink is never-ready.\n\
|
|
Mutation oracle: replace flush_terminal_frames in cancel arm with \
|
|
bare unbounded ws_send.send loop → never-ready sink blocks → \
|
|
advance() does not unblock → task never exits → panic",
|
|
);
|
|
}
|
|
|
|
/// Cancellation during ordinary traffic (no blocked data send, cancel fires in
|
|
/// the select! arm) with a queued FI denial must route through the bounded
|
|
/// terminal-drain path. The writer task must exit within the flush timeout.
|
|
///
|
|
/// This is the "ordinary traffic cancellation" path: the select! `cancel.cancelled()`
|
|
/// arm fires before any data send is in flight. The fix routes it through
|
|
/// `flush_terminal_frames` which includes the terminal_ctrl_rx drain with timeout.
|
|
///
|
|
/// ## Mutation oracle
|
|
///
|
|
/// Remove `&mut terminal_ctrl_rx` from the `flush_terminal_frames` call in the
|
|
/// `cancel.cancelled()` arm → terminal drain skipped → FI denial lost → the denial
|
|
/// frame is not delivered to the sink. (The mock sink check catches the omission if
|
|
/// the sink is observable; for the never-ready-sink variant the timing test catches it.)
|
|
#[tokio::test(start_paused = true)]
|
|
async fn cancellation_during_select_with_fi_denial_routes_through_bounded_path() {
|
|
// Use no queued data — cancel fires directly in the select! arm.
|
|
let (_data_tx, data_rx) = mpsc::channel::<WsMessage>(1);
|
|
let (_ctrl_tx, ctrl_rx) = mpsc::channel(1);
|
|
let (terminal_ctrl_tx, terminal_ctrl_rx) = mpsc::channel::<WsMessage>(1);
|
|
let (_restart_tx, restart_rx) = mpsc::channel(1);
|
|
let cancel = CancellationToken::new();
|
|
let ready_polled = Arc::new(Notify::new());
|
|
|
|
let denial = crate::nip_fi_session::denial_frame(
|
|
crate::nip_fi_session::NipFiWsRoute::Root,
|
|
buzz_auth::DenialClass::AuthorizationDenied,
|
|
);
|
|
terminal_ctrl_tx.try_send(denial).expect("queue FI denial");
|
|
|
|
let writer = tokio::spawn(send_loop_inner(
|
|
NeverReadySink {
|
|
ready_polled: Arc::clone(&ready_polled),
|
|
},
|
|
data_rx,
|
|
ctrl_rx,
|
|
terminal_ctrl_rx,
|
|
restart_rx,
|
|
cancel.clone(),
|
|
ordinary_disconnect_reason(),
|
|
));
|
|
|
|
// Cancel before the select! ever receives any data — fires the cancel arm.
|
|
cancel.cancel();
|
|
tokio::task::yield_now().await;
|
|
tokio::time::advance(WS_TERMINAL_FLUSH_TIMEOUT + Duration::from_millis(1)).await;
|
|
writer.await.expect(
|
|
"merge-integrity: writer must exit from select! cancel arm within bounded flush.\n\
|
|
Mutation oracle: remove &mut terminal_ctrl_rx from flush_terminal_frames in \
|
|
the cancel.cancelled() arm → denial skipped → for never-ready sink, \
|
|
flush_terminal_frames still times out, but the terminal drain path is absent",
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_loop_batches_queued_data_frames_into_one_flush() {
|
|
let (data_tx, data_rx) = mpsc::channel(MAX_WS_SEND_BATCH);
|
|
let (_ctrl_tx, ctrl_rx) = mpsc::channel(1);
|
|
for i in 0..5 {
|
|
data_tx
|
|
.send(WsMessage::Text(format!("data-{i}").into()))
|
|
.await
|
|
.expect("queue data frame");
|
|
}
|
|
|
|
let (sink, state) = MockSink::new(Some(1));
|
|
let (_restart_tx, restart_rx) = mpsc::channel(1);
|
|
send_loop_inner(
|
|
sink,
|
|
data_rx,
|
|
ctrl_rx,
|
|
mpsc::channel(1).1,
|
|
restart_rx,
|
|
CancellationToken::new(),
|
|
ordinary_disconnect_reason(),
|
|
)
|
|
.await;
|
|
|
|
let state = state.lock().expect("mock sink poisoned");
|
|
assert_eq!(state.flush_count, 1);
|
|
assert_eq!(
|
|
text_payloads(&state.messages),
|
|
vec!["data-0", "data-1", "data-2", "data-3", "data-4"]
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_loop_batch_one_preserves_single_frame_flush_behavior() {
|
|
let (data_tx, data_rx) = mpsc::channel(1);
|
|
let (_ctrl_tx, ctrl_rx) = mpsc::channel(1);
|
|
data_tx
|
|
.send(WsMessage::Text("single".into()))
|
|
.await
|
|
.expect("queue data frame");
|
|
|
|
let (sink, state) = MockSink::new(Some(1));
|
|
let (_restart_tx, restart_rx) = mpsc::channel(1);
|
|
send_loop_inner(
|
|
sink,
|
|
data_rx,
|
|
ctrl_rx,
|
|
mpsc::channel(1).1,
|
|
restart_rx,
|
|
CancellationToken::new(),
|
|
ordinary_disconnect_reason(),
|
|
)
|
|
.await;
|
|
|
|
let state = state.lock().expect("mock sink poisoned");
|
|
assert_eq!(state.flush_count, 1);
|
|
assert_eq!(text_payloads(&state.messages), vec!["single"]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_loop_drains_control_before_batched_data_without_reordering() {
|
|
let (data_tx, data_rx) = mpsc::channel(MAX_WS_SEND_BATCH);
|
|
let (ctrl_tx, ctrl_rx) = mpsc::channel(1);
|
|
data_tx
|
|
.send(WsMessage::Text("data-0".into()))
|
|
.await
|
|
.expect("queue data frame");
|
|
data_tx
|
|
.send(WsMessage::Text("data-1".into()))
|
|
.await
|
|
.expect("queue data frame");
|
|
ctrl_tx
|
|
.send(WsMessage::Text("control".into()))
|
|
.await
|
|
.expect("queue control frame");
|
|
|
|
let (sink, state) = MockSink::new(Some(2));
|
|
let (_restart_tx, restart_rx) = mpsc::channel(1);
|
|
send_loop_inner(
|
|
sink,
|
|
data_rx,
|
|
ctrl_rx,
|
|
mpsc::channel(1).1,
|
|
restart_rx,
|
|
CancellationToken::new(),
|
|
ordinary_disconnect_reason(),
|
|
)
|
|
.await;
|
|
|
|
let state = state.lock().expect("mock sink poisoned");
|
|
assert_eq!(state.flush_count, 2);
|
|
assert_eq!(
|
|
text_payloads(&state.messages),
|
|
vec!["control", "data-0", "data-1"]
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_loop_acknowledges_restart_after_flushing_exactly_one_1012() {
|
|
let (_data_tx, data_rx) = mpsc::channel(1);
|
|
let (_ctrl_tx, ctrl_rx) = mpsc::channel(1);
|
|
let (restart_tx, restart_rx) = mpsc::channel(1);
|
|
let (flushed_tx, flushed_rx) = tokio::sync::oneshot::channel();
|
|
restart_tx
|
|
.send(RestartClose {
|
|
flushed: flushed_tx,
|
|
})
|
|
.await
|
|
.expect("queue restart close");
|
|
|
|
let (sink, state) = MockSink::new(None);
|
|
send_loop_inner(
|
|
sink,
|
|
data_rx,
|
|
ctrl_rx,
|
|
mpsc::channel(1).1,
|
|
restart_rx,
|
|
CancellationToken::new(),
|
|
ordinary_disconnect_reason(),
|
|
)
|
|
.await;
|
|
|
|
assert_eq!(flushed_rx.await, Ok(true));
|
|
let state = state.lock().expect("mock sink poisoned");
|
|
assert_eq!(state.flush_count, 1, "ack follows the close flush");
|
|
assert_eq!(state.messages.len(), 1, "writer exits after restart close");
|
|
match &state.messages[0] {
|
|
WsMessage::Close(Some(close)) => {
|
|
assert_eq!(close.code, axum::extract::ws::close_code::RESTART);
|
|
assert_eq!(close.reason.as_str(), "relay restarting");
|
|
}
|
|
other => panic!("expected one 1012 restart close, got {other:?}"),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_loop_reports_restart_flush_failure() {
|
|
let (_data_tx, data_rx) = mpsc::channel(1);
|
|
let (_ctrl_tx, ctrl_rx) = mpsc::channel(1);
|
|
let (restart_tx, restart_rx) = mpsc::channel(1);
|
|
let (flushed_tx, flushed_rx) = tokio::sync::oneshot::channel();
|
|
restart_tx
|
|
.send(RestartClose {
|
|
flushed: flushed_tx,
|
|
})
|
|
.await
|
|
.expect("queue restart close");
|
|
|
|
let (sink, state) = MockSink::new(Some(1));
|
|
send_loop_inner(
|
|
sink,
|
|
data_rx,
|
|
ctrl_rx,
|
|
mpsc::channel(1).1,
|
|
restart_rx,
|
|
CancellationToken::new(),
|
|
ordinary_disconnect_reason(),
|
|
)
|
|
.await;
|
|
|
|
assert_eq!(flushed_rx.await, Ok(false));
|
|
let state = state.lock().expect("mock sink poisoned");
|
|
assert_eq!(state.flush_count, 1);
|
|
assert_eq!(state.messages.len(), 1, "no fallback close is appended");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_loop_sends_policy_close_when_community_is_deleted() {
|
|
let (_data_tx, data_rx) = mpsc::channel(1);
|
|
let (_ctrl_tx, ctrl_rx) = mpsc::channel(1);
|
|
let (_restart_tx, restart_rx) = mpsc::channel(1);
|
|
let cancel = CancellationToken::new();
|
|
cancel.cancel();
|
|
|
|
let (sink, state) = MockSink::new(None);
|
|
send_loop_inner(
|
|
sink,
|
|
data_rx,
|
|
ctrl_rx,
|
|
mpsc::channel(1).1,
|
|
restart_rx,
|
|
cancel,
|
|
deleted_community_disconnect_reason(),
|
|
)
|
|
.await;
|
|
|
|
let state = state.lock().expect("mock sink poisoned");
|
|
assert_eq!(state.messages.len(), 1);
|
|
match &state.messages[0] {
|
|
WsMessage::Close(Some(close)) => {
|
|
assert_eq!(close.code, axum::extract::ws::close_code::POLICY);
|
|
assert_eq!(close.reason.as_str(), "community deleted");
|
|
}
|
|
other => panic!("expected one 1008 deletion close, got {other:?}"),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_loop_sends_policy_close_when_community_is_archived() {
|
|
let (_data_tx, data_rx) = mpsc::channel(1);
|
|
let (_ctrl_tx, ctrl_rx) = mpsc::channel(1);
|
|
let (_restart_tx, restart_rx) = mpsc::channel(1);
|
|
let cancel = CancellationToken::new();
|
|
cancel.cancel();
|
|
|
|
let (sink, state) = MockSink::new(None);
|
|
send_loop_inner(
|
|
sink,
|
|
data_rx,
|
|
ctrl_rx,
|
|
mpsc::channel(1).1,
|
|
restart_rx,
|
|
cancel,
|
|
archived_community_disconnect_reason(),
|
|
)
|
|
.await;
|
|
|
|
let state = state.lock().expect("mock sink poisoned");
|
|
assert_eq!(state.messages.len(), 1);
|
|
match &state.messages[0] {
|
|
WsMessage::Close(Some(close)) => {
|
|
assert_eq!(close.code, axum::extract::ws::close_code::POLICY);
|
|
assert_eq!(close.reason.as_str(), "community archived");
|
|
}
|
|
other => panic!("expected one 1008 archive close, got {other:?}"),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_loop_sends_bare_close_for_ordinary_cancellation() {
|
|
let (_data_tx, data_rx) = mpsc::channel(1);
|
|
let (_ctrl_tx, ctrl_rx) = mpsc::channel(1);
|
|
let (_restart_tx, restart_rx) = mpsc::channel(1);
|
|
let cancel = CancellationToken::new();
|
|
cancel.cancel();
|
|
|
|
let (sink, state) = MockSink::new(None);
|
|
send_loop_inner(
|
|
sink,
|
|
data_rx,
|
|
ctrl_rx,
|
|
mpsc::channel(1).1,
|
|
restart_rx,
|
|
cancel,
|
|
ordinary_disconnect_reason(),
|
|
)
|
|
.await;
|
|
|
|
let state = state.lock().expect("mock sink poisoned");
|
|
assert_eq!(state.messages.as_slice(), [WsMessage::Close(None)]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_loop_flushes_queued_control_before_close_on_cancel() {
|
|
// A ban disconnect queues its `OK false "blocked: …"` reason frame on
|
|
// the control channel and then cancels the token (B3). The biased
|
|
// select polls the cancel branch first, so the reason frame would be
|
|
// stranded unless the cancel branch drains ctrl before emitting Close.
|
|
// This test exercises `send_loop_inner` end-to-end to prove the reason
|
|
// frame reaches the client, in order, ahead of the Close.
|
|
let (_data_tx, data_rx) = mpsc::channel(1);
|
|
let (ctrl_tx, ctrl_rx) = mpsc::channel(1);
|
|
ctrl_tx
|
|
.send(WsMessage::Text("blocked: you are banned".into()))
|
|
.await
|
|
.expect("queue ban reason frame");
|
|
|
|
let cancel = CancellationToken::new();
|
|
cancel.cancel();
|
|
|
|
let (sink, state) = MockSink::new(None);
|
|
let (_restart_tx, restart_rx) = mpsc::channel(1);
|
|
send_loop_inner(
|
|
sink,
|
|
data_rx,
|
|
ctrl_rx,
|
|
mpsc::channel(1).1,
|
|
restart_rx,
|
|
cancel,
|
|
ordinary_disconnect_reason(),
|
|
)
|
|
.await;
|
|
|
|
let state = state.lock().expect("mock sink poisoned");
|
|
assert_eq!(
|
|
state.messages.len(),
|
|
2,
|
|
"reason frame then Close, nothing else"
|
|
);
|
|
match &state.messages[0] {
|
|
WsMessage::Text(text) => {
|
|
assert_eq!(text.as_str(), "blocked: you are banned")
|
|
}
|
|
other => panic!("expected the ban reason frame first, got {other:?}"),
|
|
}
|
|
assert!(
|
|
matches!(state.messages[1], WsMessage::Close(None)),
|
|
"ordinary cancellation retains the bare Close after the reason frame"
|
|
);
|
|
}
|
|
|
|
// ── NIP-FI session deadline — production function falsifiability ──────────
|
|
//
|
|
// These tests call `compute_session_deadline` directly (the production path
|
|
// used by `handle_connection`) with real `VerifiedAssertion` fixtures.
|
|
// Deleting or mutating `compute_session_deadline` turns these red.
|
|
|
|
#[test]
|
|
fn deadline_exp_is_earliest_selects_exp() {
|
|
use buzz_auth::VerifiedAssertion;
|
|
use chrono::{Duration, Utc};
|
|
|
|
let now = Utc::now();
|
|
let exp = now + Duration::seconds(100);
|
|
let iat_max_age = now + Duration::seconds(300);
|
|
let key_hard = now + Duration::seconds(200);
|
|
// authority_deadlines = [exp, iat_max_age, key_hard] → min = exp
|
|
let assertion = VerifiedAssertion::for_test(None, vec![exp, iat_max_age, key_hard]);
|
|
let lifetime = std::time::Duration::from_secs(400);
|
|
let deadline = compute_session_deadline(&assertion, now, Some(lifetime));
|
|
// exp < key_hard < lifetime; upstream = exp, partition >> exp → exp wins.
|
|
assert_eq!(deadline, exp, "exp is earliest upstream term");
|
|
}
|
|
|
|
#[test]
|
|
fn deadline_max_connection_lifetime_is_earliest_selects_partition() {
|
|
use buzz_auth::VerifiedAssertion;
|
|
use chrono::{Duration, Utc};
|
|
|
|
let now = Utc::now();
|
|
let exp = now + Duration::seconds(400);
|
|
let iat_max_age = now + Duration::seconds(300);
|
|
let key_hard = now + Duration::seconds(200);
|
|
// authority_deadlines = [exp, iat_max_age, key_hard] → upstream = key_hard (200s)
|
|
// lifetime partition = now + 100s < key_hard → partition wins.
|
|
let assertion = VerifiedAssertion::for_test(None, vec![exp, iat_max_age, key_hard]);
|
|
let lifetime = std::time::Duration::from_secs(100);
|
|
let deadline = compute_session_deadline(&assertion, now, Some(lifetime));
|
|
// partition (now+100s) < upstream (now+200s) → partition wins.
|
|
let expected_partition = now + Duration::seconds(100);
|
|
// Allow 1s of wall-clock slack in the test.
|
|
let delta = if deadline > expected_partition {
|
|
(deadline - expected_partition).num_milliseconds().abs()
|
|
} else {
|
|
(expected_partition - deadline).num_milliseconds().abs()
|
|
};
|
|
assert!(delta < 1000, "partition term should win; delta={delta}ms");
|
|
}
|
|
|
|
#[test]
|
|
fn deadline_no_lifetime_returns_upstream_only() {
|
|
use buzz_auth::VerifiedAssertion;
|
|
use chrono::{Duration, Utc};
|
|
|
|
let now = Utc::now();
|
|
let exp = now + Duration::seconds(600);
|
|
let key_hard = now + Duration::seconds(3600);
|
|
let assertion = VerifiedAssertion::for_test(None, vec![exp, key_hard]);
|
|
let deadline = compute_session_deadline(&assertion, now, None);
|
|
assert_eq!(deadline, exp, "no lifetime → upstream (exp) only");
|
|
}
|
|
|
|
// ── NIP-FI expiry notice delivered on terminal_ctrl_tx before cancel ─────
|
|
//
|
|
// The expiry task queues `restricted: authorization denied` on
|
|
// `terminal_ctrl_tx` (capacity-1, prioritised) BEFORE cancellation via the
|
|
// gate. This test invokes the production
|
|
// `nip_fi_session::spawn_nip_fi_expiry_task` constructor (Root route):
|
|
// an already-expired deadline fires immediately; the terminal channel carries
|
|
// the denial frame; the cancel fires afterward.
|
|
//
|
|
// Mutation evidence:
|
|
// A) Change the enqueue in `spawn_nip_fi_expiry_task` back to `ctrl_tx` →
|
|
// `terminal_rx.try_recv()` returns `Err`; test panics at "terminal
|
|
// channel must contain the denial frame".
|
|
// B) Delete `cancel.cancel()` inside gate.expire() →
|
|
// `cancel.is_cancelled()` is false; test panics at "expiry task must
|
|
// cancel the connection".
|
|
|
|
#[tokio::test]
|
|
async fn expiry_notice_queued_on_ctrl_before_cancel() {
|
|
use tokio::sync::mpsc;
|
|
|
|
let (terminal_ctrl_tx, mut terminal_ctrl_rx) = mpsc::channel::<WsMessage>(1);
|
|
let cancel = CancellationToken::new();
|
|
|
|
// Already-expired deadline → fires immediately.
|
|
let deadline = chrono::Utc::now() - chrono::Duration::seconds(10);
|
|
|
|
let gate = crate::nip_fi_gate::SessionAdmissionGate::new(deadline, cancel.clone());
|
|
|
|
// Invoke the production shared constructor — Root route.
|
|
let expiry_task = crate::nip_fi_session::spawn_nip_fi_expiry_task(
|
|
deadline,
|
|
gate,
|
|
crate::state::CommunityConnectionControl::new(cancel.clone()),
|
|
terminal_ctrl_tx,
|
|
crate::nip_fi_session::NipFiWsRoute::Root,
|
|
);
|
|
|
|
tokio::time::timeout(std::time::Duration::from_secs(2), expiry_task)
|
|
.await
|
|
.expect("expiry task must complete within 2s")
|
|
.expect("expiry task must not panic");
|
|
|
|
// terminal_ctrl_rx must contain the denial frame.
|
|
let terminal_frame = terminal_ctrl_rx
|
|
.try_recv()
|
|
.expect("terminal channel must contain the denial frame before cancel");
|
|
match terminal_frame {
|
|
WsMessage::Text(text) => {
|
|
// Root route: NOTICE format ["NOTICE", <message>].
|
|
let v: serde_json::Value = serde_json::from_str(&text).expect("valid JSON");
|
|
let payload = v.get(1).and_then(|c| c.as_str()).unwrap_or("");
|
|
assert_eq!(
|
|
payload,
|
|
buzz_auth::DenialClass::AuthorizationDenied.nostr_text(),
|
|
"terminal frame must carry the exact authorization_denied text"
|
|
);
|
|
}
|
|
other => panic!("terminal frame must be Text, got {other:?}"),
|
|
}
|
|
// Cancel must have fired after the terminal send.
|
|
assert!(
|
|
cancel.is_cancelled(),
|
|
"expiry task must cancel the connection"
|
|
);
|
|
}
|
|
|
|
// ── B2: frame-admission fence and AUTH TOCTOU ─────────────────────────────
|
|
//
|
|
// Once the NIP-FI expiry task calls cancel(), no further message dispatch
|
|
// should occur — even if a frame was already buffered in the recv queue
|
|
// before cancel fired.
|
|
//
|
|
// The fence is the `if conn.cancel.is_cancelled() { return; }` check at the
|
|
// top of `handle_text_message`. These tests exercise two windows:
|
|
//
|
|
// 1. A buffered REQ/EVENT/COUNT frame that arrives after cancel fires.
|
|
// 2. An AUTH message dispatched while cancel is already set
|
|
// (the TOCTOU window where auth_state.write() is acquired, cancel is
|
|
// checked under the lock, and the write is skipped if cancelled).
|
|
//
|
|
// Mutation evidence:
|
|
// A) Remove `if conn.cancel.is_cancelled() { return; }` from
|
|
// `handle_text_message` → the EVENT test receives a frame on send_rx
|
|
// (an OK or NOTICE) → the assertion panics.
|
|
// B) Remove `if conn.cancel.is_cancelled() { return; }` from the AUTH
|
|
// handler (inside the write guard) → the AUTH test's
|
|
// `not Authenticated` assertion may still hold due to the DB path, but
|
|
// the top-level handle_text_message fence is the true gate.
|
|
|
|
#[tokio::test]
|
|
async fn b2_cancelled_connection_event_frame_not_dispatched() {
|
|
use std::collections::HashMap;
|
|
use tokio::sync::mpsc;
|
|
|
|
// Pre-cancel the token — simulates the expiry task having already fired.
|
|
let cancel = CancellationToken::new();
|
|
cancel.cancel();
|
|
|
|
let (send_tx, mut send_rx) = mpsc::channel::<WsMessage>(8);
|
|
let (ctrl_tx, _ctrl_rx) = mpsc::channel::<WsMessage>(8);
|
|
let (terminal_ctrl_tx, _terminal_ctrl_rx) = mpsc::channel::<WsMessage>(1);
|
|
|
|
let conn = Arc::new(ConnectionState {
|
|
conn_id: uuid::Uuid::new_v4(),
|
|
tenant: buzz_core::tenant::TenantContext::resolved(
|
|
buzz_core::tenant::CommunityId::from_uuid(uuid::Uuid::nil()),
|
|
"test.local".to_string(),
|
|
),
|
|
remote_addr: "127.0.0.1:1234".parse().unwrap(),
|
|
auth_state: StdMutex::new(AuthState::Failed),
|
|
subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
|
|
send_tx,
|
|
ctrl_tx,
|
|
terminal_ctrl_tx,
|
|
cancel: cancel.clone(),
|
|
backpressure_count: Arc::new(std::sync::atomic::AtomicU8::new(0)),
|
|
grace_limit: 3,
|
|
nip_fi_assertion: None,
|
|
session_deadline: None,
|
|
nip_fi_gate: crate::nip_fi_gate::SessionAdmissionGate::off_mode(cancel.clone()),
|
|
community_control: crate::state::CommunityConnectionControl::new(cancel.clone()),
|
|
});
|
|
|
|
let state = crate::state::tests::test_state().await;
|
|
// A plausible EVENT frame — the handler would normally send OK/NOTICE.
|
|
let event = nostr::EventBuilder::new(nostr::Kind::TextNote, "b2 test")
|
|
.sign_with_keys(&Keys::generate())
|
|
.unwrap();
|
|
let raw = serde_json::json!(["EVENT", event]).to_string();
|
|
|
|
handle_text_message(raw, Arc::clone(&conn), Arc::clone(&state)).await;
|
|
|
|
// No frame must be sent — the fence must return before any handler runs.
|
|
assert!(
|
|
send_rx.try_recv().is_err(),
|
|
"B2: a pre-cancelled connection must not dispatch an EVENT frame to any handler"
|
|
);
|
|
}
|
|
|
|
// ── B3: send_loop writer delivers denial-then-Close through real send path ─
|
|
//
|
|
// These tests drive the real `send_loop_inner` against a sink that records
|
|
// every frame, saturate ctrl_tx, enqueue a denial frame on terminal_ctrl_tx,
|
|
// then cancel the token. The sink is non-blocking (MockSink), so send_loop
|
|
// runs to completion synchronously after cancel fires.
|
|
//
|
|
// Assertion: the denial frame appears in the output BEFORE the Close frame.
|
|
// This proves the queue-then-cancel ordering holds through the actual writer
|
|
// code path, not just through a channel try_recv check.
|
|
//
|
|
// Mutation evidence:
|
|
// A) In send_loop_inner's cancel branch, swap the terminal drain and the
|
|
// ctrl drain → denial frame position flips → assertion panics.
|
|
// B) Remove the terminal drain entirely → denial frame absent → assertion
|
|
// panics on the "denial frame must precede Close" check.
|
|
|
|
#[tokio::test]
|
|
async fn b3_root_pairing_denial_precedes_close_through_send_loop() {
|
|
use crate::nip_fi_session::NipFiWsRoute;
|
|
|
|
let (data_tx, data_rx) = mpsc::channel::<WsMessage>(16);
|
|
let (ctrl_tx, ctrl_rx) = mpsc::channel::<WsMessage>(8);
|
|
let (terminal_ctrl_tx, terminal_ctrl_rx) = mpsc::channel::<WsMessage>(1);
|
|
let (_restart_tx, restart_rx) = mpsc::channel(1);
|
|
let cancel = CancellationToken::new();
|
|
|
|
// Saturate ctrl_tx so an ordinary send couldn't carry the denial frame.
|
|
for i in 0..8u8 {
|
|
ctrl_tx
|
|
.try_send(WsMessage::Text(format!("ordinary-{i}").into()))
|
|
.expect("ctrl_tx has capacity 8");
|
|
}
|
|
drop(data_tx); // no data traffic in this test
|
|
|
|
// Enqueue the denial frame on the terminal channel, then cancel.
|
|
// This is the queue-then-cancel pattern the pairing denial path uses.
|
|
terminal_ctrl_tx
|
|
.try_send(crate::nip_fi_session::denial_frame(
|
|
NipFiWsRoute::Root,
|
|
buzz_auth::DenialClass::AuthorizationDenied,
|
|
))
|
|
.expect("terminal channel is empty");
|
|
cancel.cancel();
|
|
|
|
let (sink, state_arc) = MockSink::new(None);
|
|
send_loop_inner(
|
|
sink,
|
|
data_rx,
|
|
ctrl_rx,
|
|
terminal_ctrl_rx,
|
|
restart_rx,
|
|
cancel,
|
|
ordinary_disconnect_reason(),
|
|
)
|
|
.await;
|
|
|
|
let state = state_arc.lock().expect("mock sink poisoned");
|
|
// The first frame written must be the denial frame.
|
|
// The last frame written must be Close (or None close).
|
|
let msgs = &state.messages;
|
|
assert!(
|
|
!msgs.is_empty(),
|
|
"send_loop must write at least the denial frame + Close"
|
|
);
|
|
// Find the denial frame.
|
|
let denial_pos = msgs
|
|
.iter()
|
|
.position(|m| matches!(m, WsMessage::Text(t) if t.contains("authorization denied")));
|
|
let close_pos = msgs.iter().rposition(|m| matches!(m, WsMessage::Close(_)));
|
|
|
|
let denial_pos = denial_pos.expect("denial frame must appear in send_loop output");
|
|
let close_pos = close_pos.expect("Close frame must appear in send_loop output");
|
|
assert!(
|
|
denial_pos < close_pos,
|
|
"B3: denial frame (pos {denial_pos}) must precede Close frame (pos {close_pos})"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn b3_expiry_denial_precedes_close_through_send_loop() {
|
|
use crate::nip_fi_session::{spawn_nip_fi_expiry_task, NipFiWsRoute};
|
|
use chrono::Utc;
|
|
|
|
let (data_tx, data_rx) = mpsc::channel::<WsMessage>(16);
|
|
let (ctrl_tx, ctrl_rx) = mpsc::channel::<WsMessage>(8);
|
|
let (terminal_ctrl_tx, terminal_ctrl_rx) = mpsc::channel::<WsMessage>(1);
|
|
let (_restart_tx, restart_rx) = mpsc::channel(1);
|
|
let cancel = CancellationToken::new();
|
|
|
|
// Saturate ctrl_tx.
|
|
for i in 0..8u8 {
|
|
ctrl_tx
|
|
.try_send(WsMessage::Text(format!("ordinary-{i}").into()))
|
|
.expect("ctrl_tx has capacity 8");
|
|
}
|
|
drop(data_tx);
|
|
|
|
// Arm the expiry task with an already-expired deadline. It will
|
|
// immediately enqueue the denial frame on the terminal channel and
|
|
// cancel the token.
|
|
let already_expired = Utc::now() - chrono::Duration::seconds(1);
|
|
let gate = crate::nip_fi_gate::SessionAdmissionGate::new(already_expired, cancel.clone());
|
|
let control = crate::state::CommunityConnectionControl::new(cancel.clone());
|
|
let disconnect_reason = control.disconnect_reason();
|
|
let expiry_handle = spawn_nip_fi_expiry_task(
|
|
already_expired,
|
|
gate,
|
|
control,
|
|
terminal_ctrl_tx,
|
|
NipFiWsRoute::Root,
|
|
);
|
|
// Wait for the expiry task to fire before we run the send_loop.
|
|
expiry_handle.await.expect("expiry task must complete");
|
|
|
|
let (sink, state_arc) = MockSink::new(None);
|
|
send_loop_inner(
|
|
sink,
|
|
data_rx,
|
|
ctrl_rx,
|
|
terminal_ctrl_rx,
|
|
restart_rx,
|
|
cancel,
|
|
disconnect_reason,
|
|
)
|
|
.await;
|
|
|
|
let state = state_arc.lock().expect("mock sink poisoned");
|
|
let msgs = &state.messages;
|
|
assert!(
|
|
!msgs.is_empty(),
|
|
"send_loop must write at least the denial frame + Close"
|
|
);
|
|
let denial_pos = msgs
|
|
.iter()
|
|
.position(|m| matches!(m, WsMessage::Text(t) if t.contains("authorization denied")));
|
|
let close_pos = msgs.iter().rposition(|m| matches!(m, WsMessage::Close(_)));
|
|
|
|
let denial_pos = denial_pos.expect("expiry denial frame must appear in send_loop output");
|
|
let close_pos = close_pos.expect("Close frame must appear in send_loop output");
|
|
assert!(
|
|
denial_pos < close_pos,
|
|
"B3: expiry denial frame (pos {denial_pos}) must precede Close frame (pos {close_pos})"
|
|
);
|
|
// Expiry is `authorization_denied`: the same 1008 as every other row.
|
|
assert_eq!(
|
|
msgs[close_pos],
|
|
CommunityDisconnectReason::AuthorizationDenied.close_message()
|
|
);
|
|
}
|
|
|
|
// ── F3: bootstrap deadline witness — root WS route ──────────────────────────
|
|
//
|
|
// Fix 3: the NIP-FI gate and expiry task are created BEFORE the
|
|
// `is_community_active` bootstrap await in `handle_connection`, so a
|
|
// session deadline that fires during a slow DB check still terminates the
|
|
// root WebSocket connection on time.
|
|
//
|
|
// This test passes a pre-built already-expired gate + a pre-fired expiry
|
|
// task to `handle_active_connection` via `pre_built = Some(...)`.
|
|
// The expiry task fires immediately (deadline in the past), queues the
|
|
// denial NOTICE on the terminal channel, and cancels the token.
|
|
// The test asserts the client receives the auth challenge, then the
|
|
// authorization-denied NOTICE, all within 2 s.
|
|
//
|
|
// Mutation oracle:
|
|
// Remove the `if let Some((gate, tx, rx, expiry)) = pre_built` branch
|
|
// from `handle_active_connection` (always use the else branch) → the
|
|
// pre-built terminal channel carrying the denial frame is discarded →
|
|
// the send_loop drains a fresh (empty) terminal channel → denial NOTICE
|
|
// never appears → `received_denial` stays false → assertion panics.
|
|
#[tokio::test]
|
|
async fn f3_root_pre_built_expired_gate_terminates_connection() {
|
|
use axum::{routing::get, Router};
|
|
use buzz_auth::VerifiedAssertion;
|
|
use chrono::{Duration, Utc};
|
|
use std::sync::Arc;
|
|
use tokio::net::TcpListener;
|
|
use tokio_tungstenite::connect_async;
|
|
|
|
let key = nostr::Keys::generate();
|
|
|
|
// Deadline in the past — expiry fires immediately on spawn.
|
|
let deadline = Utc::now() - Duration::seconds(1);
|
|
let assertion = VerifiedAssertion::for_test(Some(key.public_key()), vec![deadline]);
|
|
|
|
let state = crate::state::tests::test_state().await;
|
|
let tenant = TenantContext::resolved(
|
|
buzz_core::tenant::CommunityId::from_uuid(uuid::Uuid::nil()),
|
|
"test.local".to_string(),
|
|
);
|
|
|
|
let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::<()>();
|
|
let state_c = Arc::clone(&state);
|
|
let tenant_c = tenant.clone();
|
|
let assertion_c = assertion.clone();
|
|
|
|
let listener = TcpListener::bind("127.0.0.1:0")
|
|
.await
|
|
.expect("F3-root: bind listener");
|
|
let addr = listener.local_addr().expect("F3-root: local addr");
|
|
|
|
let server = tokio::spawn(async move {
|
|
let app = Router::new().route(
|
|
"/",
|
|
get({
|
|
let state_i = Arc::clone(&state_c);
|
|
let tenant_i = tenant_c.clone();
|
|
let assertion_i = assertion_c.clone();
|
|
move |ws: axum::extract::ws::WebSocketUpgrade| {
|
|
let state_i = Arc::clone(&state_i);
|
|
let tenant_i = tenant_i.clone();
|
|
let assertion_i = assertion_i.clone();
|
|
let conn_time = chrono::Utc::now();
|
|
let conn_id = uuid::Uuid::new_v4();
|
|
async move {
|
|
ws.on_upgrade(move |socket| async move {
|
|
let conn_cancel = CancellationToken::new();
|
|
|
|
// Build pre-built bundle: gate + expiry pre-fired
|
|
// (simulates handle_connection arming before bootstrap).
|
|
let (pre_tx, pre_rx) = mpsc::channel::<WsMessage>(1);
|
|
let pre_gate = crate::nip_fi_gate::SessionAdmissionGate::new(
|
|
deadline,
|
|
conn_cancel.clone(),
|
|
);
|
|
let control =
|
|
crate::state::CommunityConnectionControl::new(conn_cancel);
|
|
let pre_expiry = crate::nip_fi_session::spawn_nip_fi_expiry_task(
|
|
deadline,
|
|
Arc::clone(&pre_gate),
|
|
control.clone(),
|
|
pre_tx.clone(),
|
|
crate::nip_fi_session::NipFiWsRoute::Root,
|
|
);
|
|
handle_active_connection(
|
|
socket,
|
|
state_i,
|
|
"127.0.0.1:9999".parse().unwrap(),
|
|
tenant_i,
|
|
conn_id,
|
|
control,
|
|
Some(assertion_i),
|
|
conn_time,
|
|
Some((pre_gate, pre_tx, pre_rx, Some(pre_expiry))),
|
|
)
|
|
.await
|
|
})
|
|
}
|
|
}
|
|
}),
|
|
);
|
|
|
|
let _ = ready_tx.send(());
|
|
axum::serve(listener, app).await.expect("test server");
|
|
});
|
|
|
|
let _ = tokio::time::timeout(std::time::Duration::from_secs(2), ready_rx)
|
|
.await
|
|
.expect("F3-root: server ready");
|
|
|
|
let (mut client, _) = connect_async(format!("ws://{addr}/"))
|
|
.await
|
|
.expect("F3-root: connect");
|
|
|
|
// The root WS path sends an AUTH challenge first, then the pre-fired
|
|
// expiry drains through the send_loop as authorization-denied NOTICE.
|
|
let mut received_denial = false;
|
|
|
|
for _ in 0..5 {
|
|
let frame =
|
|
tokio::time::timeout(std::time::Duration::from_secs(2), client.next()).await;
|
|
|
|
match frame {
|
|
Ok(Some(Ok(tokio_tungstenite::tungstenite::Message::Text(t))))
|
|
if t.contains("authorization denied") =>
|
|
{
|
|
received_denial = true;
|
|
break;
|
|
}
|
|
Ok(Some(Ok(tokio_tungstenite::tungstenite::Message::Close(_))))
|
|
| Ok(Some(Err(_)))
|
|
| Ok(None)
|
|
| Err(_) => break,
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
assert!(
|
|
received_denial,
|
|
"F3-root: must receive authorization denied NOTICE before close;\n Mutation oracle: remove pre_built branch from handle_active_connection \
|
|
→ pre-built terminal channel discarded → denial never sent → \
|
|
assertion panics"
|
|
);
|
|
|
|
server.abort();
|
|
let _ = server.await;
|
|
}
|
|
|
|
/// Fix 3 (F3): bootstrap-drain through the real outer wrapper (`handle_connection`).
|
|
///
|
|
/// The `run_registered_community_connection` wrapper in `handle_connection` provides
|
|
/// an `on_not_run` closure that drains the pre-terminal channel and closes the socket
|
|
/// when the community-active check fails or cancellation fires during bootstrap.
|
|
///
|
|
/// Scenario: FI assertion with past deadline → expiry task fires immediately and
|
|
/// cancels the token. The community-active check never completes (lazy pool, no DB).
|
|
/// The `on_not_run` path drains the denial frame through the real WebSocket.
|
|
///
|
|
/// ## Mutation oracle
|
|
///
|
|
/// Replace the `on_not_run` closure body with `|| async {}` → the socket is
|
|
/// dropped without sending the denial → client receives only Close → assertion panics.
|
|
#[tokio::test]
|
|
async fn f3_root_outer_wrapper_delivers_denial_on_bootstrap_cancellation() {
|
|
use axum::{routing::get, Router};
|
|
use buzz_auth::VerifiedAssertion;
|
|
use chrono::{Duration, Utc};
|
|
use futures_util::StreamExt as _;
|
|
use std::sync::Arc;
|
|
use tokio::net::TcpListener;
|
|
use tokio_tungstenite::connect_async;
|
|
|
|
// Past deadline → expiry fires immediately; cancellation beats any DB check.
|
|
let key = nostr::Keys::generate();
|
|
let deadline = Utc::now() - Duration::seconds(2);
|
|
let assertion = VerifiedAssertion::for_test(Some(key.public_key()), vec![deadline]);
|
|
|
|
let state = crate::state::tests::test_state().await;
|
|
let tenant = TenantContext::resolved(
|
|
buzz_core::tenant::CommunityId::from_uuid(uuid::Uuid::nil()),
|
|
"test.local".to_string(),
|
|
);
|
|
|
|
let listener = TcpListener::bind("127.0.0.1:0")
|
|
.await
|
|
.expect("F3-outer-root: bind listener");
|
|
let addr = listener.local_addr().expect("F3-outer-root: local addr");
|
|
let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::<()>();
|
|
|
|
let state_c = Arc::clone(&state);
|
|
let tenant_c = tenant.clone();
|
|
let assertion_c = assertion.clone();
|
|
|
|
let server = tokio::spawn(async move {
|
|
let app = Router::new().route(
|
|
"/",
|
|
get({
|
|
let state_i = Arc::clone(&state_c);
|
|
let tenant_i = tenant_c.clone();
|
|
let assertion_i = assertion_c.clone();
|
|
move |ws: axum::extract::ws::WebSocketUpgrade| {
|
|
let state_i = Arc::clone(&state_i);
|
|
let tenant_i = tenant_i.clone();
|
|
let assertion_i = assertion_i.clone();
|
|
let conn_time = chrono::Utc::now();
|
|
async move {
|
|
ws.on_upgrade(move |socket| async move {
|
|
// Call the REAL outer wrapper — includes
|
|
// run_registered_community_connection with its
|
|
// on_not_run drain closure.
|
|
handle_connection(
|
|
socket,
|
|
state_i,
|
|
"127.0.0.1:9999".parse().unwrap(),
|
|
tenant_i,
|
|
Some(assertion_i),
|
|
None,
|
|
conn_time,
|
|
)
|
|
.await
|
|
})
|
|
}
|
|
}
|
|
}),
|
|
);
|
|
let _ = ready_tx.send(());
|
|
axum::serve(listener, app).await.expect("test server");
|
|
});
|
|
|
|
let _ = tokio::time::timeout(std::time::Duration::from_secs(2), ready_rx)
|
|
.await
|
|
.expect("F3-outer-root: server ready");
|
|
|
|
let (mut client, _) = connect_async(format!("ws://{addr}/"))
|
|
.await
|
|
.expect("F3-outer-root: connect");
|
|
|
|
// Expect the authorization-denied NOTICE before the close.
|
|
// The server sends an AUTH challenge first; then when expiry fires the
|
|
// on_not_run closure delivers the denial frame before dropping the socket.
|
|
let mut received_denial = false;
|
|
for _ in 0..8 {
|
|
let frame =
|
|
tokio::time::timeout(std::time::Duration::from_secs(3), client.next()).await;
|
|
match frame {
|
|
Ok(Some(Ok(tokio_tungstenite::tungstenite::Message::Text(t))))
|
|
if t.contains("authorization denied") =>
|
|
{
|
|
received_denial = true;
|
|
break;
|
|
}
|
|
Ok(Some(Ok(tokio_tungstenite::tungstenite::Message::Close(_))))
|
|
| Ok(Some(Err(_)))
|
|
| Ok(None)
|
|
| Err(_) => break,
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
assert!(
|
|
received_denial,
|
|
"F3-outer-root: on_not_run must drain and deliver the FI denial frame \
|
|
before the socket is dropped.\n\
|
|
Mutation oracle: replace the on_not_run closure body with `|| async {{}}` \
|
|
→ socket dropped without drain → client sees only Close → assertion panics"
|
|
);
|
|
|
|
server.abort();
|
|
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<tokio::net::TcpStream>,
|
|
>;
|
|
|
|
/// 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<AppState>, 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<Arc<crate::handlers::close::test_seam::Pause>>,
|
|
/// 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<tokio::sync::oneshot::Sender<crate::nip_fi_gate::SessionEffectPermit>>,
|
|
/// 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<super::test_seam::Hold>,
|
|
}
|
|
|
|
/// 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<AppState>, 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::<WsMessage>(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<String>, 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<AppState>,
|
|
assertion: Option<VerifiedAssertion>,
|
|
) -> impl std::future::Future<Output = std::net::SocketAddr> {
|
|
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"
|
|
);
|
|
}
|
|
}
|