bevy_symbios_multiuser 0.7.0

//! XRPC Relay — a WebRTC signaling broker built on Axum with optional
//! ATProto JWT authentication and DID-based signature verification.
//!
//! The relay accepts WebSocket connections, optionally verifies ATProto JWT
//! bearer tokens, and routes WebRTC SDP offers/answers and ICE candidates
//! between peers using an in-memory connection map.
//!
//! # Authentication
//!
//! When [`RelayConfig::auth_required`] is `true`, every connecting client must
//! present a valid ATProto access JWT. The relay accepts the token from two
//! sources, checked in order:
//!
//! 1. `Authorization: Bearer <token>` header (native clients).
//! 2. `Sec-WebSocket-Protocol: access_token, <token>` header (WASM clients
//!    using the subprotocol trick — the relay echoes the selected subprotocol
//!    back per RFC 6455).
//!
//! Bearer tokens are intentionally **not** accepted via query string: query
//! parameters are routinely captured in plaintext by reverse proxy access logs,
//! load balancers, and intermediate firewalls, which would leak ATProto session
//! credentials into operator-side logs the user never consented to share.
//!
//! The relay resolves the issuer's DID document (via `plc.directory` for
//! `did:plc`, or HTTPS for `did:web` — domain-only DIDs use
//! `/.well-known/did.json`, path-based DIDs like `did:web:example.com:u:alice`
//! use `/{path}/did.json`), extracts the `#atproto` signing key (P-256/ES256
//! or secp256k1/ES256K), and cryptographically verifies the JWT signature.
//! Resolved keys are cached in memory with a 5-minute TTL.
//! The authenticated DID becomes the peer's session identity.
//!
//! When `auth_required` is `false`, authentication is opportunistic — clients
//! presenting a valid, signature-verified token are identified by their DID,
//! while clients that present no token fall back to random UUIDs. An explicitly
//! invalid token is rejected with HTTP 401 regardless of `auth_required`.
//! Tokens are only trusted when a DID resolver is configured.
//!
//! # Room Isolation
//!
//! The URL path used during WebSocket upgrade determines the peer's **room**.
//! For example, `wss://relay/game_A` and `wss://relay/game_B` are separate
//! rooms — peers only see `PeerJoined`/`PeerLeft` events and can only exchange
//! signals with other peers in the same room. Cross-room signals are dropped.
//! Connecting to `/` (or with no path) places the peer in a `"default"` room.
//! Room paths are percent-decoded so that `/my%20room` and `/my room` resolve
//! to the same room.
//!
//! Peer state is keyed by `(room, session_id)`, so the same authenticated
//! identity may legitimately hold concurrent connections in different rooms
//! (e.g. a lobby and a game room). Within a single room, a reconnect from the
//! same identity still replaces the prior entry and emits a `PeerLeft` to the
//! room so the WebRTC mesh state machine can recover cleanly.
//!
//! # Hardening
//!
//! - **HTTP request timeout** — A `tower-http` `TimeoutLayer` drops any HTTP
//!   connection that has not completed the request (including header parsing and
//!   WebSocket upgrade) within 10 seconds, mitigating Slowloris-style attacks
//!   that trickle headers slowly to hold TCP connections without ever reaching
//!   the WebSocket handler.
//! - **Connection limits** — [`RelayConfig::max_peers`] caps the number of
//!   concurrent connections (default `512`). The limit is enforced via an atomic
//!   counter that reserves a slot *before* async identity extraction, preventing
//!   TOCTOU bypasses from concurrent handshakes. An RAII `ConnectionGuard`
//!   ensures the counter is decremented even if the WebSocket upgrade callback
//!   is never executed (e.g. TCP drops during the HTTP handshake). New
//!   connections are rejected with HTTP 503 once the limit is reached.
//! - **Message size cap** — Incoming WebSocket messages are limited to 64 KiB.
//!   SDP offers/answers and ICE candidates are typically a few KiB at most.
//! - **Control signal filtering** — Clients cannot forge `PeerJoined`/`PeerLeft`
//!   control signals; only the relay may originate these.
//! - **SSRF protection** — `did:web` domain resolution validates against
//!   private/loopback IPs and pins the resolved address to prevent DNS rebinding.
//! - **DID document size limit** — Responses are streamed with an incremental
//!   256 KiB cap, aborting before buffering oversized payloads.
//! - **Idle timeout** — WebSocket connections that receive no messages for 120
//!   seconds are disconnected, preventing Slowloris-style attacks that hold
//!   connection slots indefinitely.
//! - **Handshake timeout** — The authentication/identity extraction phase is
//!   capped at 15 seconds, preventing connection slot exhaustion from DIDs that
//!   tarpit the HTTP fetch.
//! - **WebSocket write timeout** — Every outbound write (signaling envelope or
//!   server-side Ping frame) is wrapped in a 5-second timeout. Without this, an
//!   attacker that opens a connection but never drains their TCP receive buffer
//!   would cause `ws_tx.send` to block indefinitely; client-side Pings would
//!   keep the idle-timeout from firing, permanently holding a connection slot.
//! - **Self-targeting rejection** — SDP offers/answers addressed to the sender's
//!   own session ID are dropped, preventing pointless self-negotiation loops.
//! - **Invalid message disconnect** — Peers that send 10 cumulative invalid
//!   messages (malformed JSON, binary frames, forged control signals) are
//!   disconnected, preventing log exhaustion attacks.
//! - **Negative DID cache** — Authoritative DID resolution failures (404,
//!   malformed document, SSRF-blocked IP, unsupported key type) are cached
//!   for 60 seconds, preventing attackers from using the relay as a DDoS
//!   reflector by spamming handshakes with the same DID pointing at a victim
//!   server. **Transient** failures (DNS timeout, concurrency-limit
//!   rejection, 5xx origin response) are deliberately **not** cached so that
//!   a local bottleneck cannot be amplified into a 60-second outage for a
//!   DID that remains perfectly valid — the legitimate user whose handshake
//!   races a bottleneck is free to retry as soon as the transient clears.
//! - **DID request coalescing** — The key cache uses [`moka::future::Cache`]
//!   with `try_get_with`, deduplicating concurrent lookups for the same DID.
//!   This prevents the relay from amplifying connection bursts into outbound
//!   HTTP floods against DID hosting servers.
//! - **Domain client cap** — Per-domain `reqwest::Client` instances (DNS-pinned
//!   for SSRF protection) are capped at 100. Each client holds a connection
//!   pool and background workers, so the low cap prevents resource exhaustion
//!   from an attacker feeding many unique `did:web` domains.
//! - **Server-side WebSocket pings** — The relay sends Ping frames every 30
//!   seconds. Browsers cannot initiate WebSocket pings (the API only supports
//!   *responding* to pings), so without server-side pings, idle WASM clients
//!   would be reaped by the idle timeout.
//! - **Backpressure** — When the per-peer relay channel (256 slots) is full,
//!   signals are dropped and each sender's per-target strike counter is
//!   incremented; per-sender log emission is silenced after 50 strikes to
//!   bound log noise. An **aggregate** counter per target (shared across all
//!   senders) accumulates a strike on every `TrySendError::Full` and drips
//!   on every successful send. Once the aggregate crosses 256 strikes
//!   (matching `RELAY_CHANNEL_CAPACITY`), the target's write task is
//!   signalled to shut down, closing the WebSocket with code 1013 ("Try
//!   Again Later"). This is deliberate: WebRTC does not retransmit SDP
//!   offers or ICE candidates over the signaling channel, so a silently
//!   dropped message permanently stalls the mesh — disconnecting the target
//!   surfaces a hard error that the client's reconnect logic can recover
//!   from. The tight per-sender burst limit (16) means a single flooding
//!   peer contributes at most 16 of the 256 strikes required, so no single
//!   attacker can kick an arbitrary target on its own.
//! - **Handshake slot budget** — At most `max_peers / 4` connections may be
//!   in the authentication/DID-resolution phase simultaneously. This prevents
//!   attackers from exhausting all connection slots by tarpitting the DID
//!   fetch with slow-responding servers. When `max_peers == 0` (unlimited),
//!   the budget falls back to a fixed cap so the tarpit protection is never
//!   silently disabled by an unlimited-peers configuration.
//! - **Per-target burst limiting** — Each sender may route at most 16 messages
//!   to the same target within a single per-target window. This prevents one
//!   sender from filling a target's relay channel (256 slots) with garbage,
//!   which would cause legitimate signalling messages from other peers to be
//!   silently dropped (WebRTC negotiation sabotage). Set to 16: comfortably
//!   above legitimate mesh setup (~1 SDP + ~10 ICE per target) while keeping
//!   a single sender's contribution to ~6% of channel capacity, raising the
//!   coordination cost of stuffing a victim's channel to 16 attackers.
//! - **Per-sender rate limiting** — Each peer is rate-limited via a token bucket
//!   with a burst capacity of `max(1024, max_peers × 16)`, capped at 16,384,
//!   and a steady-state refill of 20 tokens per second. The burst scales with
//!   the operator's configured room size so a full-mesh WebRTC init (N SDP
//!   offers + multiple ICE candidates each) never trips the limiter on
//!   legitimate traffic, while the low refill rate caps sustained throughput.
//!   A peer that still exhausts the budget is **disconnected**: the signaling
//!   channel does not retransmit, so silently dropping signals would leave
//!   the mesh stalled — surfacing a hard error lets the client's reconnect
//!   logic rebuild the mesh cleanly.
//! - **Per-domain DID fetch concurrency limit** — Each `did:web` domain is
//!   limited to 10 concurrent in-flight fetches. Slots are released as soon as
//!   each fetch completes (via RAII guard), so attacker requests that fail
//!   quickly cannot permanently exhaust the budget for legitimate users.
//! - **Global `did:web` fetch concurrency limit** — Total concurrent `did:web`
//!   fetches across all domains are capped at 50. DNS resolution runs on the
//!   pure-Rust `hickory-resolver` with its own UDP/TCP sockets (no OS-blocking
//!   `getaddrinfo`), so dropping the future closes the socket and releases
//!   the concurrency guard synchronously — an attacker with a tarpit
//!   nameserver cannot hold slots past the client's 15 s handshake timeout.
//!   Slots are also freed on normal completion, making the limit resistant
//!   to unauthenticated DoS.
//! - **Peer ID length validation** — The `peer_id` field in incoming
//!   [`SignalEnvelope`] messages is capped at 512 bytes after deserialization.
//!   DIDs and UUIDs are well under this limit; oversized values are rejected
//!   as invalid messages to prevent per-target map bloat and log-output
//!   amplification.
//! - **Unique target cap** — Each sender may address at most
//!   `clamp(max_peers, 256, 4096)` unique targets within a single per-target
//!   rate window (4,096 when `max_peers = 0`/unlimited). Legitimate peers
//!   target at most the number of peers in their room; an attacker forging
//!   random target IDs to bloat the per-target counter map is disconnected
//!   when the cap is exceeded.
//! - **JWT audience validation** — When [`RelayConfig::service_did`] is set,
//!   the relay validates the JWT `aud` claim against the configured value.
//!   This prevents cross-service token replay attacks where a JWT issued for
//!   one relay is presented to a different relay.
//!
//! # Usage
//!
//! ```rust,no_run
//! use bevy_symbios_multiuser::relay::{RelayConfig, run_relay};
//!
//! #[tokio::main]
//! async fn main() {
//!     let config = RelayConfig {
//!         bind_addr: "0.0.0.0:3536".to_string(),
//!         auth_required: false,
//!         max_peers: 512,
//!         service_did: None,
//!     };
//!     run_relay(config).await.expect("relay crashed");
//! }
//! ```

