arcature 2026.2.1

//! Typed WebSocket wrapper over `axum::extract::ws` (PROGRAM.md §AP2.1-8).
//!
//! A thin typed layer that adds the realtime safety core the raw extractor
//! does not provide: explicit origin policy, explicit per-connection
//! authorization (channel names never implicitly authorize), connection-
//! limit enforcement, frame/message size limits, bounded-buffer
//! backpressure (via the broadcast channel), heartbeat, graceful shutdown,
//! and stable tracing spans. Raw `axum::extract::ws` remains a first-class
//! escape hatch (AGENTS.md §16): every entry point the wrapper uses is a
//! public, documented seam, and an application can call
//! `WebSocketUpgrade::on_upgrade` directly and skip this wrapper.
//!
//! # No proprietary protocol
//!
//! The wrapper transports the application's serialized bytes (see
//! [`crate::realtime::channel::ChannelPayload`]); it owns no wire format,
//! no `X-Arcature-*` header, and no page protocol (AGENTS.md §15). A stock
//! browser WebSocket client works without knowing Arcature exists.
//!
//! # Lifecycle / reconnect
//!
//! The server does not own client reconnect. The contract: the client owns
//! its reconnect/backoff strategy; the server owns admission (origin/auth/
//! limit), one live broadcast subscription per connection, a heartbeat to
//! detect dead peers, and a graceful `1001` (Going Away) close on server
//! drain. The client should expect a clean close on server drain and
//! reconnect per its own policy.

use std::time::Duration;

use crate::axum::body::Bytes;
use crate::axum::extract::ws::{self, CloseFrame, Message, WebSocket, WebSocketUpgrade};
use crate::axum::http::{HeaderMap, StatusCode};
use crate::axum::response::{IntoResponse, Response};
use crate::realtime::channel::Broadcast;
use crate::realtime::error::{ChannelError, ProtocolHint, RealtimeError, admission_status};
use crate::realtime::origin::{OriginDecision, OriginPolicy};
use crate::realtime::registry::Registry;
use crate::realtime::shutdown::ShutdownConfig;

/// The authorization hook for a realtime connection. The application
/// implements this to decide whether a given connection (identified by its
/// headers — the `Origin` is checked separately by the [`OriginPolicy`])
/// may subscribe to a channel.
///
/// Channel names never implicitly authorize (PROGRAM.md §AP2.1-8): the
/// application must return `Ok(channel)` to admit, or `Err(())` to deny.
/// The hook is `async` and `Send` because it typically consults a session
/// or a policy service.
///
/// The hook returns the [`Broadcast`] the connection subscribes to; the
/// wrapper does not look up channels by name (no implicit-by-name
/// authorization). The application maps its own channel identity to the
/// `Broadcast` it owns.
///
/// The return type is `Option<Broadcast>`: `Some(broadcast)` admits the
/// connection to that channel, `None` denies it (the wrapper maps `None`
/// to [`crate::realtime::error::RealtimeError::Unauthorized`]). Denial
/// carries no associated data, so `Option` is more honest than
/// `Result<_, ()>` (AGENTS.md §18: no future-proof variants, no `()`
/// error that conveys nothing).
pub trait Authorizer: Clone + Send + Sync + 'static {
    /// Decide whether the connection may subscribe, returning the
    /// `Broadcast` to subscribe to. `headers` is the request headers
    /// (e.g. `Cookie` for a session); `channel_id` is an application-
    /// defined channel identifier (e.g. a route parameter the handler
    /// extracted and passed in). The wrapper does not interpret it.
    /// Returns `Some(broadcast)` to admit, `None` to deny.
    fn authorize(
        &self,
        headers: &HeaderMap,
        channel_id: &str,
    ) -> impl std::future::Future<Output = Option<Broadcast>> + Send;
}

