diff --git a/README.md b/README.md index ec96304..da8471a 100644 --- a/README.md +++ b/README.md @@ -256,6 +256,49 @@ process registry: `urus.server` (the listener-pool supervisor) and `urus.listener.{i}` (each accept loop). `smarm::whereis(name)` resolves them from any actor inside the runtime. +### WebSocket (v0.4) + +A route handler accepts the upgrade by handing `Conn::upgrade` a +`WsHandler`; after the `101` the connection actor leaves HTTP and drives +the handler from a duplex select loop (smarm RFC 008 fd arms: first-of +outbound-channel / fd-readable, one actor, no fd dup): + +```rust +struct Echo; + +impl urus::WsHandler for Echo { + fn on_message(&mut self, msg: urus::Message, sender: &urus::WsSender) { + let _ = sender.send(msg); // or .text(..) / .binary(..) / .ping() / .close(code, reason) + } + fn on_close(&mut self, code: Option, reason: &str) { + // what the PEER said: Some(code) iff its close frame arrived; + // None for EOF / write failure / handshake timeout. + let _ = (code, reason); + } +} + +Router::new().get("/echo", |c: Conn, _n: Next| c.upgrade(Echo)) +``` + +Callbacks run **inside the connection actor** — a slow `on_message` +stops reads, which is backpressure, not a bug. A handler that wants +concurrency spawns its own actor and hands it a `WsSender` clone (the +SSE-producer pattern); every `WsSender` method returns `Err(WsClosed)` +once the connection ends. While a large outbound frame is mid-write the +actor isn't reading — accepted v1 trade-off of the one-actor topology. + +Control frames are invisible to the handler in v1: pings are +auto-ponged, a peer close is auto-echoed and surfaces as `on_close`. +Protocol violations close with the mapped RFC 6455 code (1002/1009/1007; +handler panic → 1011). Caps: `Config.max_frame_payload` (1 MiB default, +enforced from the frame header before payload buffers) and +`Config.max_message_bytes` (4 MiB default, spans fragments). + +Like SSE, an open WebSocket has no request clock: idle is fine, a dead +client is caught when a write stalls past `write_timeout`, and graceful +shutdown force-stops the connection at the drain deadline. See +`examples/ws_echo.rs`. + ## Examples ### CRUD with Actor Ownership @@ -326,9 +369,10 @@ This avoids the complexity of async ecosystems while maintaining full concurrenc ## Roadmap See [`ROADMAP.md`](ROADMAP.md). Summary: v0.2 (supervised listener pool, -graceful shutdown, enforced timeouts, registry names) and v0.3 (streaming -bodies, chunked requests, SSE) are done; next up are WebSocket (v0.4), -PubSub (v0.5), and Phoenix-style channels (v0.6). +graceful shutdown, enforced timeouts, registry names), v0.3 (streaming +bodies, chunked requests, SSE) and v0.4 (WebSocket: handshake, frame +codec, duplex handler API) are done; next up are PubSub (v0.5) and +Phoenix-style channels (v0.6). Refer to `urus-spec.md` and `urus-v1-build-notes.md` in the artifact persistence for the original design and implementation notes. diff --git a/ROADMAP.md b/ROADMAP.md index 177ba4e..c82262e 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -159,7 +159,7 @@ force-stopped at the drain deadline on shutdown; full suite green. --- -## v0.4 — WebSocket upgrade path +## v0.4 — WebSocket upgrade path ✅ (landed 2026-06, chunks 1-3) 1. **Handshake** ✅ (49538cc) — `Conn::upgrade()` (handler arg comes with chunk 3): §4.2.1 validation, 101 + accept key, marker field on @@ -175,22 +175,46 @@ force-stopped at the drain deadline on shutdown; full suite green. close-code wire validation, `Assembler` for fragmentation with whole-message UTF-8 check, `FrameError`→close-code map (1002/1009/1007). §5.7 examples pinned verbatim. -3. **Duplex topology** (the design decision from the assessment): - - **(a) Two actors, dup'd fd**: reader actor on the dup, writer actor - on the original, writer owns a `Receiver`; the user handler - talks to both via channels. Proven trick (listener pool), works today. - - **(b) One actor, smarm first-of(fd-readable, channel)**: fd arms in - select landed with RFC 008 (`393cdd0`) — no longer blocked. Simpler - topology, no dup; use the fallible `try_select` twins (registration - errors, incl. `AlreadyExists`, surface as `Err` and the wait fully - retires). - Was "default (a), revisit if (b) lands" — (b) landed, so this is a - live choice for the v0.4 design pass. Note RFC 008 is phase 1: - single-waiter-per-fd stands, `dup` remains the documented answer for - true simultaneous read+write waits. -4. Handler API sketch (settle in design pass): trait with - `on_message/on_close` + a `WsSender` handle — gen_server-shaped, so a - ws connection *is* an actor you can monitor, link, and register. +3. **Duplex wiring** ✅ — topology decision: **(b) one actor, RFC 008 + fd arms** (urus exists to put smarm through its paces; first fd-arm + consumer is the point). As landed: + - `ws::duplex::run_duplex` replaces the HTTP loop after the 101: + `try_select(&[&outbound_rx, &FdArm::readable(fd)])` per iteration — + outbound at index 0 (select is priority-ordered: owed writes drain + before reading more). Accepted gap documented: mid-write of a large + outbound frame the actor isn't reading. + - Handler API (the deferred questions, as answered): in-actor + callbacks — `WsHandler { on_message, on_close }` runs inside the + select loop; concurrency = spawn your own actor with a `WsSender` + clone (the SSE-producer pattern). `Conn::upgrade(self, handler)` + flat signature (builder deferred until per-upgrade options exist, + e.g. subprotocols). `WsSender`: Clone, `send/text/binary/ping/ + close → Err(WsClosed)`, unbounded channel like SSE (slow client + bounded by write_timeout). Control frames invisible v1: auto-pong + inline (§5.5.2), peer close auto-echoed (§5.5.1, server closes TCP + first per §7.1.1) surfacing only as `on_close(Some(code), reason)`; + wordless endings (EOF, write failure, handshake timeout) = + `on_close(None, "")`. `on_ping`/`on_pong` can land later as default + trait methods, non-breaking. + - Server-initiated close (WsSender::close or FrameError→1002/1009/ + 1007, handler panic→1011): close frame out, then a bounded drain + (one write_timeout budget) for the peer's echo, discarding data + (§1.4). Handler panics re-raise smarm's stop sentinel before being + treated as panics — the pipeline's catch_unwind dance, replicated. + - Caps: `Config.max_frame_payload` (1 MiB) / `max_message_bytes` + (4 MiB) — Config fields per the max_body_bytes precedent; header + cap checked before payload buffers. + - Buf handover: bytes pipelined past the upgrade head carry into the + duplex loop (tested). + - Registry: conn stays **Busy** for the ws lifetime (an open ws is + in-flight work); graceful shutdown force-stops it out of the select + park at the drain deadline (tested — the in-runtime request_stop + wake covers select parks, not just channel recv). + - Validated against a real client (python websocket-client): echo, + ping/pong payload, close handshake. `examples/ws_echo.rs` shows + both handler shapes (/echo in-actor, /clock producer-spawning). +4. **Handler API** ✅ — folded into chunk 3 above (was: settle in design + pass; settled: gen_server-shaped trait + WsSender handle). ## v0.5 — PubSub diff --git a/examples/ws_echo.rs b/examples/ws_echo.rs new file mode 100644 index 0000000..0a96e44 --- /dev/null +++ b/examples/ws_echo.rs @@ -0,0 +1,103 @@ +//! WebSocket echo server (v0.4). +//! +//! cargo run --example ws_echo +//! websocat ws://127.0.0.1:8080/echo +//! +//! Demonstrates the chunk-3 handler model: +//! - `EchoHandler` runs INSIDE the connection actor's select loop — echo +//! is exactly the workload that wants zero extra moving parts. +//! - `/clock` shows the other shape: the handler spawns a producer actor +//! at `on_message("start")` and hands it a `WsSender` clone — the +//! SSE-producer pattern, verbatim. The producer exits when its send +//! fails (`WsClosed`: client gone or server shutting down). +//! +//! Routing happens before the upgrade, so one server can host both +//! endpoints — `Conn::upgrade(handler)` is just another thing a route +//! handler returns. + +use std::time::Duration; + +use urus::{ + serve_with_shutdown, shutdown_handle, Config, Conn, Message, Next, Pipeline, Router, + WsHandler, WsSender, +}; + +// --------------------------------------------------------------------------- +// /echo — everything comes straight back +// --------------------------------------------------------------------------- + +struct EchoHandler; + +impl WsHandler for EchoHandler { + fn on_message(&mut self, msg: Message, sender: &WsSender) { + if matches!(&msg, Message::Text(t) if t == "bye") { + let _ = sender.close(1000, "you said bye"); + return; + } + let _ = sender.send(msg); + } + + fn on_close(&mut self, code: Option, reason: &str) { + println!("ws_echo: /echo closed (code {code:?}, reason {reason:?})"); + } +} + +// --------------------------------------------------------------------------- +// /clock — "start" spawns a producer actor ticking once a second +// --------------------------------------------------------------------------- + +struct ClockHandler { + started: bool, +} + +impl WsHandler for ClockHandler { + fn on_message(&mut self, msg: Message, sender: &WsSender) { + match msg { + Message::Text(t) if t == "start" && !self.started => { + self.started = true; + let out = sender.clone(); + smarm::spawn(move || { + let mut n = 0u64; + loop { + if out.text(format!("tick {n}")).is_err() { + return; // WsClosed: connection over, we follow + } + n += 1; + smarm::sleep(Duration::from_secs(1)); + } + }); + } + _ => { + let _ = sender.text("send \"start\" to begin"); + } + } + } +} + +// --------------------------------------------------------------------------- +// main +// --------------------------------------------------------------------------- + +fn main() { + let pipeline = Pipeline::new().plug( + Router::new() + .get("/echo", |c: Conn, _n: Next| c.upgrade(EchoHandler)) + .get("/clock", |c: Conn, _n: Next| { + c.upgrade(ClockHandler { started: false }) + }), + ); + + let cfg = Config::new("127.0.0.1:8080".parse().unwrap()); + println!("ws_echo: ws://127.0.0.1:8080/echo and /clock — press Enter to shut down"); + + let (handle, signal) = shutdown_handle(); + std::thread::spawn(move || { + let mut line = String::new(); + let _ = std::io::stdin().read_line(&mut line); + println!("ws_echo: shutting down…"); + handle.shutdown(); + }); + + serve_with_shutdown(cfg, pipeline, signal).unwrap(); + println!("ws_echo: bye"); +} diff --git a/src/conn.rs b/src/conn.rs index 4a91e65..46460c8 100644 --- a/src/conn.rs +++ b/src/conn.rs @@ -366,25 +366,31 @@ impl Conn { self } - /// Accept a WebSocket upgrade on this request (RFC 6455 §4.2). + /// Accept a WebSocket upgrade on this request (RFC 6455 §4.2), + /// handing `handler` to the duplex loop that takes the socket over. /// /// On a valid handshake: status `101`, the `upgrade`/`connection`/ - /// `sec-websocket-accept` response headers, the [`WsUpgrade`] marker - /// the connection actor acts on, and the pipeline is halted. + /// `sec-websocket-accept` response headers, the [`WsUpgrade`] payload + /// the connection actor acts on, and the pipeline is halted. After + /// the 101 the connection actor leaves HTTP and drives `handler` — + /// see [`WsHandler`](crate::ws::WsHandler) for the callback contract. /// - /// On an invalid handshake: the appropriate rejection (`400`, or - /// `426` + `sec-websocket-version: 13` for a version mismatch), - /// empty body, pipeline halted; the connection stays plain HTTP with - /// normal keep-alive semantics. Pre-check with + /// On an invalid handshake: `handler` is dropped untouched and the + /// appropriate rejection goes out (`400`, or `426` + + /// `sec-websocket-version: 13` for a version mismatch), empty body, + /// pipeline halted; the connection stays plain HTTP with normal + /// keep-alive semantics. Pre-check with /// [`ws::handshake::is_upgrade_request`](crate::ws::handshake::is_upgrade_request) /// to branch before committing. /// /// [`WsUpgrade`]: crate::ws::WsUpgrade - pub fn upgrade(mut self) -> Self { + pub fn upgrade(mut self, handler: impl crate::ws::WsHandler) -> Self { use crate::ws::{handshake, Rejection}; match handshake::validate(&self) { Ok(accept) => { - self.upgrade = Some(crate::ws::WsUpgrade { _priv: () }); + self.upgrade = Some(crate::ws::WsUpgrade { + handler: Box::new(handler), + }); self.put_status(101) .put_header("upgrade", "websocket") .put_header("connection", "Upgrade") diff --git a/src/conn_actor.rs b/src/conn_actor.rs index 52adb7e..40fa3a3 100644 --- a/src/conn_actor.rs +++ b/src/conn_actor.rs @@ -59,6 +59,13 @@ pub struct ConnLimits { /// A client that stops reading mid-response is dropped when its /// socket buffer fills and a write stalls past the budget. pub write_timeout: Duration, + /// WebSocket: hard cap on a single frame's payload, enforced from + /// the frame header BEFORE the payload is buffered. Violation closes + /// with 1009. + pub max_frame_payload: usize, + /// WebSocket: hard cap on a complete (reassembled) message; spans + /// fragments. Violation closes with 1009. + pub max_message_bytes: usize, } impl Default for ConnLimits { @@ -71,6 +78,8 @@ impl Default for ConnLimits { keep_alive_timeout: Duration::from_secs(60), request_timeout: Duration::from_secs(30), write_timeout: Duration::from_secs(30), + max_frame_payload: 1024 * 1024, + max_message_bytes: 4 * 1024 * 1024, } } } @@ -231,17 +240,25 @@ pub fn run_connection( .put_body(RespBody::Empty); } - // ----- 3.5 WebSocket upgrade short-circuit. ----- - // An accepted handshake (marker + 101) ends HTTP on this socket: - // write the 101 head and leave the request loop. Chunk 3 (duplex - // wiring) takes the socket over right here; until then leaving - // the loop closes it via OwnedFd::drop — the 101 is still the - // honest, testable wire artefact. The status check is defensive: - // a post-handler plug that clobbered the 101 forfeits the - // upgrade and falls through to plain HTTP. + // ----- 3.5 WebSocket upgrade. ----- + // An accepted handshake (payload + 101) ends HTTP on this socket: + // write the 101 head, hand the fd to the duplex loop with the + // boxed handler and any bytes already read past this request (a + // client may pipeline its first frame behind the handshake — those + // bytes are ws bytes now). The registry entry stays Busy for the + // whole ws lifetime: an open WebSocket is in-flight work, not + // reapable idle HTTP; graceful shutdown force-stops it out of the + // select park at the drain deadline. The status check is + // defensive: a post-handler plug that clobbered the 101 forfeits + // the upgrade and falls through to plain HTTP. if response_conn.upgrade.is_some() && response_conn.status == Some(101) { let head = parser::serialise_response(&response_conn, true); - let _ = write_all(raw, &head, Instant::now() + limits.write_timeout); + if write_all(raw, &head, Instant::now() + limits.write_timeout).is_err() { + return; + } + let upgrade = response_conn.upgrade.take().expect("checked above"); + buf.drain(..head_len + consumed_past_head); + crate::ws::duplex::run_duplex(raw, buf, upgrade.handler, &limits); return; } @@ -550,7 +567,7 @@ fn read_chunked_body( // `ErrorKind::TimedOut` when `deadline` passes before the fd turns // readable, or the last io error. -fn read_some( +pub(crate) fn read_some( fd: RawFd, buf: &mut Vec, chunk: usize, @@ -616,7 +633,7 @@ fn try_write_once(fd: RawFd, buf: &[u8]) { // are down — a client that stops reading must not pin this actor in // wait_writable forever (the write-side twin of slowloris). -fn write_all(fd: RawFd, mut buf: &[u8], deadline: Instant) -> io::Result<()> { +pub(crate) fn write_all(fd: RawFd, mut buf: &[u8], deadline: Instant) -> io::Result<()> { while !buf.is_empty() { // Park on writability before each syscall, bounded by the budget. let remaining = deadline.saturating_duration_since(Instant::now()); diff --git a/src/lib.rs b/src/lib.rs index 551a2a7..cb910ab 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -34,6 +34,7 @@ pub use conn::{Assigns, Body, Conn, HeaderMap, HttpVersion, Method, Params, Resp pub use plug::{Next, Pipeline, Plug}; pub use router::Router; pub use sse::{EventSender, SseClosed}; +pub use ws::{Message, WsClosed, WsHandler, WsSender}; pub use serve::{ serve, serve_with, serve_with_shutdown, shutdown_handle, Config, Handle, ShutdownSignal, }; diff --git a/src/serve.rs b/src/serve.rs index 9135937..98ae7b5 100644 --- a/src/serve.rs +++ b/src/serve.rs @@ -47,6 +47,12 @@ pub struct Config { /// connections are closed immediately on shutdown and do not run the /// clock out. pub drain_timeout: Duration, + /// WebSocket: cap on a single frame's payload (header-checked + /// before buffering; violation closes 1009). + pub max_frame_payload: usize, + /// WebSocket: cap on a complete reassembled message (spans + /// fragments; violation closes 1009). + pub max_message_bytes: usize, /// Number of smarm scheduler OS threads. `None` means smarm's default /// (one per CPU). Set this to a small fixed number in tests so multiple /// concurrent test servers don't oversubscribe the host. @@ -69,6 +75,8 @@ impl Config { write_timeout: Duration::from_secs(30), max_body_bytes: 16 * 1024 * 1024, drain_timeout: Duration::from_secs(30), + max_frame_payload: 1024 * 1024, + max_message_bytes: 4 * 1024 * 1024, scheduler_threads: None, } } @@ -82,6 +90,8 @@ impl Config { keep_alive_timeout: self.keep_alive_timeout, request_timeout: self.request_timeout, write_timeout: self.write_timeout, + max_frame_payload: self.max_frame_payload, + max_message_bytes: self.max_message_bytes, } } } diff --git a/src/ws/duplex.rs b/src/ws/duplex.rs new file mode 100644 index 0000000..c1276dd --- /dev/null +++ b/src/ws/duplex.rs @@ -0,0 +1,398 @@ +//! The WebSocket duplex loop — v0.4 chunk 3. +//! +//! After an accepted upgrade the connection actor calls [`run_duplex`], +//! which replaces the HTTP request loop for the rest of the socket's +//! life. ONE actor speaks both directions (the v0.4 topology decision): +//! each iteration parks on a two-arm select — +//! +//! try_select(&[&outbound_rx, &FdArm::readable(fd)]) +//! +//! — smarm's RFC 008 fd arms composed with a channel arm on one wait +//! epoch; urus is deliberately the first consumer of that machinery. +//! The outbound arm sits at index 0 (select is priority-ordered): owed +//! writes drain before we read more, which is honest backpressure. +//! Accepted gap, by design: while a large outbound frame is mid-write +//! the actor isn't reading, so a ping arriving during that write is +//! answered after it completes. +//! +//! Handler model (the v0.4 chunk-3 decisions, mirroring SSE): +//! - [`WsHandler`] callbacks run IN the connection actor, inside the +//! select loop. A slow `on_message` stops reads — backpressure, not a +//! bug. Handlers needing concurrency spawn their own actor and hand it +//! a [`WsSender`] clone, exactly like an SSE producer. +//! - [`WsSender`] is the clonable outbound handle; every method returns +//! `Err(WsClosed)` once the connection is gone. +//! - Control frames are invisible to the handler in v1: pings are +//! auto-ponged (RFC 6455 §5.5.2 MUST), pongs are absorbed, a peer +//! close is auto-echoed (§5.5.1) and surfaces only as `on_close`. +//! `on_ping`/`on_pong` hooks can land later as default trait methods +//! without breaking anyone. +//! +//! Liveness mirrors SSE: there is no request clock on an open +//! WebSocket. An idle connection parks in the select (stoppable by a +//! draining registry — graceful shutdown force-stops it at the drain +//! deadline); a dead client is caught when a write stalls past +//! `write_timeout`. Reads have no deadline of their own — the select +//! only wakes the read path when bytes are pending. + +use crate::conn_actor::{read_some, write_all, ConnLimits}; +use crate::ws::frame::{self, Assembler, Frame, FrameError, Message, Opcode}; + +use smarm::FdArm; + +use std::os::fd::RawFd; +use std::time::Instant; + +// --------------------------------------------------------------------------- +// Handler trait +// --------------------------------------------------------------------------- + +/// Application callbacks for an upgraded WebSocket connection. Both run +/// inside the connection actor's select loop — see the module docs for +/// what that implies (backpressure, spawn-for-concurrency). +/// +/// A panic in `on_message` closes the connection with a `1011 internal +/// error` close frame and skips `on_close` (the handler is not re-entered +/// after it panicked). +pub trait WsHandler: Send + 'static { + /// A complete data message (fragments already reassembled, text + /// already UTF-8 validated) arrived from the client. + fn on_message(&mut self, msg: Message, sender: &WsSender); + + /// The connection is over. `code`/`reason` are what the PEER said: + /// `Some` iff a close frame was received from the client (whether it + /// initiated or echoed ours); `None` for everything wordless — EOF, + /// write failure, or a close handshake that timed out. Called exactly + /// once, last, on every exit path except a handler panic or a + /// force-stop unwind. + fn on_close(&mut self, code: Option, reason: &str) { + let _ = (code, reason); + } +} + +// --------------------------------------------------------------------------- +// WsSender — the clonable outbound handle +// --------------------------------------------------------------------------- + +/// The connection is gone (client disconnect, write timeout, close +/// handshake completed, or server shutdown). Producers should exit when +/// they see this — the exact contract of `SseClosed`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct WsClosed; + +impl std::fmt::Display for WsClosed { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("WebSocket connection closed") + } +} +impl std::error::Error for WsClosed {} + +pub(crate) enum Outbound { + Msg(Message), + Ping, + Close(u16, String), +} + +/// Outbound handle for an upgraded connection. Clonable — hand copies to +/// producer actors freely; the connection outlives none of them (sends +/// after it ends return `Err(WsClosed)`). +/// +/// v1 sends each message as a single unfragmented frame; outbound +/// fragmentation is not implemented. +#[derive(Clone)] +pub struct WsSender { + pub(crate) tx: smarm::Sender, +} + +impl WsSender { + /// Queue a data message for the client. + pub fn send(&self, msg: Message) -> Result<(), WsClosed> { + self.tx.send(Outbound::Msg(msg)).map_err(|_| WsClosed) + } + + /// Sugar: queue a text message. + pub fn text(&self, s: impl Into) -> Result<(), WsClosed> { + self.send(Message::Text(s.into())) + } + + /// Sugar: queue a binary message. + pub fn binary(&self, b: impl Into>) -> Result<(), WsClosed> { + self.send(Message::Binary(b.into())) + } + + /// Queue a ping (empty payload). Application-level liveness beyond + /// the built-in write-timeout detection; the pong is absorbed by the + /// loop in v1. + pub fn ping(&self) -> Result<(), WsClosed> { + self.tx.send(Outbound::Ping).map_err(|_| WsClosed) + } + + /// Initiate the close handshake: a close frame with `code`/`reason` + /// goes out, the loop then waits (bounded by `write_timeout`) for the + /// peer's echo before dropping the socket. `reason` is truncated to + /// fit the 125-byte control payload on a char boundary. + pub fn close(&self, code: u16, reason: &str) -> Result<(), WsClosed> { + self.tx + .send(Outbound::Close(code, reason.to_string())) + .map_err(|_| WsClosed) + } +} + +// --------------------------------------------------------------------------- +// The loop +// --------------------------------------------------------------------------- + +/// Run the duplex loop until the connection ends. `buf` carries any bytes +/// already read past the upgrade request's head — a client may pipeline +/// its first frame behind the handshake, and those bytes belong to us +/// now. +pub(crate) fn run_duplex( + raw: RawFd, + mut buf: Vec, + mut handler: Box, + limits: &ConnLimits, +) { + let (tx, rx) = smarm::channel::(); + // We hold this sender for handler callbacks, so the rx arm can never + // report closed while the loop runs — no closed-arm starvation case. + let sender = WsSender { tx }; + let mut asm = Assembler::new(limits.max_message_bytes); + + loop { + // ----- 1. Decode everything already buffered. ----- + // Runs before the first select so a pipelined first frame is + // served without waiting for fresh readability. + loop { + let (frame, consumed) = + match frame::decode(&buf, true, limits.max_frame_payload) { + Ok(Some(hit)) => hit, + Ok(None) => break, // incomplete; need more bytes + Err(e) => { + let peer = close_and_drain(raw, e, &mut buf, limits); + finish(&mut handler, peer); + return; + } + }; + buf.drain(..consumed); + + match frame.opcode { + Opcode::Ping => { + // §5.5.2 MUST pong, payload echoed. Written inline — + // "as soon as practical" beats channel ordering. + let pong = Frame::new(Opcode::Pong, frame.payload); + if write_frame(raw, &pong, limits).is_err() { + finish(&mut handler, None); + return; + } + } + Opcode::Pong => {} // invisible in v1 + Opcode::Close => { + // Peer-initiated close: echo its status code (§5.5.1) + // and end. The server side closes the TCP connection + // first by design (§7.1.1 puts that burden on us). + let echo_payload = + frame.payload.get(..2).map(<[u8]>::to_vec).unwrap_or_default(); + let peer = match frame::parse_close_payload(&frame.payload) { + Ok(cp) => { + let _ = write_frame( + raw, + &Frame::new(Opcode::Close, echo_payload), + limits, + ); + cp + } + Err(e) => { + // Malformed close payload: answer with the + // mapped code instead of an echo. + let _ = write_frame( + raw, + &Frame::new( + Opcode::Close, + frame::close_payload(e.close_code(), ""), + ), + limits, + ); + None + } + }; + finish_with(&mut handler, peer); + return; + } + _ => match asm.push(frame) { + Ok(Some(msg)) => { + // A panicking handler must not swallow smarm's stop + // sentinel — same re-raise dance as the pipeline. + let r = std::panic::catch_unwind( + std::panic::AssertUnwindSafe(|| { + handler.on_message(msg, &sender) + }), + ); + if r.is_err() { + smarm::preempt::check_cancelled(); + // Genuine handler panic: 1011 out, no further + // handler calls (no on_close). + let _ = close_and_drain_code( + raw, 1011, "", &mut buf, limits, + ); + return; + } + } + Ok(None) => {} // mid-fragmentation + Err(e) => { + let peer = close_and_drain(raw, e, &mut buf, limits); + finish(&mut handler, peer); + return; + } + }, + } + } + + // ----- 2. Park on first-of outbound / fd-readable. ----- + // Outbound at index 0: select is priority-ordered, so owed writes + // go out before we read more. try_select over select: an fd-arm + // registration failure (EBADF after a peer RST, EMFILE on the + // epoll set) is a connection event, not a crash. NOTE from the + // RFC 008 review: an `AlreadyExists` here would mean a stale + // waiter survived select's eager cleanup — that is a smarm bug + // candidate; capture it with --features smarm-trace and report, + // do not paper over. + let fd_arm = FdArm::readable(raw); + let idx = match smarm::try_select(&[&rx, &fd_arm]) { + Ok(i) => i, + Err(_) => { + finish(&mut handler, None); + return; + } + }; + + if idx == 0 { + // Outbound arm won. Single receiver + we hold a sender: a + // ready arm means a message is queued (Empty/closed are + // impossible here, but stay defensive). + let Ok(Some(out)) = rx.try_recv() else { continue }; + match out { + Outbound::Msg(msg) => { + let f = match msg { + Message::Text(s) => Frame::new(Opcode::Text, s.into_bytes()), + Message::Binary(b) => Frame::new(Opcode::Binary, b), + }; + if write_frame(raw, &f, limits).is_err() { + finish(&mut handler, None); + return; + } + } + Outbound::Ping => { + let f = Frame::new(Opcode::Ping, Vec::new()); + if write_frame(raw, &f, limits).is_err() { + finish(&mut handler, None); + return; + } + } + Outbound::Close(code, reason) => { + // We initiate: close out, bounded wait for the echo. + let peer = close_and_drain_code(raw, code, &reason, &mut buf, limits); + finish(&mut handler, peer); + return; + } + } + } else { + // Readable. The deadline is nominal — the fd already reported + // ready, this returns without a fresh park in practice. + match read_some( + raw, + &mut buf, + limits.initial_read_buf, + Instant::now() + limits.write_timeout, + ) { + Ok(0) => { + // EOF without a close frame: abnormal but common. + finish(&mut handler, None); + return; + } + Ok(_) => {} + Err(_) => { + finish(&mut handler, None); + return; + } + } + } + } +} + +/// `on_close(None, "")` — wordless endings (EOF, write failure, errors). +fn finish(handler: &mut Box, peer: Option<(u16, String)>) { + finish_with(handler, peer); +} + +/// `on_close` with whatever the peer said, panic-shielded like +/// `on_message` (stop sentinel re-raised; a panic here just ends the +/// already-ending connection). +fn finish_with(handler: &mut Box, peer: Option<(u16, String)>) { + let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + match &peer { + Some((code, reason)) => handler.on_close(Some(*code), reason), + None => handler.on_close(None, ""), + } + })); + if r.is_err() { + smarm::preempt::check_cancelled(); + } +} + +/// One frame onto the wire under a fresh `write_timeout` budget — the +/// streaming-chunk discipline, verbatim. +fn write_frame(raw: RawFd, f: &Frame, limits: &ConnLimits) -> std::io::Result<()> { + write_all(raw, &f.encode(), Instant::now() + limits.write_timeout) +} + +/// Server-initiated close for a frame/assembly error: map to the §7.4 +/// code and run the handshake tail. +fn close_and_drain( + raw: RawFd, + e: FrameError, + buf: &mut Vec, + limits: &ConnLimits, +) -> Option<(u16, String)> { + close_and_drain_code(raw, e.close_code(), "", buf, limits) +} + +/// Write our close frame, then wait — bounded by ONE `write_timeout` +/// budget — for the peer's close frame, discarding data frames (§1.4: an +/// endpoint that has sent close may discard incoming data). Returns what +/// the peer's close said, `None` if it never arrived (EOF, timeout, +/// garbage). The socket drops at the caller either way. +fn close_and_drain_code( + raw: RawFd, + code: u16, + reason: &str, + buf: &mut Vec, + limits: &ConnLimits, +) -> Option<(u16, String)> { + let ours = Frame::new(Opcode::Close, frame::close_payload(code, reason)); + if write_frame(raw, &ours, limits).is_err() { + return None; + } + + let deadline = Instant::now() + limits.write_timeout; + loop { + loop { + match frame::decode(buf, true, limits.max_frame_payload) { + Ok(Some((frame, consumed))) => { + buf.drain(..consumed); + if frame.opcode == Opcode::Close { + return frame::parse_close_payload(&frame.payload) + .ok() + .flatten(); + } + // Data/ping mid-handshake: discarded. + } + Ok(None) => break, + Err(_) => return None, // garbage during the tail: give up + } + } + match read_some(raw, buf, limits.initial_read_buf, deadline) { + Ok(0) | Err(_) => return None, + Ok(_) => {} + } + } +} diff --git a/src/ws/mod.rs b/src/ws/mod.rs index 0f59f80..8ac494c 100644 --- a/src/ws/mod.rs +++ b/src/ws/mod.rs @@ -12,17 +12,24 @@ //! vendoring. Base64 here is the 20-byte *encode* direction only and is //! not cryptographic, so that stays in-tree (`handshake::b64`). +pub mod duplex; pub mod frame; pub mod handshake; +pub use duplex::{WsClosed, WsHandler, WsSender}; pub use frame::{Frame, FrameError, Message, Opcode}; pub use handshake::Rejection; -/// Marker carried on a [`Conn`](crate::Conn) whose handler accepted a -/// WebSocket upgrade. Opaque: chunk 3 grows this into the handler/duplex -/// handoff payload, so nothing outside the crate should depend on its -/// shape. -#[derive(Debug)] +/// Payload carried on a [`Conn`](crate::Conn) whose handler accepted a +/// WebSocket upgrade: the boxed [`WsHandler`] the duplex loop will drive +/// (the chunk-3 growth this was reserved for). Still opaque outside the +/// crate. pub struct WsUpgrade { - pub(crate) _priv: (), + pub(crate) handler: Box, +} + +impl std::fmt::Debug for WsUpgrade { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("WsUpgrade { handler: .. }") + } } diff --git a/tests/integration.rs b/tests/integration.rs index e5cd91c..756d7ce 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -1019,88 +1019,264 @@ fn sse_heartbeats_on_silence() { // closes the socket; the tests assert the wire artefact, not ws traffic. // --------------------------------------------------------------------------- +const WS_HANDSHAKE: &[u8] = + b"GET /ws HTTP/1.1\r\nHost: x\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\ + Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n"; + +/// Echo handler: data messages come straight back; the text "bye" makes +/// the SERVER initiate the close handshake (1000/"goodbye"). +struct Echo; + +impl urus::WsHandler for Echo { + fn on_message(&mut self, msg: urus::Message, sender: &urus::WsSender) { + if matches!(&msg, urus::Message::Text(t) if t == "bye") { + let _ = sender.close(1000, "goodbye"); + return; + } + let _ = sender.send(msg); + } +} + fn ws_pipeline() -> Pipeline { - Pipeline::new().plug(Router::new().get("/ws", |c: Conn, _n: Next| c.upgrade())) + Pipeline::new().plug(Router::new().get("/ws", |c: Conn, _n: Next| c.upgrade(Echo))) } -#[test] -fn ws_handshake_101_with_accept_key() { - let port = spawn_server(ws_pipeline()); - let resp = send_request( - port, - b"GET /ws HTTP/1.1\r\nHost: x\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\ - Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n", - ); - assert_eq!(http_status(&resp), 101); - let head = String::from_utf8_lossy(&resp).to_lowercase(); - assert!(head.contains("upgrade: websocket"), "{head}"); - assert!(head.contains("connection: upgrade"), "{head}"); - // RFC 6455 §1.3 worked example key -> known accept value. - assert!( - head.contains("sec-websocket-accept: s3pplmbitxaq9kygzzhzrbk+xoo="), - "{head}" - ); - // 1xx carries no body framing. - assert!(!head.contains("content-length"), "{head}"); - assert!(!head.contains("transfer-encoding"), "{head}"); - assert_eq!(http_body(&resp), b"", "no body after a 101"); -} +// ----- client-side helpers (the test IS the ws client) ----- -#[test] -fn ws_handshake_wrong_version_426() { - let port = spawn_server(ws_pipeline()); - let resp = send_request( - port, - b"GET /ws HTTP/1.1\r\nHost: x\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\ - Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 8\r\n\r\n", - ); - assert_eq!(http_status(&resp), 426); - let head = String::from_utf8_lossy(&resp).to_lowercase(); - assert!(head.contains("sec-websocket-version: 13"), "{head}"); -} +use urus::ws::frame::{self as wsframe, Frame, Opcode}; -#[test] -fn ws_handshake_missing_key_400_keeps_http_alive() { - let port = spawn_server(ws_pipeline()); - // A rejected upgrade is a normal HTTP response: same connection must - // serve a second request afterwards. +/// Complete the handshake on a fresh socket; panics on anything but 101. +fn ws_connect(port: u16) -> TcpStream { let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap(); s.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); - s.write_all( - b"GET /ws HTTP/1.1\r\nHost: x\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\ - Sec-WebSocket-Version: 13\r\n\r\n", - ) - .unwrap(); - let mut buf = [0u8; 1024]; - let n = s.read(&mut buf).unwrap(); - let first = String::from_utf8_lossy(&buf[..n]); - assert!(first.starts_with("HTTP/1.1 400"), "{first}"); + s.write_all(WS_HANDSHAKE).unwrap(); + let head = String::from_utf8(read_response_head(&mut s)).unwrap(); + assert!(head.starts_with("HTTP/1.1 101"), "handshake: {head}"); + s +} - s.write_all( - b"GET /ws HTTP/1.1\r\nHost: x\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\ - Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n", - ) - .unwrap(); - let n = s.read(&mut buf).unwrap(); - let second = String::from_utf8_lossy(&buf[..n]); - assert!(second.starts_with("HTTP/1.1 101"), "{second}"); +fn ws_send(s: &mut TcpStream, frame: &Frame) { + s.write_all(&wsframe::encode_masked(frame, [0x11, 0x22, 0x33, 0x44])) + .unwrap(); +} + +/// Read one server frame (server frames are unmasked). `buf` carries +/// partial bytes across calls. +fn ws_read_frame(s: &mut TcpStream, buf: &mut Vec) -> Frame { + loop { + if let Some((f, consumed)) = + wsframe::decode(buf, false, 16 * 1024 * 1024).expect("bad server frame") + { + buf.drain(..consumed); + return f; + } + let mut tmp = [0u8; 4096]; + let n = s.read(&mut tmp).expect("read mid-frame"); + assert!(n > 0, "EOF mid-frame; buffered: {buf:?}"); + buf.extend_from_slice(&tmp[..n]); + } +} + +/// Assert the socket yields EOF (the server closed). +fn ws_expect_eof(s: &mut TcpStream, buf: &[u8]) { + assert!(buf.is_empty(), "unconsumed server bytes at EOF check: {buf:?}"); + let mut rest = Vec::new(); + let _ = s.read_to_end(&mut rest); // EOF or reset; both are closed + assert!(rest.is_empty(), "unexpected bytes before EOF: {rest:?}"); } #[test] -fn ws_accepted_upgrade_leaves_http_loop() { - // Chunk-1 semantics: after the 101 the actor exits and the fd drops — - // the client sees EOF after the 101 head. (Chunk 3 replaces the EOF - // with the ws duplex loop; this test will be rewritten then.) +fn ws_echo_roundtrip_text_and_binary() { let port = spawn_server(ws_pipeline()); - let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap(); - s.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); - s.write_all( - b"GET /ws HTTP/1.1\r\nHost: x\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\ - Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n", - ) - .unwrap(); + let mut s = ws_connect(port); let mut buf = Vec::new(); - s.read_to_end(&mut buf).unwrap(); // EOF must arrive, not a hang - let head = String::from_utf8_lossy(&buf); + + ws_send(&mut s, &Frame::new(Opcode::Text, "hello")); + let f = ws_read_frame(&mut s, &mut buf); + assert_eq!(f.opcode, Opcode::Text); + assert_eq!(f.payload, b"hello"); + + ws_send(&mut s, &Frame::new(Opcode::Binary, vec![0u8, 159, 146, 150])); + let f = ws_read_frame(&mut s, &mut buf); + assert_eq!(f.opcode, Opcode::Binary); + assert_eq!(f.payload, vec![0u8, 159, 146, 150]); +} + +/// A first frame pipelined in the SAME write as the handshake must be +/// served: bytes read past the 101 request head carry into the duplex +/// loop (the buf-handover requirement). +#[test] +fn ws_pipelined_first_frame_served() { + let port = spawn_server(ws_pipeline()); + let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap(); + s.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); + + let mut wire = WS_HANDSHAKE.to_vec(); + wire.extend_from_slice(&wsframe::encode_masked( + &Frame::new(Opcode::Text, "early"), + [9, 9, 9, 9], + )); + s.write_all(&wire).unwrap(); + + let head = String::from_utf8(read_response_head(&mut s)).unwrap(); assert!(head.starts_with("HTTP/1.1 101"), "{head}"); + let mut buf = Vec::new(); + let f = ws_read_frame(&mut s, &mut buf); + assert_eq!(f.payload, b"early"); +} + +/// §5.5.2: a ping is answered with a pong echoing the payload, with no +/// handler involvement. +#[test] +fn ws_ping_gets_pong_with_payload() { + let port = spawn_server(ws_pipeline()); + let mut s = ws_connect(port); + let mut buf = Vec::new(); + + ws_send(&mut s, &Frame::new(Opcode::Ping, "abc")); + let f = ws_read_frame(&mut s, &mut buf); + assert_eq!(f.opcode, Opcode::Pong); + assert_eq!(f.payload, b"abc"); + + // The connection is still a working websocket afterwards. + ws_send(&mut s, &Frame::new(Opcode::Text, "still here")); + assert_eq!(ws_read_frame(&mut s, &mut buf).payload, b"still here"); +} + +/// Client-initiated close: the server echoes the status code (§5.5.1) +/// and closes the TCP connection (§7.1.1 — server closes first). +#[test] +fn ws_client_close_is_echoed_then_eof() { + let port = spawn_server(ws_pipeline()); + let mut s = ws_connect(port); + let mut buf = Vec::new(); + + ws_send( + &mut s, + &Frame::new(Opcode::Close, wsframe::close_payload(1000, "done")), + ); + let f = ws_read_frame(&mut s, &mut buf); + assert_eq!(f.opcode, Opcode::Close); + let (code, _) = wsframe::parse_close_payload(&f.payload).unwrap().unwrap(); + assert_eq!(code, 1000); + ws_expect_eof(&mut s, &buf); +} + +/// Server-initiated close via WsSender::close: the close frame carries +/// our code/reason; after the client echoes, the socket closes. +#[test] +fn ws_server_initiated_close_handshake() { + let port = spawn_server(ws_pipeline()); + let mut s = ws_connect(port); + let mut buf = Vec::new(); + + ws_send(&mut s, &Frame::new(Opcode::Text, "bye")); + let f = ws_read_frame(&mut s, &mut buf); + assert_eq!(f.opcode, Opcode::Close); + let (code, reason) = wsframe::parse_close_payload(&f.payload).unwrap().unwrap(); + assert_eq!(code, 1000); + assert_eq!(reason, "goodbye"); + + // Echo the close; server completes the handshake and drops the fd. + ws_send( + &mut s, + &Frame::new(Opcode::Close, wsframe::close_payload(1000, "")), + ); + ws_expect_eof(&mut s, &buf); +} + +/// A protocol violation (unmasked client frame) starts the close +/// handshake with 1002. +#[test] +fn ws_unmasked_client_frame_closes_1002() { + let port = spawn_server(ws_pipeline()); + let mut s = ws_connect(port); + let mut buf = Vec::new(); + + // Server-style (unmasked) encoding from the client = violation. + s.write_all(&Frame::new(Opcode::Text, "naughty").encode()).unwrap(); + let f = ws_read_frame(&mut s, &mut buf); + assert_eq!(f.opcode, Opcode::Close); + let (code, _) = wsframe::parse_close_payload(&f.payload).unwrap().unwrap(); + assert_eq!(code, 1002); + + // Complete the handshake so the server returns inside the test budget. + ws_send( + &mut s, + &Frame::new(Opcode::Close, wsframe::close_payload(1002, "")), + ); + ws_expect_eof(&mut s, &buf); +} + +/// A frame whose HEADER already exceeds max_frame_payload is rejected +/// before any payload is buffered: 1009 with no payload bytes sent. +#[test] +fn ws_oversize_frame_header_closes_1009() { + let port = spawn_server(ws_pipeline()); + let mut s = ws_connect(port); + let mut buf = Vec::new(); + + // fin+text, masked, 64-bit length claiming 2 MiB (cap is 1 MiB). + let header: &[u8] = &[ + 0x81, 0xFF, 0, 0, 0, 0, 0, 0x20, 0, 0, // len = 0x200000 + 9, 9, 9, 9, // mask key + ]; + s.write_all(header).unwrap(); + let f = ws_read_frame(&mut s, &mut buf); + assert_eq!(f.opcode, Opcode::Close); + let (code, _) = wsframe::parse_close_payload(&f.payload).unwrap().unwrap(); + assert_eq!(code, 1009); + + ws_send( + &mut s, + &Frame::new(Opcode::Close, wsframe::close_payload(1009, "")), + ); + ws_expect_eof(&mut s, &buf); +} + +/// Fragmented text reassembles into one message; control frames may +/// interleave mid-fragmentation (§5.4) without disturbing assembly. +#[test] +fn ws_fragmented_message_with_interleaved_ping() { + let port = spawn_server(ws_pipeline()); + let mut s = ws_connect(port); + let mut buf = Vec::new(); + + let first = Frame { fin: false, opcode: Opcode::Text, payload: b"hel".to_vec() }; + ws_send(&mut s, &first); + ws_send(&mut s, &Frame::new(Opcode::Ping, "mid")); + let cont = Frame { fin: true, opcode: Opcode::Continuation, payload: b"lo".to_vec() }; + ws_send(&mut s, &cont); + + // Pong for the interleaved ping arrives first (written inline on + // receipt), then the reassembled echo. + let f = ws_read_frame(&mut s, &mut buf); + assert_eq!(f.opcode, Opcode::Pong); + assert_eq!(f.payload, b"mid"); + let f = ws_read_frame(&mut s, &mut buf); + assert_eq!(f.opcode, Opcode::Text); + assert_eq!(f.payload, b"hello"); +} + +/// An open WebSocket is in-flight work: graceful shutdown waits for the +/// drain deadline, then force-stops the conn actor parked in the duplex +/// select. serve returns; the client sees the socket drop. +#[test] +fn shutdown_force_stops_open_ws() { + let (port, handle, done_rx) = + spawn_server_with_handle(ws_pipeline(), Duration::from_millis(300)); + let mut s = ws_connect(port); + let mut buf = Vec::new(); + + // Prove the duplex loop is live before shutting down. + ws_send(&mut s, &Frame::new(Opcode::Text, "alive?")); + assert_eq!(ws_read_frame(&mut s, &mut buf).payload, b"alive?"); + + handle.shutdown(); + done_rx + .recv_timeout(Duration::from_secs(5)) + .expect("serve did not return: open ws not force-stopped at drain deadline"); + + let mut rest = Vec::new(); + let _ = s.read_to_end(&mut rest); // EOF or reset; both fine }