feat(ws): duplex wiring + handler API + echo example (v0.4 chunk 3)
Topology (b) as agreed: ONE actor per connection via smarm RFC 008 fd
arms. After the 101, ws::duplex::run_duplex replaces the HTTP loop:
try_select(&[&outbound_rx, &FdArm::readable(fd)]) per iteration,
outbound at index 0 (priority order = owed writes drain before reads;
honest backpressure). Accepted gap documented: mid-write of a large
outbound frame the actor isn't reading.
Handler API (deferred questions, as answered this session):
- WsHandler { on_message, on_close } runs IN the conn actor's select
loop; concurrency = spawn an actor with a WsSender clone (the SSE
producer pattern). Conn::upgrade(self, handler) flat signature;
WsUpgrade grows from marker to Box<dyn WsHandler> payload.
- 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, drain timeout) = on_close(None, ""). on_ping/on_pong can
land later as default methods, non-breaking.
- Server-initiated close (WsSender::close, FrameError->1002/1009/1007,
handler panic->1011): close out, bounded drain (one write_timeout
budget) for the peer echo, data discarded (1.4). Handler panics
re-raise smarm's stop sentinel first (the pipeline catch_unwind
dance, replicated).
- Caps: Config.max_frame_payload (1 MiB) / max_message_bytes (4 MiB).
Conn actor: 101 branch now hands fd + leftover buf (pipelined first
frame carries over; tested) + boxed handler to run_duplex; registry
entry stays Busy for the ws lifetime, so graceful shutdown force-stops
the conn out of the select park at the drain deadline (tested — the
in-runtime request_stop wake covers select parks too).
Tests: chunk-1 EOF test rewritten into a 9-test duplex suite (echo,
pipelined first frame, ping/pong, both close directions, 1002 unmasked,
1009 header-cap, fragmentation with interleaved ping, shutdown
force-stop). Suite 59u+40i+2d. Hammer: 35x lifecycle+ws subset + 3 full
+ 1 full under smarm-trace, all green; no AlreadyExists out of
try_select (the fresh eager-cleanup path held). Validated against
python websocket-client (echo, pong payload, close 1000).
examples/ws_echo.rs: /echo in-actor + /clock producer-spawning.
This commit is contained in:
@@ -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<u16>, 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.
|
||||
|
||||
|
||||
+41
-17
@@ -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<Frame>`; 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
|
||||
|
||||
|
||||
@@ -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<u16>, 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");
|
||||
}
|
||||
+15
-9
@@ -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")
|
||||
|
||||
+28
-11
@@ -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<u8>,
|
||||
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());
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<u16>, 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<Outbound>,
|
||||
}
|
||||
|
||||
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<String>) -> Result<(), WsClosed> {
|
||||
self.send(Message::Text(s.into()))
|
||||
}
|
||||
|
||||
/// Sugar: queue a binary message.
|
||||
pub fn binary(&self, b: impl Into<Vec<u8>>) -> 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<u8>,
|
||||
mut handler: Box<dyn WsHandler>,
|
||||
limits: &ConnLimits,
|
||||
) {
|
||||
let (tx, rx) = smarm::channel::<Outbound>();
|
||||
// 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<dyn WsHandler>, 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<dyn WsHandler>, 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<u8>,
|
||||
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<u8>,
|
||||
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(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
+13
-6
@@ -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<dyn WsHandler>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for WsUpgrade {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str("WsUpgrade { handler: .. }")
|
||||
}
|
||||
}
|
||||
|
||||
+245
-69
@@ -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",
|
||||
)
|
||||
fn ws_send(s: &mut TcpStream, frame: &Frame) {
|
||||
s.write_all(&wsframe::encode_masked(frame, [0x11, 0x22, 0x33, 0x44]))
|
||||
.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}");
|
||||
}
|
||||
|
||||
/// Read one server frame (server frames are unmasked). `buf` carries
|
||||
/// partial bytes across calls.
|
||||
fn ws_read_frame(s: &mut TcpStream, buf: &mut Vec<u8>) -> 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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user