ws-kit

Generic authenticated typed WebSocket toolkit for Rust — BroadcastHub, RoomManager, heartbeat, and token extraction for Axum.

Features

• axum (default) — Axum ws + tower, TokenExtractor::extract_token(&Parts, _)
• compression — application-layer DEFLATE for message payloads (see Compression)
• redis — multi-node room fan-out over Redis pub/sub (see Multi-node rooms)
• metrics — emit the StatsRecorder counters through the metrics facade
• tracing — deprecated no-op (kept for 0.3.x compatibility)

Quick start

use ws_kit::{hub::BroadcastHub, codec::Codec};
use serde::{Serialize, Deserialize};

#[derive(Clone, Serialize, Deserialize, Debug, PartialEq)]
struct Chat { user: String, body: String }

#[tokio::main]
async fn main() {
    let hub = BroadcastHub::<Chat>::new(1024);
    let mut rx = hub.subscribe();
    hub.broadcast(Chat { user: "alice".into(), body: "hi".into() }).unwrap();
    let msg = rx.recv().await.unwrap();
    assert_eq!(msg.user, "alice");
}

Rooms (text and binary since 0.4.0):

use ws_kit::room::RoomManager;
use ws_kit::codec::Frame;
let m = RoomManager::new();
let room = m.get_or_create("lobby");
room.join(1, "alice".to_string());
room.broadcast("hello".to_string()).unwrap();          // text
room.broadcast(vec![0xde, 0xad, 0xbe, 0xef]).unwrap(); // binary — zero re-encode

Token extraction:

use ws_kit::extractor::{TokenExtractor, TokenSourceKind};
let ex = TokenExtractor::default(); // Bearer, ?token=, ?access_token=
let tok = ex.extract_from_parts(Some("Bearer xyz"), None, "token=qry");

Origin validation (CSWSH defense):

use ws_kit::config::WsConfig;
use ws_kit::origin_allowed_in_parts;
use axum::http::StatusCode;

// Default (empty list) ALLOWS ALL origins — set your real origins:
let cfg = WsConfig::builder().allow_origin("https://app.example.com").build();

// In the upgrade handler, BEFORE ws.on_upgrade:
if !origin_allowed_in_parts(&parts, &cfg.allowed_origins) {
    return StatusCode::FORBIDDEN.into_response();
}

Exact match after normalization (lowercase scheme/host, default ports omitted); no wildcard or suffix matching. See REQUIREMENTS.md (REQ-WSKIT-200/201) and THREAT-MODEL.md.

Codec:

use ws_kit::codec::Codec;
use serde::{Serialize, Deserialize};
#[derive(Serialize, Deserialize, PartialEq, Debug)] struct Msg { t: String }
let m = Msg { t: "hi".into() };
let s = m.encode();
let d = Msg::decode(&s).unwrap();

Binary messages (0.4.0)

The wire message model is Frame:

• Frame::Text(String) — UTF-8 payloads; Codec/JsonCodec operate here.
• Frame::Binary(Bytes) — opaque bytes, carried verbatim.

Routing and room broadcast accept both (impl Into<Frame> — String, &str, Bytes, Vec<u8> convert with no re-encoding of owned payloads), and BroadcastHub<Frame> works out of the box (Frame is Serialize).

Security posture: a Frame::Binary payload is never parsed as text or JSON by any ws-kit code path — the JSON decoders take &str, so the type system enforces the separation. Text is guaranteed UTF-8 end to end (Frame::Text can only be built from a String). Authentication (TokenExtractor) reads the HTTP upgrade request and is frame-kind-agnostic.

Compression