pub(crate) mod auth;
pub(crate) mod did_resolver;
mod handler;

use dashmap::DashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize};
use tokio::sync::{Notify, mpsc};

// Re-export protocol types so existing `use relay::SignalEnvelope` still works.
pub use crate::protocol::{SignalEnvelope, SignalPayload};

/// Configuration for the relay server.
#[derive(Debug, Clone)]
pub struct RelayConfig {
    /// The address to bind the server to (e.g. `"0.0.0.0:3536"`).
    pub bind_addr: String,
    /// If `true`, reject WebSocket connections that do not present a valid
    /// ATProto JWT. The token is accepted via the `Authorization: Bearer`
    /// header (native clients) or the `Sec-WebSocket-Protocol: access_token,
    /// <jwt>` subprotocol trick (WASM clients); it is **not** accepted via
    /// query string (query parameters leak into reverse-proxy access logs).
    /// When `false`, authentication is opportunistic: valid tokens are used
    /// for identity while clients that present no token receive a random UUID.
    pub auth_required: bool,
    /// Maximum number of concurrent peer connections. New connections are
    /// rejected with HTTP 503 once this limit is reached. `0` means unlimited.
    /// A reasonable starting value is `512`.
    pub max_peers: usize,
    /// The relay's own service DID (e.g. `did:web:relay.example.com`). When set,
    /// JWT `aud` claims are validated against this value to prevent cross-service
    /// token replay attacks. When `None`, audience validation is skipped.
    pub service_did: Option<String>,
}

