gatewayd: ws frames, a strict decoder and a masked encoder

Implemented-By: OpenCode session (model recorded in docs/implementer-log.md)
This commit is contained in:
2026-09-23 20:50:00 -07:00
parent 02681069b6
commit f649ca9e3d
4 changed files with 727 additions and 0 deletions
+256
View File
@@ -0,0 +1,256 @@
//! WebSocket frames (RFC 6455, section 5), without I/O. `Decoder` is fed the bytes as they arrive,
//! in any pieces, and yields whole messages; `encode` builds our masked frames. Everything the
//! server sends is untrusted: every length is checked before anything is allocated.
use crate::ws::WsError;
/// The largest message we accept, counted from the length fields.
pub const MAX_MESSAGE: usize = 1 << 20;
pub const CONTINUATION: u8 = 0x0;
pub const TEXT: u8 = 0x1;
pub const CLOSE: u8 = 0x8;
pub const PING: u8 = 0x9;
pub const PONG: u8 = 0xA;
/// A whole message from the server.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Incoming {
Text(String),
Ping(Vec<u8>),
Pong(Vec<u8>),
/// A close frame: the status code if there is one, and the reason.
Close(Option<u16>, String),
}
/// Reassembles frames into messages.
#[derive(Debug, Default)]
pub struct Decoder {
buf: Vec<u8>,
/// A text message whose first frame has come and whose last has not.
partial: Option<Vec<u8>>,
}
/// A parsed header: what it says and how long it is.
struct Header {
fin: bool,
opcode: u8,
header_len: usize,
payload_len: usize,
}
impl Decoder {
pub fn new() -> Decoder {
Decoder::default()
}
/// Append bytes as they arrived.
pub fn feed(&mut self, bytes: &[u8]) {
self.buf.extend_from_slice(bytes);
}
/// The next whole message, `None` if more bytes are needed, or the error that ends the
/// connection. After an error, do not call again. Written for you: it is the glue.
pub fn next_message(&mut self) -> Result<Option<Incoming>, WsError> {
loop {
let Some(header) = self.header()? else {
return Ok(None);
};
let total = header.header_len.saturating_add(header.payload_len);
if self.buf.len() < total {
return Ok(None);
}
let payload: Vec<u8> = self
.buf
.get(header.header_len..total)
.unwrap_or_default()
.to_vec();
self.buf.drain(..total);
match header.opcode {
PING => return Ok(Some(Incoming::Ping(payload))),
PONG => return Ok(Some(Incoming::Pong(payload))),
CLOSE => return close(&payload).map(Some),
_ => {
// A text or continuation frame: a whole message, or wait for the next frame.
if let Some(text) = self.data_frame(header.opcode, header.fin, payload)? {
return Ok(Some(Incoming::Text(text)));
}
}
}
}
}
/// The next frame's header, once all of it is in `buf`, checked against every rule that does
/// not need the payload. It only reads `buf`. Written for you: it is the glue.
fn header(&self) -> Result<Option<Header>, WsError> {
let (Some(&b0), Some(&b1)) = (self.buf.first(), self.buf.get(1)) else {
return Ok(None);
};
check_first_bytes(b0, b1)?;
let fin = b0 & 0x80 != 0;
let opcode = b0 & 0x0F;
let Some((header_len, payload_len)) = self.length(b1 & 0x7F)? else {
return Ok(None);
};
if opcode & 0x8 != 0 {
check_control(fin, payload_len)?;
} else {
self.check_data(opcode, payload_len)?;
}
Ok(Some(Header {
fin,
opcode,
header_len,
payload_len,
}))
}
/// The header's length and the payload's, from the 7-bit length `short` in byte 1 and the
/// bytes after it. `Ok(None)` while those bytes have not all arrived.
fn length(&self, short: u8) -> Result<Option<(usize, usize)>, WsError> {
match short {
0..=125 => Ok(Some((2, usize::from(short)))),
126 => {
let Some(bytes) = self.buf.get(2..4) else {
return Ok(None);
};
let len = u16::from_be_bytes([bytes[0], bytes[1]]);
if len < 126 {
return Err(WsError::Protocol(
"a length not in its shortest form".to_string(),
));
}
Ok(Some((4, usize::from(len))))
}
127 => {
let Some(bytes) = self.buf.get(2..10) else {
return Ok(None);
};
let mut b = [0u8; 8];
b.copy_from_slice(bytes);
let len = u64::from_be_bytes(b);
if len >> 63 != 0 {
return Err(WsError::Protocol(
"a 64-bit length with its top bit set".to_string(),
));
}
if len <= 0xFFFF {
return Err(WsError::Protocol(
"a length not in its shortest form".to_string(),
));
}
Ok(Some((10, usize::try_from(len).unwrap_or(usize::MAX))))
}
_ => Err(WsError::Protocol(
"a length byte that is neither 126 nor 127".to_string(),
)),
}
}
/// The rules for a text or continuation frame that need what came before.
fn check_data(&self, opcode: u8, payload_len: usize) -> Result<(), WsError> {
if opcode == TEXT && self.partial.is_some() {
return Err(WsError::Protocol(
"a new message inside an unfinished one".to_string(),
));
}
if opcode == CONTINUATION && self.partial.is_none() {
return Err(WsError::Protocol(
"a continuation with nothing to continue".to_string(),
));
}
let so_far = self.partial.as_ref().map_or(0, Vec::len);
if payload_len > MAX_MESSAGE.saturating_sub(so_far) {
return Err(WsError::TooLarge);
}
Ok(())
}
/// A text or continuation frame's payload, taken out of `buf`: the whole message when `fin`
/// is set, or `None` when more frames must come.
fn data_frame(
&mut self,
opcode: u8,
fin: bool,
payload: Vec<u8>,
) -> Result<Option<String>, WsError> {
let mut message = if opcode == TEXT {
Vec::new()
} else {
self.partial.take().unwrap_or_default()
};
message.extend_from_slice(&payload);
if !fin {
self.partial = Some(message);
return Ok(None);
}
match String::from_utf8(message) {
Ok(text) => Ok(Some(text)),
Err(_) => Err(WsError::Protocol("text is not UTF-8".to_string())),
}
}
}
/// The rules on the first two bytes alone.
fn check_first_bytes(b0: u8, b1: u8) -> Result<(), WsError> {
if b0 & 0x70 != 0 {
return Err(WsError::Protocol("a reserved bit is set".to_string()));
}
if b1 & 0x80 != 0 {
return Err(WsError::Protocol(
"a frame from the server is masked".to_string(),
));
}
if !matches!(b0 & 0x0F, CONTINUATION | TEXT | CLOSE | PING | PONG) {
return Err(WsError::Protocol(format!("opcode {}", b0 & 0x0F)));
}
Ok(())
}
/// The rules for a control frame (close, ping, pong).
fn check_control(fin: bool, payload_len: usize) -> Result<(), WsError> {
if !fin {
return Err(WsError::Protocol("a fragmented control frame".to_string()));
}
if payload_len > 125 {
return Err(WsError::Protocol(
"a control frame over 125 bytes".to_string(),
));
}
Ok(())
}
/// A close frame's payload.
fn close(payload: &[u8]) -> Result<Incoming, WsError> {
match payload {
[] => Ok(Incoming::Close(None, String::new())),
[_] => Err(WsError::Protocol("a close frame of one byte".to_string())),
[a, b, reason @ ..] => {
let code = u16::from_be_bytes([*a, *b]);
match String::from_utf8(reason.to_vec()) {
Ok(reason) => Ok(Incoming::Close(Some(code), reason)),
Err(_) => Err(WsError::Protocol(
"a close reason that is not UTF-8".to_string(),
)),
}
}
}
}
/// One whole frame from us: FIN set, masked with `mask`.
pub fn encode(opcode: u8, payload: &[u8], mask: [u8; 4]) -> Vec<u8> {
let mut out = vec![0x80 | (opcode & 0x0F)];
let len = payload.len();
if len < 126 {
out.push(0x80 | u8::try_from(len).unwrap_or(0));
} else if let Ok(len16) = u16::try_from(len) {
out.push(0x80 | 126);
out.extend_from_slice(&len16.to_be_bytes());
} else {
out.push(0x80 | 127);
out.extend_from_slice(&u64::try_from(len).unwrap_or(u64::MAX).to_be_bytes());
}
out.extend_from_slice(&mask);
out.extend(payload.iter().zip(mask.iter().cycle()).map(|(b, m)| b ^ m));
out
}
+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 frame;
pub mod handshake;
/// Why a WebSocket ended or could not start. Every one of these ends the connection; `gatewayd`
+469
View File
@@ -0,0 +1,469 @@
//! WebSocket frames, adversarially (M4a spec, section 6). Everything a server sends is untrusted:
//! each hostile frame must end the connection with an error, never a panic, and a length must be
//! refused before anything waits for or allocates its payload. A seeded property test compares the
//! decoder with a deliberately naive one written here, on valid streams and on random mutations of
//! them, fed in random pieces. The seed is printed on failure. Do not edit.
use gatewayd::ws::WsError;
use gatewayd::ws::frame::{
CLOSE, CONTINUATION, Decoder, Incoming, MAX_MESSAGE, PING, PONG, TEXT, encode,
};
/// How a server frame's length is written: the shortest form, or a longer one on purpose.
#[derive(Clone, Copy)]
enum Len {
Short,
Force16,
Force64,
}
/// A frame as a server sends it (unmasked unless `masked`).
fn frame(fin: bool, rsv: u8, opcode: u8, masked: bool, payload: &[u8], form: Len) -> Vec<u8> {
let mut out = vec![(if fin { 0x80 } else { 0 }) | (rsv << 4) | opcode];
let m = if masked { 0x80 } else { 0 };
let len = payload.len();
match form {
Len::Short if len < 126 => out.push(m | len as u8),
Len::Short if len <= 0xFFFF => {
out.push(m | 126);
out.extend_from_slice(&(len as u16).to_be_bytes());
}
Len::Force16 => {
out.push(m | 126);
out.extend_from_slice(&(len as u16).to_be_bytes());
}
_ => {
out.push(m | 127);
out.extend_from_slice(&(len as u64).to_be_bytes());
}
}
if masked {
out.extend_from_slice(&[1, 2, 3, 4]);
}
out.extend_from_slice(payload);
out
}
fn text(s: &str) -> Vec<u8> {
frame(true, 0, TEXT, false, s.as_bytes(), Len::Short)
}
/// Feed `bytes` in pieces of `step` and collect every message, stopping at the first error.
fn decode(bytes: &[u8], step: usize) -> (Vec<Incoming>, Option<String>) {
let mut d = Decoder::new();
let mut got = Vec::new();
for piece in bytes.chunks(step.max(1)) {
d.feed(piece);
loop {
match d.next_message() {
Ok(Some(m)) => got.push(m),
Ok(None) => break,
Err(e) => return (got, Some(e.to_string())),
}
}
}
(got, None)
}
fn fails(bytes: &[u8], why: &str) {
for step in [1, 2, 3, 7, bytes.len().max(1)] {
let (_, err) = decode(bytes, step);
assert!(err.is_some(), "{why} (fed {step} at a time) was accepted");
}
}
#[test]
fn plain_messages() {
let mut bytes = text("hello");
bytes.extend(frame(true, 0, PING, false, b"p1", Len::Short));
bytes.extend(frame(true, 0, PONG, false, b"", Len::Short));
bytes.extend(frame(
true,
0,
CLOSE,
false,
&[0x03, 0xE8, b'b', b'y', b'e'],
Len::Short,
));
let (got, err) = decode(&bytes, bytes.len());
assert_eq!(err, None);
assert_eq!(
got,
[
Incoming::Text("hello".into()),
Incoming::Ping(b"p1".to_vec()),
Incoming::Pong(Vec::new()),
Incoming::Close(Some(1000), "bye".into()),
]
);
assert_eq!(
decode(&frame(true, 0, CLOSE, false, b"", Len::Short), 1).0,
[Incoming::Close(None, String::new())]
);
}
#[test]
fn fragments_reassemble_with_control_frames_between_and_utf8_split_across_them() {
let snow = "snow ☃ man";
let bytes_of = snow.as_bytes();
let cut = snow.find('☃').unwrap() + 1; // inside the three-byte character
let mut bytes = frame(false, 0, TEXT, false, &bytes_of[..cut], Len::Short);
bytes.extend(frame(true, 0, PING, false, b"mid", Len::Short));
bytes.extend(frame(
false,
0,
CONTINUATION,
false,
&bytes_of[cut..cut + 1],
Len::Short,
));
bytes.extend(frame(
true,
0,
CONTINUATION,
false,
&bytes_of[cut + 1..],
Len::Short,
));
for step in 1..=bytes.len() {
let (got, err) = decode(&bytes, step);
assert_eq!(err, None, "step {step}");
assert_eq!(
got,
[Incoming::Ping(b"mid".to_vec()), Incoming::Text(snow.into())],
"step {step}"
);
}
}
#[test]
fn lengths_in_every_form() {
for len in [0usize, 1, 125, 126, 127, 65_535, 65_536, 100_000] {
let body = "x".repeat(len);
let (got, err) = decode(&text(&body), 4096);
assert_eq!(err, None, "{len}");
assert_eq!(got, [Incoming::Text(body)], "{len}");
}
}
#[test]
fn hostile_frames_end_the_connection() {
for rsv in [1, 2, 4] {
fails(
&frame(true, rsv, TEXT, false, b"x", Len::Short),
"a reserved bit",
);
}
fails(
&frame(true, 0, TEXT, true, b"x", Len::Short),
"a masked frame from the server",
);
for op in [2u8, 3, 7, 11, 15] {
fails(
&frame(true, 0, op, false, b"x", Len::Short),
"an unknown or binary opcode",
);
}
fails(
&frame(true, 0, PING, false, &[0u8; 126], Len::Short),
"a control frame over 125 bytes",
);
fails(
&frame(false, 0, PING, false, b"x", Len::Short),
"a fragmented control frame",
);
fails(
&frame(true, 0, CONTINUATION, false, b"x", Len::Short),
"a continuation with nothing to continue",
);
let mut inside = frame(false, 0, TEXT, false, b"a", Len::Short);
inside.extend(text("b"));
fails(&inside, "a new message inside an unfinished one");
fails(
&frame(true, 0, TEXT, false, b"x", Len::Force16),
"a 16-bit length for 1 byte",
);
fails(
&frame(true, 0, TEXT, false, &[b'y'; 200], Len::Force64),
"a 64-bit length for 200 bytes",
);
fails(
&frame(true, 0, TEXT, false, &[0xff, 0xfe], Len::Short),
"text that is not UTF-8",
);
fails(
&frame(true, 0, CLOSE, false, &[3], Len::Short),
"a close frame of one byte",
);
fails(
&frame(true, 0, CLOSE, false, &[3, 232, 0xff], Len::Short),
"a close reason that is not UTF-8",
);
}
#[test]
fn huge_lengths_are_refused_from_the_header_alone() {
// Only the header is fed: the decoder must refuse without waiting for a payload.
let top_bit = [0x81u8, 127, 0x80, 0, 0, 0, 0, 0, 0, 1];
let mut d = Decoder::new();
d.feed(&top_bit);
assert!(
d.next_message().is_err(),
"a 64-bit length with its top bit set"
);
let too_big = (MAX_MESSAGE as u64) + 1;
let mut head = vec![0x81u8, 127];
head.extend_from_slice(&too_big.to_be_bytes());
let mut d = Decoder::new();
d.feed(&head);
assert!(matches!(d.next_message(), Err(WsError::TooLarge)));
let mut head = vec![0x81u8, 127];
head.extend_from_slice(&0x7FFF_FFFF_FFFF_FFFFu64.to_be_bytes());
let mut d = Decoder::new();
d.feed(&head);
assert!(matches!(d.next_message(), Err(WsError::TooLarge)));
// Across fragments: the sum counts.
let half = MAX_MESSAGE / 2 + 1;
let mut d = Decoder::new();
d.feed(&frame(false, 0, TEXT, false, &vec![b'a'; half], Len::Short));
assert!(matches!(d.next_message(), Ok(None)));
let mut second = vec![0x00u8, 127];
second.extend_from_slice(&(half as u64).to_be_bytes());
d.feed(&second);
assert!(matches!(d.next_message(), Err(WsError::TooLarge)));
// Exactly the limit is fine.
let (got, err) = decode(&text(&"z".repeat(MAX_MESSAGE)), 65_536);
assert_eq!(err, None);
assert_eq!(got.len(), 1);
}
#[test]
fn our_frames_are_masked_and_decode_back() {
let mask = [0x11, 0x22, 0x33, 0x44];
for len in [0usize, 5, 125, 126, 65_535, 65_536] {
let payload: Vec<u8> = (0..len).map(|i| (i % 251) as u8).collect();
let bytes = encode(TEXT, &payload, mask);
assert_eq!(bytes[0], 0x80 | TEXT, "FIN and the opcode");
assert_ne!(bytes[1] & 0x80, 0, "the mask bit");
let (len_field, header) = match bytes[1] & 0x7F {
126 => (u16::from_be_bytes([bytes[2], bytes[3]]) as usize, 4),
127 => (
u64::from_be_bytes(bytes[2..10].try_into().unwrap()) as usize,
10,
),
n => (n as usize, 2),
};
assert_eq!(len_field, len);
let shortest = if len < 126 {
2
} else if len <= 0xFFFF {
4
} else {
10
};
assert_eq!(header, shortest, "the shortest length form");
assert_eq!(&bytes[header..header + 4], &mask);
let unmasked: Vec<u8> = bytes[header + 4..]
.iter()
.zip(mask.iter().cycle())
.map(|(b, m)| b ^ m)
.collect();
assert_eq!(unmasked, payload);
}
assert_eq!(encode(PONG, b"p", mask)[0], 0x80 | PONG);
}
// ---------- the property test ----------
struct Rng(u64);
impl Rng {
fn next(&mut self) -> u64 {
let mut x = self.0;
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
self.0 = x;
x
}
fn below(&mut self, n: usize) -> usize {
(self.next() % n.max(1) as u64) as usize
}
}
/// The naive decoder: the whole buffer at once, the rules written out plainly.
fn naive(bytes: &[u8]) -> (Vec<Incoming>, bool) {
let mut out = Vec::new();
let mut i = 0usize;
let mut partial: Option<Vec<u8>> = None;
while i < bytes.len() {
if bytes.len() - i < 2 {
return (out, false);
}
let (b0, b1) = (bytes[i], bytes[i + 1]);
let (fin, rsv, op, masked, short) = (
b0 >> 7 == 1,
(b0 >> 4) & 7,
b0 & 15,
b1 >> 7 == 1,
(b1 & 127) as usize,
);
if rsv != 0 || masked || ![0, 1, 8, 9, 10].contains(&op) {
return (out, true);
}
let (hl, len) = if short == 126 {
if bytes.len() - i < 4 {
return (out, false);
}
let l = u16::from_be_bytes([bytes[i + 2], bytes[i + 3]]) as usize;
if l < 126 {
return (out, true);
}
(4, l)
} else if short == 127 {
if bytes.len() - i < 10 {
return (out, false);
}
let l = u64::from_be_bytes(bytes[i + 2..i + 10].try_into().unwrap());
if l >> 63 == 1 || l <= 0xFFFF {
return (out, true);
}
(10, l as usize)
} else {
(2, short)
};
let control = op >= 8;
if control && (!fin || len > 125) {
return (out, true);
}
if !control {
if (op == 1 && partial.is_some()) || (op == 0 && partial.is_none()) {
return (out, true);
}
if partial.as_ref().map_or(0, |p| p.len()) + len > MAX_MESSAGE {
return (out, true);
}
}
if bytes.len() - i - hl < len {
return (out, false);
}
let payload = bytes[i + hl..i + hl + len].to_vec();
i += hl + len;
match op {
9 => out.push(Incoming::Ping(payload)),
10 => out.push(Incoming::Pong(payload)),
8 => {
if payload.len() == 1 {
return (out, true);
}
if payload.is_empty() {
out.push(Incoming::Close(None, String::new()));
} else {
match String::from_utf8(payload[2..].to_vec()) {
Ok(r) => out.push(Incoming::Close(
Some(u16::from_be_bytes([payload[0], payload[1]])),
r,
)),
Err(_) => return (out, true),
}
}
}
_ => {
let mut m = if op == 1 {
Vec::new()
} else {
partial.take().unwrap()
};
m.extend_from_slice(&payload);
if fin {
match String::from_utf8(m) {
Ok(t) => out.push(Incoming::Text(t)),
Err(_) => return (out, true),
}
} else {
partial = Some(m);
}
}
}
}
(out, false)
}
/// A random valid stream: text messages split into random fragments, with control frames between.
fn valid_stream(rng: &mut Rng) -> Vec<u8> {
let mut bytes = Vec::new();
for _ in 0..1 + rng.below(6) {
let len = [0, 1, 50, 125, 126, 300, 70_000][rng.below(7)];
let body: String = (0..len)
.map(|k| {
if (k + rng.below(3)).is_multiple_of(29) {
'é'
} else {
'a'
}
})
.collect();
let raw = body.as_bytes();
let parts = 1 + rng.below(3);
let mut cuts: Vec<usize> = (0..parts - 1).map(|_| rng.below(raw.len() + 1)).collect();
cuts.sort();
let mut start = 0;
for (k, cut) in cuts
.iter()
.copied()
.chain(std::iter::once(raw.len()))
.enumerate()
{
let op = if k == 0 { TEXT } else { CONTINUATION };
bytes.extend(frame(
k == parts - 1,
0,
op,
false,
&raw[start..cut],
Len::Short,
));
start = cut;
if rng.below(3) == 0 {
bytes.extend(frame(
true,
0,
PING,
false,
&[rng.below(256) as u8; 3],
Len::Short,
));
}
}
}
bytes
}
#[test]
fn random_streams_agree_with_the_naive_decoder() {
for case in 0..300u64 {
let seed = 0x9E37_79B9_7F4A_7C15 ^ (case * 7919 + 1);
let mut rng = Rng(seed);
let mut bytes = valid_stream(&mut rng);
if case % 2 == 1 {
// Mutate: flip a few random bits, so most streams break somewhere different.
for _ in 0..1 + rng.below(4) {
let at = rng.below(bytes.len());
bytes[at] ^= 1 << rng.below(8);
}
}
let (want, want_err) = naive(&bytes);
let step = 1 + rng.below(4096);
let (got, got_err) = decode(&bytes, step);
assert_eq!(got, want, "seed {seed:#x}, step {step}: messages differ");
assert_eq!(
got_err.is_some(),
want_err,
"seed {seed:#x}, step {step}: {got_err:?} vs naive error {want_err}"
);
}
}