/// A trivial `Authorizer` that admits every connection to a single shared
/// broadcast. Intended for tests and for applications that wire their own
/// authorization upstream (e.g. via an Axum layer that rejects before the
/// handler runs). Using it in production without an upstream authz layer
/// would leave channel admission open — the application is responsible
/// for not doing that.
#[derive(Clone)]
pub struct AllowAll {
    broadcast: Broadcast,
}

impl AllowAll {
    #[must_use]
    pub fn new(broadcast: Broadcast) -> Self {
        Self { broadcast }
    }
}

impl Authorizer for AllowAll {
    async fn authorize(&self, _headers: &HeaderMap, _channel_id: &str) -> Option<Broadcast> {
        Some(self.broadcast.clone())
    }
}

/// Frame/message size and timing limits for a WebSocket connection. These
/// are attacker-facing bounds (AGENTS.md §29): an oversized message is
/// rejected by axum (the upgrade is configured with `max_message_size`)
/// and a heartbeat that does not return within the timeout closes the
/// connection as a dead peer.
#[derive(Clone, Copy, Debug)]
pub struct WsLimits {
    /// Maximum message size accepted from the client. Oversize messages
    /// cause axum to close the connection (the upgrade is configured with
    /// [`WebSocketUpgrade::max_message_size`]).
    pub max_message_size: usize,
    /// Maximum frame size accepted from the client.
    pub max_frame_size: usize,
    /// Heartbeat interval (server sends a `Ping`); the client's `Pong` is
    /// handled automatically by axum. Set to zero to disable the heartbeat.
    pub heartbeat_interval: Duration,
    /// If no message (including the auto-`Pong`) is received within this
    /// duration, the connection is closed as dead.
    pub heartbeat_timeout: Duration,
}

impl WsLimits {
    /// Conservative defaults for an attacker-facing endpoint. Message cap
    /// 64 KiB, frame cap 64 KiB, 20s heartbeat, 40s timeout.
    #[must_use]
    pub fn conservative() -> Self {
        Self {
            max_message_size: 64 * 1024,
            max_frame_size: 64 * 1024,
            heartbeat_interval: Duration::from_secs(20),
            heartbeat_timeout: Duration::from_secs(40),
        }
    }
}

/// The configuration for a WebSocket endpoint. Constructed once and cloned
/// into every handler (cheaply — it is a small `Arc`-backed bundle).
#[derive(Clone)]
pub struct WebSocketEndpoint<A> {
    authorizer: A,
    origin: OriginPolicy,
    registry: Registry,
    limits: WsLimits,
    shutdown: ShutdownConfig,
}

impl<A: Authorizer> WebSocketEndpoint<A> {
    /// Construct a WebSocket endpoint with the given authorizer, origin
    /// policy, connection registry, and limits. The shutdown config is
    /// shared with the SSE endpoint (the application passes the same
    /// [`Registry`] and [`ShutdownConfig`]).
    #[must_use]
    pub fn new(
        authorizer: A,
        origin: OriginPolicy,
        registry: Registry,
        limits: WsLimits,
        shutdown: ShutdownConfig,
    ) -> Self {
        Self {
            authorizer,
            origin,
            registry,
            limits,
            shutdown,
        }
    }