/// A connected peer's sender handle paired with a unique connection ID.
///
/// The connection ID distinguishes multiple WebSocket connections from the
/// same user (e.g. reconnects), preventing stale cleanup from clobbering a
/// newer connection.
#[derive(Clone)]
pub struct PeerEntry {
    /// Channel sender for delivering signaling envelopes to this peer's
    /// WebSocket write task.
    pub tx: mpsc::Sender<SignalEnvelope>,
    /// Unique ID for this specific WebSocket connection, used to distinguish
    /// reconnects and prevent stale cleanup from clobbering a newer session.
    pub conn_id: uuid::Uuid,
    /// Aggregate backpressure strike counter across all senders targeting
    /// this peer. Incremented by `TrySendError::Full` and decremented by
    /// successful sends. Once it crosses the kick threshold, [`Self::shutdown`]
    /// is fired so the peer's write task exits and the client reconnect logic
    /// can rebuild the WebRTC mesh — preferable to silently blackholing SDP
    /// offers/ICE candidates which never retransmit on the signaling channel.
    pub backpressure_strikes: Arc<AtomicU64>,
    /// Force-shutdown signal for this peer's write task. Cleanly breaks the
    /// write loop (peer receives a WebSocket close) so reconnect logic fires
    /// instead of the mesh stalling forever in the "connecting" state.
    pub shutdown: Arc<Notify>,
}

