kyle/boxmaker
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gatewayd: ws conn, messages, pings, closing and a dead peer

Browse Source 
Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)
This commit is contained in:
kyle
2026-09-23 20:54:38 -07:00
parent f649ca9e3d
commit 74e08a5b1a
5 changed files with 583 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/gatewayd/src/ws/conn.rs
+208
View File
@@ -0,0 +1,208 @@
//! One WebSocket connection to Mattermost (M4a spec, section 6): open it, send text, and poll for
//! the next text message while answering pings, sending our own, and noticing a dead peer.
use std::io::{Read, Write};
use std::time::{Duration, Instant};
use crate::config::ServerUrl;
use crate::net::{Connector, Stream};
use crate::ws::WsError;
use crate::ws::frame::{CLOSE, Decoder, Incoming, PING, PONG, TEXT, encode};
use crate::ws::handshake::handshake;
/// Mattermost's WebSocket path.
pub const PATH: &str = "/api/v4/websocket";
/// How often we ping, and how long silence may last.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Timing {
pub ping_every: Duration,
pub dead_after: Duration,
}
pub struct Ws {
stream: Stream,
decoder: Decoder,
random: Box<dyn Read + Send>,
timing: Timing,
last_heard: Instant,
last_ping: Instant,
}
impl Ws {
/// Connect, and complete the handshake with `token`. `random` supplies the key and every mask
/// (in `gatewayd`, `/dev/urandom`).
pub fn open(
connector: &Connector,
token: &str,
timing: Timing,
mut random: Box<dyn Read + Send>,
) -> Result<Ws, WsError> {
// 1. `connector.connect(timing.dead_after)`; its error becomes
// WsError::Handshake(e.to_string()).
// 2. `handshake(&mut stream, &host_header(connector.server()), PATH, token, &mut random)?`.
// 3. A Ws with a new Decoder, and last_heard and last_ping both now.
let mut stream = connector
.connect(timing.dead_after)
.map_err(|e| WsError::Handshake(e.to_string()))?;
handshake(
&mut stream,
&host_header(connector.server()),
PATH,
token,
&mut random,
)?;
let now = Instant::now();
Ok(Ws {
stream,
decoder: Decoder::new(),
random,
timing,
last_heard: now,
last_ping: now,
})
}
fn send(&mut self, opcode: u8, payload: &[u8]) -> Result<(), WsError> {
// 4 mask bytes from `random` (read_exact), then write `encode(opcode, payload, mask)` and
// flush.
let mut mask = [0u8; 4];
self.random.read_exact(&mut mask)?;
let frame = encode(opcode, payload, mask);
self.stream.write_all(&frame)?;
self.stream.flush()?;
Ok(())
}
pub fn send_text(&mut self, text: &str) -> Result<(), WsError> {
// `send` with TEXT.
self.send(TEXT, text.as_bytes())
}
/// The next text message, or `None` after about `wait` with none. Pings are answered and sent
/// here; a close frame is answered and ends the connection (`Closed`); silence past the
/// dead-after limit is `Dead`. Written for you: it is the glue.
pub fn poll(&mut self, wait: Duration) -> Result<Option<String>, WsError> {
let until = Instant::now() + wait;
loop {
if let Some(text) = self.take_messages()? {
return Ok(Some(text));
}
let now = Instant::now();
self.keep_alive(now)?;
if now >= until {
return Ok(None);
}
let timeout = self.read_timeout(until, now);
self.read_some(timeout)?;
}
}
/// Act on every whole message the decoder holds, until one is text.
fn take_messages(&mut self) -> Result<Option<String>, WsError> {
// `while let Some(message) = self.decoder.next_message()? { match message { ... } }`:
// - Text(text) -> return Ok(Some(text)).
// - Ping(payload) -> `self.send(PONG, &payload)?`.
// - Pong(_) -> nothing.
// - Close(code, _) -> the reply is the code as 2 big-endian bytes, or empty:
// `code.map(|c| c.to_be_bytes().to_vec()).unwrap_or_default()`. Then
// `let _ = self.send(CLOSE, &reply);` (the peer may be gone) and Err(WsError::Closed).
// After the loop: Ok(None).
while let Some(message) = self.decoder.next_message()? {
match message {
Incoming::Text(text) => return Ok(Some(text)),
Incoming::Ping(payload) => self.send(PONG, &payload)?,
Incoming::Pong(_) => {}
Incoming::Close(code, _) => {
let reply = code.map(|c| c.to_be_bytes().to_vec()).unwrap_or_default();
let _ = self.send(CLOSE, &reply);
return Err(WsError::Closed);
}
}
}
Ok(None)
}
/// The dead-peer check, and our ping when one is due.
fn keep_alive(&mut self, now: Instant) -> Result<(), WsError> {
// 1. `now.duration_since(self.last_heard) >= self.timing.dead_after` -> Err(WsError::Dead).
// 2. `now.duration_since(self.last_ping) >= self.timing.ping_every` ->
// `self.send(PING, b"")?` and `self.last_ping = now`.
// 3. Ok(()).
if now.duration_since(self.last_heard) >= self.timing.dead_after {
return Err(WsError::Dead);
}
if now.duration_since(self.last_ping) >= self.timing.ping_every {
self.send(PING, b"")?;
self.last_ping = now;
}
Ok(())
}
/// How long the next read may wait: until the soonest of the end of `wait`, the next ping and
/// the dead-after limit, and never less than 1 ms.
fn read_timeout(&self, until: Instant, now: Instant) -> Duration {
// `(self.last_ping + self.timing.ping_every).saturating_duration_since(now)` is the time to
// the next ping; the same with last_heard and dead_after; and
// `until.saturating_duration_since(now)`. The least of the three (`.min`), then
// `.max(Duration::from_millis(1))`.
let next_ping = (self.last_ping + self.timing.ping_every).saturating_duration_since(now);
let next_dead = (self.last_heard + self.timing.dead_after).saturating_duration_since(now);
let until_left = until.saturating_duration_since(now);
next_ping
.min(next_dead)
.min(until_left)
.max(Duration::from_millis(1))
}
/// One read, at most `timeout` long, fed to the decoder.
fn read_some(&mut self, timeout: Duration) -> Result<(), WsError> {
// 1. `self.stream.set_read_timeout(Some(timeout))?`.
// 2. `let mut buf = [0u8; 16 * 1024];` and `match self.stream.read(&mut buf)`:
// - Ok(0) -> Err(WsError::Closed).
// - Ok(n) -> `self.decoder.feed(buf.get(..n).unwrap_or_default())`,
// `self.last_heard = Instant::now()`, Ok(()).
// - Err(e) whose `e.kind()` is WouldBlock, TimedOut or Interrupted -> Ok(()): nothing
// came, and `poll` goes round.
// - any other Err(e) -> Err(WsError::Io(e)).
self.stream.set_read_timeout(Some(timeout))?;
let mut buf = [0u8; 16 * 1024];
match self.stream.read(&mut buf) {
Ok(0) => Err(WsError::Closed),
Ok(n) => {
self.decoder.feed(buf.get(..n).unwrap_or_default());
self.last_heard = Instant::now();
Ok(())
}
Err(e)
if matches!(
e.kind(),
std::io::ErrorKind::WouldBlock
| std::io::ErrorKind::TimedOut
| std::io::ErrorKind::Interrupted
) =>
{
Ok(())
}
Err(e) => Err(WsError::Io(e)),
}
}
/// Send a close frame, best effort, and drop the connection.
pub fn close(mut self) {
// Send CLOSE with 1000 as 2 big-endian bytes; ignore the error.
let _ = self.send(CLOSE, &1000u16.to_be_bytes());
}
}
/// The `Host` header for a server: the port is written only when it is not the scheme's default.
pub fn host_header(server: &ServerUrl) -> String {
// The host alone when the port is the default for the scheme (443 for tls, 80 otherwise), else
// "host:port".
let default = if server.tls { 443 } else { 80 };
if server.port == default {
server.host.clone()
} else {
format!("{}:{}", server.host, server.port)
}
}
crates/gatewayd/src/ws/mod.rs
+1
View File
@@ -1,6 +1,7 @@
//! The WebSocket client (RFC 6455; M4a spec, section 6): the handshake, the frame codec, and the
//! connection that uses them.
pub mod conn;
pub mod frame;
pub mod handshake;
crates/gatewayd/tests/support/ws_server.rs
+132
View File
@@ -0,0 +1,132 @@
//! A scripted WebSocket server for tests: it accepts one handshake per connection and then lets
//! the test send raw frames and read the client's. Built on `tls_server`. Do not edit.
#![allow(dead_code)] // each test file uses a different part of this module
use std::io::{Read, Write};
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use gatewayd::ws::handshake::accept_for;
use rustls::ServerConfig;
use crate::tls_server::{Conn, serve};
pub struct Peer {
pub conn: Box<dyn Conn>,
/// The request head the client sent, for tests that check it.
pub request: String,
}
/// A frame from the client: opcode, whether it was masked, and the unmasked payload.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientFrame {
pub opcode: u8,
pub masked: bool,
pub mask: [u8; 4],
pub payload: Vec<u8>,
}
impl Peer {
pub fn send(&mut self, bytes: &[u8]) {
let _ = self.conn.write_all(bytes);
let _ = self.conn.flush();
}
/// An unmasked server frame, FIN set.
pub fn frame(&mut self, opcode: u8, payload: &[u8]) {
let mut out = vec![0x80 | opcode];
if payload.len() < 126 {
out.push(payload.len() as u8);
} else {
out.push(126);
out.extend_from_slice(&(payload.len() as u16).to_be_bytes());
}
out.extend_from_slice(payload);
self.send(&out);
}
pub fn text(&mut self, text: &str) {
self.frame(0x1, text.as_bytes());
}
fn read_exact(&mut self, n: usize) -> Option<Vec<u8>> {
let mut buf = vec![0u8; n];
self.conn.read_exact(&mut buf).ok()?;
Some(buf)
}
/// The next frame from the client, or `None` when it has gone.
pub fn read_frame(&mut self) -> Option<ClientFrame> {
let head = self.read_exact(2)?;
let opcode = head[0] & 0x0F;
let masked = head[1] & 0x80 != 0;
let len = match head[1] & 0x7F {
126 => u16::from_be_bytes(self.read_exact(2)?.try_into().ok()?) as usize,
127 => u64::from_be_bytes(self.read_exact(8)?.try_into().ok()?) as usize,
n => n as usize,
};
let mask: [u8; 4] = if masked {
self.read_exact(4)?.try_into().ok()?
} else {
[0; 4]
};
let raw = self.read_exact(len)?;
let payload = raw
.iter()
.zip(mask.iter().cycle())
.map(|(b, m)| b ^ m)
.collect();
Some(ClientFrame {
opcode,
masked,
mask,
payload,
})
}
pub fn pause(&self, d: Duration) {
std::thread::sleep(d);
}
}
/// Serve WebSocket connections: complete the handshake (or answer `refuse_with` instead), then run
/// `script` on the connection.
pub fn serve_ws<F>(
tls: Option<Arc<ServerConfig>>,
refuse_with: Option<&'static str>,
script: F,
) -> SocketAddr
where
F: Fn(Peer) + Send + Sync + 'static,
{
serve(tls, move |mut conn| {
let mut head = Vec::new();
let mut byte = [0u8; 1];
while !head.ends_with(b"\r\n\r\n") {
if conn.read(&mut byte).map(|n| n == 0).unwrap_or(true) {
return;
}
head.push(byte[0]);
}
let request = String::from_utf8_lossy(&head).into_owned();
if let Some(reply) = refuse_with {
let _ = conn.write_all(reply.as_bytes());
return;
}
let key = request
.lines()
.find_map(|l| l.strip_prefix("Sec-WebSocket-Key: "))
.unwrap_or_default()
.trim()
.to_string();
let reply = format!(
"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {}\r\n\r\n",
accept_for(&key)
);
let _ = conn.write_all(reply.as_bytes());
let _ = conn.flush();
script(Peer { conn, request });
})
}
crates/gatewayd/tests/ws_conn.rs
+241
View File
@@ -0,0 +1,241 @@
//! A WebSocket connection against a scripted server, plain and TLS: messages in order, pings both
//! ways, a dead peer, closing, and hostile input (M4a spec, section 6). Do not edit.
#[path = "support/tls_server.rs"]
mod tls_server;
#[path = "support/ws_server.rs"]
mod ws_server;
use std::io::Cursor;
use std::sync::mpsc;
use std::time::{Duration, Instant};
use gatewayd::config::ServerUrl;
use gatewayd::net::Connector;
use gatewayd::ws::WsError;
use gatewayd::ws::conn::{Timing, Ws, host_header};
use tls_server::{fixture, server_config};
use ws_server::{ClientFrame, serve_ws};
const SLOW: Timing = Timing {
ping_every: Duration::from_secs(60),
dead_after: Duration::from_secs(60),
};
/// Plenty of deterministic "random" bytes: the key, then masks.
fn random() -> Box<Cursor<Vec<u8>>> {
Box::new(Cursor::new(
(0..4096u32).map(|i| (i * 37 % 251) as u8).collect(),
))
}
fn open(port: u16, tls: bool, timing: Timing) -> Result<Ws, WsError> {
let url = ServerUrl {
tls,
host: "localhost".to_string(),
port,
};
let ca = tls.then(|| fixture("test-ca.pem"));
let c = Connector::new(url, ca.as_deref()).unwrap();
Ws::open(&c, "TOKEN", timing, random())
}
/// What the server saw, within 5 s: a missing frame fails the test instead of hanging it.
fn got<T>(rx: &mpsc::Receiver<T>) -> T {
rx.recv_timeout(Duration::from_secs(5))
.expect("the server saw nothing within 5 s")
}
/// Poll until a text message or an error, for at most 5 s.
fn next(ws: &mut Ws) -> Result<String, WsError> {
let until = Instant::now() + Duration::from_secs(5);
loop {
if let Some(t) = ws.poll(Duration::from_millis(200))? {
return Ok(t);
}
assert!(Instant::now() < until, "no message within 5 s");
}
}
#[test]
fn messages_arrive_in_order_plain_and_over_tls() {
for tls in [false, true] {
let (tx, rx) = mpsc::channel();
let config = tls.then(|| server_config("server"));
let addr = serve_ws(config, None, move |mut p| {
tx.send(p.request.clone()).unwrap();
p.text("{\"event\":\"hello\"}");
p.send(&[0x01, 0x03, b'o', b'n', b'e']);
p.send(&[0x80, 0x04, b'-', b't', b'w', b'o']);
p.text("three");
p.pause(Duration::from_secs(2));
});
let mut ws = open(addr.port(), tls, SLOW).unwrap();
let request = got(&rx);
assert!(
request.contains("Authorization: Bearer TOKEN\r\n"),
"{request}"
);
assert!(
request.starts_with("GET /api/v4/websocket HTTP/1.1\r\n"),
"{request}"
);
assert!(
request.contains(&format!("Host: localhost:{}\r\n", addr.port())),
"{request}"
);
assert_eq!(next(&mut ws).unwrap(), "{\"event\":\"hello\"}");
assert_eq!(next(&mut ws).unwrap(), "one-two");
assert_eq!(next(&mut ws).unwrap(), "three");
}
}
#[test]
fn a_ping_is_answered_with_the_same_payload() {
let (tx, rx) = mpsc::channel();
let addr = serve_ws(None, None, move |mut p| {
p.frame(0x9, b"are you there");
tx.send(p.read_frame()).unwrap();
p.text("after");
p.pause(Duration::from_secs(2));
});
let mut ws = open(addr.port(), false, SLOW).unwrap();
assert_eq!(next(&mut ws).unwrap(), "after");
let pong: ClientFrame = got(&rx).expect("a pong");
assert_eq!(
(pong.opcode, pong.masked, pong.payload.as_slice()),
(0xA, true, &b"are you there"[..])
);
}
#[test]
fn we_ping_on_schedule_and_every_frame_is_masked_differently() {
let (tx, rx) = mpsc::channel();
let addr = serve_ws(None, None, move |mut p| {
for _ in 0..3 {
tx.send(p.read_frame()).unwrap();
p.frame(0xA, b"");
}
p.pause(Duration::from_secs(2));
});
let timing = Timing {
ping_every: Duration::from_millis(100),
dead_after: Duration::from_secs(5),
};
let mut ws = open(addr.port(), false, timing).unwrap();
ws.send_text("first").unwrap();
let started = Instant::now();
while started.elapsed() < Duration::from_millis(350) {
let _ = ws.poll(Duration::from_millis(50)).unwrap();
}
let frames: Vec<ClientFrame> = (0..3).map(|_| got(&rx).unwrap()).collect();
assert_eq!(
(frames[0].opcode, frames[0].payload.as_slice()),
(0x1, &b"first"[..])
);
assert_eq!(frames[1].opcode, 0x9, "a ping after ping_every");
assert_eq!(frames[2].opcode, 0x9);
assert!(frames.iter().all(|f| f.masked));
assert_ne!(
frames[0].mask, frames[1].mask,
"a fresh mask for every frame"
);
}
#[test]
fn silence_is_a_dead_peer() {
let addr = serve_ws(None, None, |p| p.pause(Duration::from_secs(10)));
let timing = Timing {
ping_every: Duration::from_secs(60),
dead_after: Duration::from_millis(300),
};
let mut ws = open(addr.port(), false, timing).unwrap();
let started = Instant::now();
let err = loop {
match ws.poll(Duration::from_millis(100)) {
Ok(_) => assert!(
started.elapsed() < Duration::from_secs(3),
"never declared dead"
),
Err(e) => break e,
}
};
assert!(matches!(err, WsError::Dead), "{err}");
assert!(started.elapsed() < Duration::from_secs(1));
}
#[test]
fn a_peer_that_trickles_is_alive_and_its_message_arrives() {
let addr = serve_ws(None, None, |mut p| {
for b in [0x81u8, 0x05, b'd', b'r', b'i', b'p', b's'] {
p.send(&[b]);
p.pause(Duration::from_millis(100));
}
p.pause(Duration::from_secs(2));
});
let timing = Timing {
ping_every: Duration::from_secs(60),
dead_after: Duration::from_millis(400),
};
let mut ws = open(addr.port(), false, timing).unwrap();
assert_eq!(next(&mut ws).unwrap(), "drips");
}
#[test]
fn a_close_frame_is_answered_and_ends_the_connection() {
let (tx, rx) = mpsc::channel();
let addr = serve_ws(None, None, move |mut p| {
p.frame(0x8, &[0x03, 0xE8]);
tx.send(p.read_frame()).unwrap();
});
let mut ws = open(addr.port(), false, SLOW).unwrap();
assert!(matches!(next(&mut ws), Err(WsError::Closed)));
let reply = got(&rx).expect("a close in reply");
assert_eq!(
(reply.opcode, reply.payload.as_slice()),
(0x8, &[0x03u8, 0xE8][..])
);
}
#[test]
fn a_dropped_connection_is_closed() {
let addr = serve_ws(None, None, drop);
let mut ws = open(addr.port(), false, SLOW).unwrap();
assert!(matches!(next(&mut ws), Err(WsError::Closed)));
}
#[test]
fn a_hostile_frame_is_an_error_not_a_panic() {
let addr = serve_ws(None, None, |mut p| {
p.send(&[0x81, 0xFF, 0x80, 0, 0, 0, 0, 0, 0, 0]);
p.pause(Duration::from_secs(2));
});
let mut ws = open(addr.port(), false, SLOW).unwrap();
assert!(matches!(next(&mut ws), Err(WsError::Protocol(_))));
}
#[test]
fn a_refused_handshake() {
let addr = serve_ws(
None,
Some("HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\n\r\n"),
|_| {},
);
assert!(matches!(
open(addr.port(), false, SLOW),
Err(WsError::Handshake(_))
));
}
#[test]
fn the_host_header_names_the_port_only_when_it_is_not_the_default() {
let url = |tls, port| ServerUrl {
tls,
host: "chat.example".to_string(),
port,
};
assert_eq!(host_header(&url(true, 443)), "chat.example");
assert_eq!(host_header(&url(false, 80)), "chat.example");
assert_eq!(host_header(&url(true, 80)), "chat.example:80");
assert_eq!(host_header(&url(false, 8065)), "chat.example:8065");
}