    /// The Axum handler. Accept a `WebSocketUpgrade`, the request headers,
    /// and an application-defined `channel_id` (the handler extracts it,
    /// e.g. from a route parameter). Performs admission (origin, authz,
    /// connection limit — all *before* the upgrade, so a rejection returns a
    /// clean HTTP status), then upgrades and runs the typed connection
    /// loop in a spawned task.
    ///
    /// The raw `WebSocketUpgrade` is the escape hatch: an application that
    /// needs full control calls `ws.on_upgrade(...)` directly and skips
    /// this wrapper. The wrapper is opt-in (AGENTS.md §16).
    pub async fn handle(
        self,
        ws: WebSocketUpgrade,
        headers: HeaderMap,
        channel_id: String,
    ) -> Response {
        // ── Admission: origin ──────────────────────────────────────────
        if self.origin.authorize(headers.get("origin")) == OriginDecision::Denied {
            tracing::debug!(
                target: "arcature::realtime::ws::admit",
                realtime_transport = "ws",
                error_category = "origin",
                "realtime upgrade rejected: origin policy"
            );
            return admission_status(&RealtimeError::Origin).into_response();
        }

        // ── Admission: authorization (explicit, not by channel name) ──
        let broadcast = match self.authorizer.authorize(&headers, &channel_id).await {
            Some(bc) => bc,
            None => {
                tracing::debug!(
                    target: "arcature::realtime::ws::admit",
                    realtime_transport = "ws",
                    error_category = "authz",
                    "realtime upgrade rejected: authorization"
                );
                return admission_status(&RealtimeError::Unauthorized).into_response();
            }
        };

        // ── Admission: connection limit (before the upgrade, so a rejection
        //    returns a clean 503, not a post-upgrade 1013). The guard is
        //    moved into the upgrade closure so its lifetime is the
        //    connection's lifetime; if the upgrade never completes, axum
        //    drops the closure and the guard, decrementing the count. ──
        let guard = match self.registry.acquire(self.shutdown.max_connections()) {
            Ok(g) => g,
            Err(_) => {
                tracing::debug!(
                    target: "arcature::realtime::ws::admit",
                    realtime_transport = "ws",
                    error_category = "limit",
                    "realtime upgrade rejected: connection limit"
                );
                return admission_status(&RealtimeError::ConnectionLimit).into_response();
            }
        };

        // ── Upgrade with the configured size limits ─────────────────────
        let limits = self.limits;
        let shutdown = self.shutdown;
        ws.max_message_size(limits.max_message_size)
            .max_frame_size(limits.max_frame_size)
            .on_upgrade(move |socket| run_connection(socket, broadcast, guard, limits, shutdown))
    }
}

