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:
@@ -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,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`
|
||||
|
||||
Reference in New Issue
Block a user