/// Shared server state holding the map of connected peers.
#[derive(Clone)]
pub struct RelayState {
    /// Two-level map: `room -> session_id -> PeerEntry`. Keying rooms at the
    /// outer level (instead of using a flat `(room, session_id)` tuple map)
    /// turns per-room broadcasts on `PeerJoined`/`PeerLeft` from an
    /// O(N_total_peers) scan into an O(K_peers_in_room) iteration, which is
    /// the primary scaling knob for a relay hosting thousands of rooms.
    ///
    /// Nesting still allows the same authenticated identity to hold concurrent
    /// connections in different rooms (each lives under a different outer
    /// key); within a single room, a reconnect from the same identity still
    /// replaces the prior entry.
    ///
    /// The inner map is wrapped in `Arc` so that a brief outer-shard read lock
    /// can be dropped immediately after cloning the handle — iteration,
    /// lookups, and per-target sends then hold only the inner shard locks,
    /// which is what keeps unrelated rooms from serialising on each other.
    pub peers: Arc<DashMap<String, Arc<DashMap<String, PeerEntry>>>>,
    /// Whether authentication is mandatory for new connections.
    pub auth_required: bool,
    /// Maximum concurrent peers (`0` = unlimited).
    pub max_peers: usize,
    /// DID document resolver for JWT signature verification.
    /// `None` disables cryptographic signature checks.
    pub did_resolver: Option<did_resolver::DidResolver>,
    /// The relay's own service DID for JWT audience validation.
    /// When set, tokens whose `aud` claim does not match are rejected.
    pub service_did: Option<String>,
    /// Atomic counter tracking active + in-handshake connections.
    /// Prevents TOCTOU bypasses where concurrent handshakes all pass the
    /// `max_peers` check before any of them insert into `peers`.
    pub active_connections: Arc<AtomicUsize>,
    /// Atomic counter tracking connections currently in the auth/DID-resolution
    /// phase. Capped at `max_peers / 4` (minimum 1) to prevent DID tarpit
    /// attacks from exhausting all connection slots while legitimate peers wait.
    pub active_handshakes: Arc<AtomicUsize>,
}

impl RelayState {
    fn new(auth_required: bool, max_peers: usize, service_did: Option<String>) -> Self {
        // Always create the DID resolver so that opportunistic authentication
        // works when `auth_required` is false: clients presenting a valid JWT
        // get identified by their DID, while unauthenticated clients fall back
        // to random UUIDs.
        Self {
            peers: Arc::new(DashMap::new()),
            auth_required,
            max_peers,
            did_resolver: Some(did_resolver::DidResolver::new()),
            service_did,
            active_connections: Arc::new(AtomicUsize::new(0)),
            active_handshakes: Arc::new(AtomicUsize::new(0)),
        }
    }
}

/// Start the relay signaling server.
///
/// Binds to the configured address and serves WebSocket connections.
/// This function runs until the server is shut down.
pub async fn run_relay(config: RelayConfig) -> Result<(), Box<dyn std::error::Error>> {
    let state = RelayState::new(config.auth_required, config.max_peers, config.service_did);

    // Accept WebSocket upgrades on any path so clients can use room-based
    // URLs (e.g. `/my_room`) as well as the canonical `/ws` endpoint.
    //
    // The tower-http TimeoutLayer wraps the entire HTTP service so that
    // connections which have not completed the HTTP request (including header
    // parsing and WebSocket upgrade) within 10 seconds are dropped. This
    // mitigates Slowloris-style attacks where an attacker trickles HTTP
    // headers slowly, holding TCP connections without ever reaching the
    // WebSocket handler or its idle timeout logic.
    let app = axum::Router::new()
        .fallback(axum::routing::get(handler::ws_handler))
        .with_state(state)
        .layer(tower_http::timeout::TimeoutLayer::with_status_code(
            axum::http::StatusCode::SERVICE_UNAVAILABLE,
            std::time::Duration::from_secs(10),
        ));

    let listener = tokio::net::TcpListener::bind(&config.bind_addr).await?;
    tracing::info!(
        addr = %config.bind_addr,
        auth_required = config.auth_required,
        "relay server listening"
    );

    axum::serve(listener, app).await?;

    Ok(())
}