Feature: compression (off by default; pulls flate2 with the pure-Rust miniz_oxide backend — the crate stays #![forbid(unsafe_code)]).

use ws_kit::codec::Frame;
use ws_kit::compression::{CompressionConfig, FrameCompressor};

let c = FrameCompressor::new(CompressionConfig::default());
let wire = c.compress_frame(Frame::from("repetitive payload ".repeat(100)))?;
let restored = c.decompress_frame(wire)?; // == original frame

Which layer — honest assessment: the right layer for WebSocket compression is the permessage-deflate handshake extension (RFC 7692). The stack ws-kit builds on — tokio-tungstenite/tungstenite 0.29 (what axum 0.8 pulls) — does not implement it and exposes no extension hook to negotiate one. Rather than hand-roll an RFC 7692 state machine against frame internals ws-kit doesn't own, 0.4.0 ships an application-layer DEFLATE codec and says so:

• Payloads ride ordinary binary frames with a 1-byte envelope (flags | deflate-or-identity). Both endpoints must run ws-kit and opt in — browsers cannot use this (they only speak permessage-deflate). It is for Rust-to-Rust links (server ↔ ws-kit client/proxy/worker).
• "Negotiation" is out-of-band configuration; for integrators signaling the capability in a custom handshake header, extension_accepted(Some(header)) matches the x-ws-kit-deflate token case-insensitively.
• Decompression-bomb guard: decompressed output is capped during inflation (CompressionConfig::max_size, default 1 MiB) — a compressed bomb never materializes in memory. Unknown envelope flags and corrupt streams fail closed (WsError::Compression).
• Oversharing is avoided on purpose: payloads below min_size (default 64 B) and incompressible payloads are sent as identity — compression never grows the wire size.

Measured on the built-in bench (cargo bench --features compression --bench compression_roundtrip, ratios are deterministic; absolute times vary with machine load): synthetic 2 KiB chat JSON → 174 B (11.9×); repetitive 16 KiB binary → 130 B (126×).

When a handshake-level permessage-deflate lands upstream in tokio-tungstenite/axum, CompressionConfig is the struct that gets wired through.

Multi-node rooms (Redis)

Feature: redis (off by default). RedisRoomRegistry keeps rooms local (ordinary RoomManager semantics) and fans broadcasts out through Redis pub/sub, so clients connected to different nodes share rooms:

       node A                                node B
 ┌─────────────────┐                  ┌─────────────────┐
 │ RoomManager (A) │                  │ RoomManager (B) │
 │   room "lobby"  │                  │   room "lobby"  │
 └───┬────────┬────┘                  └───┬────────┬────┘
     │        │                           │        │
  local    PUBLISH ws-kit:room:lobby   local    PUBLISH
  sinks          │                     sinks        │
     │        ┌───▼───────────────────────┼──────────┤
     │        │            Redis          │          │
     │        └───┬───────────────────────┴──────────┤
     │            │     (pub/sub: best-effort)       │
  every node PSUBSCRIBEs ws-kit:room:* → delivers each message
  to its LOCAL room with that id (if hosted)

let registry = std::sync::Arc::new(
    ws_kit::redis_rooms::RedisRoomRegistry::connect("redis://127.0.0.1:6379").await?,
);
// One subscriber task per node, at startup:
let sub = std::sync::Arc::clone(&registry);
tokio::spawn(async move { let _ = sub.run_subscriber().await; });

registry.broadcast("lobby", ws_kit::codec::Frame::from("hi fleet")).await?;

Topology details:

• broadcast(room, frame) delivers to local sinks and PUBLISHes once; each node's subscriber loop applies received messages to the local room of that id — only if the node hosts it (remote traffic never materializes rooms).
• Each registry carries a random node id; self-echoes are dropped, so a publisher's local receivers see each message exactly once.
• Channel prefix is configurable (connect_with_prefix) — unrelated deployments can share one Redis.

Delivery semantics — read before trusting: Redis pub/sub is at-most-once, best-effort. No persistence, no replay, no acks; a node that is disconnected or lagging misses messages, permanently. ws-kit does not provide at-least-once delivery and makes no persistence claim. Use it for ephemeral fan-out (chat, presence, live notifications). For at-least-once, the RoomRegistry trait keeps the swap to Redis Streams (or any log) local to one module.

Trust boundary: anyone who can PUBLISH to the channel prefix can inject messages into any room on any node; restrict with Redis ACLs / network isolation and use rediss:// for untrusted links.

Testing: the publish path is covered by a trait-mocked ConnectionLike (see src/redis_rooms.rs); the full two-node pub/sub loop runs as fixture-gated tests:

docker run -d -p 6379:6379 redis:7
cargo test --features redis --test redis_rooms -- --ignored

Stats & backpressure (0.4.0)

StatsRecorder — cheap-to-clone, lock-free (atomics + sharded map):

use ws_kit::{config::WsConfig, room::RoomManager, stats::StatsRecorder};
let stats = StatsRecorder::new();
let manager = RoomManager::with_stats(&WsConfig::default(), stats.clone());
// room broadcasts now record per-room + global messages_out/bytes_out/drops
let snap = stats.snapshot();
// snap.connections, snap.rooms, snap.messages_in/out, snap.bytes_in/out,
// snap.drops, snap.per_room["lobby"]

Counters: connections (your accept loop: inc_connections() / dec_connections()), rooms (wired into RoomManager create/remove/cleanup), messages/bytes in (record_message_in — from your read loop), messages/bytes out (wired into room broadcasts), drops (broadcasts with no live receiver — backpressure that lost data). With the metrics feature the same numbers are emitted as ws_kit_* counters/gauges on the metrics facade for Prometheus/OTLP export — same pattern as the breaker crate.

Concurrency testing

The connection counter's atomic logic (bounded increment CAS loop, saturating decrement) is model-checked with loom under --cfg loom — see src/loom_tests.rs:

RUSTFLAGS="--cfg loom" cargo test --release --lib --no-default-features -- loom

(--no-default-features drops the axum stack — tokio-tungstenite does not compile under --cfg loom; the models only need the counter.)

Loom exhaustively explores bounded interleavings and proves: balanced increment/decrement pairs always return the count to zero, the count never goes negative, and the connection limit is never exceeded under races. What loom does not cover: tokio::sync::broadcast and DashMap internals (not loom-compatible) — those are trusted via tokio's and dashmap's own concurrency testing.

Security

Threat model: THREAT-MODEL.md.

Performance

Measured hot-path SLOs and allocation profile: PERF-SLO.md, with every numeric claim mapped to its proof artifact in CLAIMS.md. Benchmarks run in CI (non-gating regression visibility against the saved ci baseline); the allocation profile is proven by a counting-allocator test (tests/zero_alloc_frame_rx.rs) and the hot paths are pinned by an iai-callgrind instruction gate (benches/iai_hot_path.rs).

The broadcast ring is pre-allocated at channel construction — steady-state broadcast() performs no per-message allocation beyond the caller's message value.