/// Run one typed WebSocket connection to completion. The `on_upgrade`
/// callback runs in a task axum spawns, so this function is the
/// connection's lifetime owner. The `guard` was acquired in `handle`
/// (pre-upgrade); it is dropped when this function returns, decrementing
/// the live count.
///
/// The loop:
/// 1. Subscribe to the broadcast.
/// 2. `select!` between: client messages, broadcast payloads (forwarded to
///    the client), the heartbeat ping timer, and the drain signal.
/// 3. On any exit, drop the guard (decrements the live count) and close.
async fn run_connection(
    mut socket: WebSocket,
    broadcast: Broadcast,
    guard: crate::realtime::registry::ConnectionGuard,
    limits: WsLimits,
    shutdown: ShutdownConfig,
) {
    let mut sub = broadcast.subscribe();

    let span = tracing::debug_span!(
        target: "arcature::realtime::ws",
        "ws.connection",
        realtime_transport = "ws",
        // No channel name or client identity: low-cardinality, no payload.
    );
    let _enter = span.enter();

    // Heartbeat state.
    let mut last_seen = tokio::time::Instant::now();
    let mut heartbeat: Option<tokio::time::Interval> = if limits.heartbeat_interval.is_zero() {
        None
    } else {
        let mut i = tokio::time::interval(limits.heartbeat_interval);
        // The first tick fires immediately; skip it so the first ping is
        // after one interval, not at connect time.
        i.tick().await;
        Some(i)
    };

    loop {
        if shutdown.is_draining() {
            break;
        }
        tokio::select! {
            // Drain signal: exit promptly even when blocked on a client
            // recv with no incoming message and a long heartbeat.
            _ = shutdown.drain_notified() => {
                // begin_drain fired; the top-of-loop check will exit.
                continue;
            }
            // Client → server: receive and validate.
            msg = socket.recv() => {
                match msg {
                    Some(Ok(Message::Text(_))) | Some(Ok(Message::Binary(_))) => {
                        // Client text/binary accepted to keep the connection
                        // alive; the wrapper does not interpret it. An
                        // application wanting bidirectional control messages
                        // drops to raw `axum::extract::ws`.
                        last_seen = tokio::time::Instant::now();
                    }
                    Some(Ok(Message::Pong(_))) | Some(Ok(Message::Ping(_))) => {
                        // axum auto-responds to Ping with Pong; refresh the
                        // dead-peer timer on any keepalive frame.
                        last_seen = tokio::time::Instant::now();
                    }
                    Some(Ok(Message::Close(_))) => break,
                    Some(Err(_)) => {
                        // Stream error: close with a protocol error.
                        close_protocol(&mut socket, ProtocolHint::Stream).await;
                        break;
                    }
                    None => break,
                }
            }
            // Broadcast → client: forward the payload (skipped while draining).
            recv = sub.recv(), if !shutdown.is_draining() => {
                match recv {
                    Ok(payload) => {
                        forward_payload(&mut socket, payload).await;
                    }
                    Err(ChannelError::Lagged) => {
                        // A lagged subscriber: notify the client and close.
                        // The application decides resync policy; the wrapper
                        // does not silently drop messages.
                        close_protocol(&mut socket, ProtocolHint::Malformed).await;
                        break;
                    }
                    Err(ChannelError::Closed) | Err(ChannelError::Full) => break,
                }
            }
            // Heartbeat ping timer. The precondition `heartbeat.is_some()`
            // is critical: when the heartbeat is disabled (interval zero),
            // `heartbeat` is `None` and `ping_tick(None)` completes
            // immediately on every poll. Without this guard, the `select!`
            // branch would be permanently ready, turning the loop into a
            // CPU-burning busy spin. The guard parks the branch until a
            // heartbeat interval is actually configured.
            _ = ping_tick(heartbeat.as_mut()), if heartbeat.is_some() => {
                let now = tokio::time::Instant::now();
                if now.duration_since(last_seen) > limits.heartbeat_timeout {
                    let _ = socket
                        .send(Message::Close(Some(CloseFrame {
                            code: ws::close_code::AGAIN,
                            reason: "heartbeat timeout".into(),
                        })))
                        .await;
                    break;
                }
                let _ = socket.send(Message::Ping(Bytes::new())).await;
            }
        }
    }

    // Graceful close: send a close frame; axum flushes and tears down the
    // connection. The guard drops here, decrementing the live count.
    let _ = socket
        .send(Message::Close(Some(CloseFrame {
            code: ws::close_code::AWAY,
            reason: "server draining".into(),
        })))
        .await;
    drop(guard);
}

/// Forward a broadcast payload to the client. Sends as a text frame if the
/// bytes are valid UTF-8 (the common JSON case), else as a binary frame —
/// never panics on non-UTF-8 application bytes.
async fn forward_payload(
    socket: &mut WebSocket,
    payload: crate::realtime::channel::ChannelPayload,
) {
    match std::str::from_utf8(payload.as_bytes()) {
        Ok(s) => {
            let _ = socket.send(Message::text(s.to_string())).await;
        }
        Err(_) => {
            let _ = socket
                .send(Message::binary(payload.as_bytes().to_vec()))
                .await;
        }
    }
}

/// Drive one heartbeat tick. `None` (heartbeat disabled) completes
/// immediately so the `select!` arm never blocks when the heartbeat is off.
async fn ping_tick(tick: Option<&mut tokio::time::Interval>) {
    if let Some(t) = tick {
        t.tick().await;
    }
}

async fn close_protocol(socket: &mut WebSocket, hint: ProtocolHint) {
    let _ = socket
        .send(Message::Close(Some(CloseFrame {
            code: ws::close_code::ERROR,
            reason: hint.to_string().into(),
        })))
        .await;
}

// A `StatusCode` is `IntoResponse`; re-stating the bound here makes the
// admission-path return type check explicit (and keeps a future change
// from silently dropping the bound).
const _: () = {
    fn _assert(status: StatusCode) -> Response {
        status.into_response()
    }
};