use std::io::{self, Read, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant};
use compcol::deflate::Deflate;
use compcol::limit::LimitedDecoder;
use compcol::vec::compress_to_vec;
use compcol::{Algorithm, Decoder, Status};
use purecrypto::hash::{Digest, Sha1};
use crate::error::{Error, Result};
use crate::url::Url;
const WS_GUID: &str = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
const DEFLATE_TAIL: [u8; 4] = [0x00, 0x00, 0xFF, 0xFF];
pub(crate) const OPCODE_CONT: u8 = 0x0;
pub(crate) const OPCODE_TEXT: u8 = 0x1;
pub(crate) const OPCODE_BINARY: u8 = 0x2;
pub(crate) const OPCODE_CLOSE: u8 = 0x8;
pub(crate) const OPCODE_PING: u8 = 0x9;
pub(crate) const OPCODE_PONG: u8 = 0xA;
pub(crate) const MAX_PAYLOAD_BYTES: u64 = 64 * 1024 * 1024;
pub(crate) const HANDSHAKE_DEADLINE: Duration = Duration::from_secs(60);
const PMD_OFFER: &str =
"permessage-deflate; client_no_context_takeover; server_no_context_takeover";
const NO_SUBPROTOCOLS: &[&str] = &[];
struct Pmd {
#[allow(dead_code)]
client_no_context_takeover: bool,
server_no_context_takeover: bool,
decoder: <Deflate as Algorithm>::Decoder,
}
impl Pmd {
fn inflate_message(&mut self, compressed: &[u8]) -> Result<Vec<u8>> {
if self.server_no_context_takeover {
self.decoder.reset();
}
let mut input = Vec::with_capacity(compressed.len() + DEFLATE_TAIL.len());
input.extend_from_slice(compressed);
input.extend_from_slice(&DEFLATE_TAIL);
let taken = std::mem::replace(&mut self.decoder, Deflate::decoder());
let mut limited = LimitedDecoder::new(taken, MAX_PAYLOAD_BYTES);
let result = Self::run_inflate(&mut limited, &input);
self.decoder = limited.into_inner();
let (out, stream_ended, consumed) = result?;
if stream_ended {
self.decoder.reset();
if consumed < compressed.len() {
return Err(Error::BadResponse(
"permessage-deflate: data after the final DEFLATE block".into(),
));
}
}
Ok(out)
}
fn run_inflate(
limited: &mut LimitedDecoder<<Deflate as Algorithm>::Decoder>,
input: &[u8],
) -> Result<(Vec<u8>, bool, usize)> {
let mut out: Vec<u8> = Vec::new();
let mut scratch = vec![0u8; 64 * 1024];
let mut consumed = 0usize;
let mut stream_ended = false;
loop {
let before_consumed = consumed;
let before_written = out.len();
let (p, status) = limited
.decode(&input[consumed..], &mut scratch)
.map_err(|e| {
Error::BadResponse(format!("permessage-deflate inflate failed: {e}"))
})?;
out.extend_from_slice(&scratch[..p.written]);
consumed += p.consumed;
match status {
Status::StreamEnd => {
stream_ended = true;
break;
}
Status::OutputFull => continue,
Status::InputEmpty => {
if consumed >= input.len()
|| (consumed == before_consumed && out.len() == before_written)
{
break;
}
}
}
}
Ok((out, stream_ended, consumed))
}
}
fn deflate_message(payload: &[u8]) -> Result<Vec<u8>> {
let mut out = compress_to_vec::<Deflate>(payload)
.map_err(|e| Error::BadResponse(format!("permessage-deflate deflate failed: {e}")))?;
if out.ends_with(&DEFLATE_TAIL) {
out.truncate(out.len() - DEFLATE_TAIL.len());
}
Ok(out)
}
pub fn fetch(url: &Url) -> Result<Vec<u8>> {
fetch_with(url, &crate::net::NetConfig::default())
}
pub(crate) fn fetch_with(url: &Url, cfg: &crate::net::NetConfig) -> Result<Vec<u8>> {
match url.scheme.as_str() {
"ws" => {
let mut sock = tcp_connect(url, cfg)?;
let (mut pmd, _proto) = handshake(&mut sock, url, NO_SUBPROTOCOLS)?;
read_data_and_close(&mut sock, pmd.as_mut())
}
"wss" => {
let tcp = tcp_connect(url, cfg)?;
let mut tls = cfg.tls_connect(tcp, &url.host)?;
let (mut pmd, _proto) = handshake(&mut tls, url, NO_SUBPROTOCOLS)?;
read_data_and_close(&mut tls, pmd.as_mut())
}
other => Err(Error::UnsupportedScheme(other.to_string())),
}
}
fn tcp_connect(url: &Url, cfg: &crate::net::NetConfig) -> Result<Box<dyn crate::net::NetStream>> {
let stream = cfg.connect(&url.host, url.port)?;
stream.set_read_timeout(cfg.io_timeout())?;
stream.set_write_timeout(Some(Duration::from_secs(60)))?;
Ok(stream)
}
const SEND_TIMEOUT: Duration = Duration::from_secs(60);
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WsMessage {
Text(String),
Binary(Vec<u8>),
}
impl WsMessage {
pub fn as_bytes(&self) -> &[u8] {
match self {
WsMessage::Text(s) => s.as_bytes(),
WsMessage::Binary(b) => b,
}
}
pub fn as_text(&self) -> Option<&str> {
match self {
WsMessage::Text(s) => Some(s),
WsMessage::Binary(_) => None,
}
}
pub fn into_bytes(self) -> Vec<u8> {
match self {
WsMessage::Text(s) => s.into_bytes(),
WsMessage::Binary(b) => b,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WsClose {
pub code: u16,
pub reason: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WsOpcode {
Continuation,
Text,
Binary,
Close,
Ping,
Pong,
}
impl WsOpcode {
fn to_u8(self) -> u8 {
match self {
WsOpcode::Continuation => OPCODE_CONT,
WsOpcode::Text => OPCODE_TEXT,
WsOpcode::Binary => OPCODE_BINARY,
WsOpcode::Close => OPCODE_CLOSE,
WsOpcode::Ping => OPCODE_PING,
WsOpcode::Pong => OPCODE_PONG,
}
}
fn from_u8(op: u8) -> Result<WsOpcode> {
Ok(match op {
OPCODE_CONT => WsOpcode::Continuation,
OPCODE_TEXT => WsOpcode::Text,
OPCODE_BINARY => WsOpcode::Binary,
OPCODE_CLOSE => WsOpcode::Close,
OPCODE_PING => WsOpcode::Ping,
OPCODE_PONG => WsOpcode::Pong,
other => return Err(Error::BadResponse(format!("unknown WS opcode 0x{other:x}"))),
})
}
fn is_control(self) -> bool {
matches!(self, WsOpcode::Close | WsOpcode::Ping | WsOpcode::Pong)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WsEvent {
Text(String),
Binary(Vec<u8>),
Ping(Vec<u8>),
Pong(Vec<u8>),
Close(Option<WsClose>),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WsFrame {
pub fin: bool,
pub opcode: WsOpcode,
pub payload: Vec<u8>,
}
trait ReadWrite: Read + Write + Send {}
impl<T: Read + Write + Send + ?Sized> ReadWrite for T {}
enum WsTransportKind {
Plain {
read: Mutex<Box<dyn crate::net::NetStream>>,
write: Mutex<Box<dyn crate::net::NetStream>>,
},
Tls(Box<crate::tls::TlsConn>),
Shared(Mutex<Box<dyn ReadWrite>>),
}
const SHUTDOWN_POLL: Duration = Duration::from_millis(250);
struct WsTransport {
kind: WsTransportKind,
shutdown: Arc<AtomicBool>,
read_timeout: Mutex<Option<Duration>>,
last_write: Mutex<Instant>,
broken: AtomicBool,
close_gate: Mutex<()>,
}
fn is_read_timeout(e: &io::Error) -> bool {
matches!(
e.kind(),
io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut
) && !crate::net::is_deadline_exceeded(e)
}
impl WsTransport {
fn new(kind: WsTransportKind) -> WsTransport {
WsTransport {
kind,
shutdown: Arc::new(AtomicBool::new(false)),
read_timeout: Mutex::new(Some(SEND_TIMEOUT)),
last_write: Mutex::new(Instant::now()),
broken: AtomicBool::new(false),
close_gate: Mutex::new(()),
}
}
fn write_unless_closed(&self, send_closed: &AtomicBool, frame: &[u8]) -> Result<()> {
let _gate = self.close_gate.lock().unwrap_or_else(|e| e.into_inner());
if send_closed.load(Ordering::SeqCst) {
return Err(Error::BadResponse(
"websocket: the connection is closed".into(),
));
}
self.write_all(frame).map_err(Error::Io)?;
self.flush().map_err(Error::Io)
}
fn write_close(&self, send_closed: &AtomicBool, frame: &[u8]) -> Result<bool> {
let _gate = self.close_gate.lock().unwrap_or_else(|e| e.into_inner());
if send_closed.swap(true, Ordering::SeqCst) {
return Ok(false);
}
self.write_all(frame).map_err(Error::Io)?;
self.flush().map_err(Error::Io)?;
Ok(true)
}
fn idle(&self) -> Duration {
self.last_write.lock().unwrap().elapsed()
}
fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
match &self.kind {
WsTransportKind::Shared(s) => s.lock().unwrap().read(buf),
_ => self.read_polled(buf),
}
}
fn read_polled(&self, buf: &mut [u8]) -> io::Result<usize> {
let deadline =
(*self.read_timeout.lock().unwrap()).and_then(|d| Instant::now().checked_add(d));
let mut programmed: Option<Duration> = None;
loop {
if self.shutdown.load(Ordering::SeqCst) {
return Err(io::Error::new(
io::ErrorKind::ConnectionAborted,
"websocket shut down by WsShutdown handle",
));
}
let tick = match deadline {
Some(dl) => dl
.checked_duration_since(Instant::now())
.unwrap_or(Duration::from_millis(1))
.min(SHUTDOWN_POLL)
.max(Duration::from_millis(1)),
None => SHUTDOWN_POLL,
};
if programmed != Some(tick) {
self.kind.set_read_timeout(Some(tick))?;
programmed = Some(tick);
}
match self.kind.read(buf) {
Ok(n) => return Ok(n),
Err(e) if is_read_timeout(&e) => match deadline {
Some(dl) if Instant::now() >= dl => return Err(e),
_ => continue,
},
Err(e) => return Err(e),
}
}
}
fn write_all(&self, data: &[u8]) -> io::Result<()> {
self.check_not_broken()?;
if let Err(e) = self.kind.write_all(data) {
self.broken.store(true, Ordering::SeqCst);
return Err(e);
}
*self.last_write.lock().unwrap() = Instant::now();
Ok(())
}
fn flush(&self) -> io::Result<()> {
self.check_not_broken()?;
self.kind.flush().inspect_err(|_| {
self.broken.store(true, Ordering::SeqCst);
})
}
fn check_not_broken(&self) -> io::Result<()> {
if self.broken.load(Ordering::SeqCst) {
return Err(io::Error::new(
io::ErrorKind::BrokenPipe,
"websocket: connection unusable after an earlier failed write",
));
}
Ok(())
}
fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
*self.read_timeout.lock().unwrap() = dur;
Ok(())
}
}
impl WsTransportKind {
fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
match self {
WsTransportKind::Plain { read, .. } => read.lock().unwrap().read(buf),
WsTransportKind::Tls(c) => c.read(buf),
WsTransportKind::Shared(s) => s.lock().unwrap().read(buf),
}
}
fn write_all(&self, data: &[u8]) -> io::Result<()> {
match self {
WsTransportKind::Plain { write, .. } => write.lock().unwrap().write_all(data),
WsTransportKind::Tls(c) => c.write(data),
WsTransportKind::Shared(s) => s.lock().unwrap().write_all(data),
}
}
fn flush(&self) -> io::Result<()> {
match self {
WsTransportKind::Plain { write, .. } => write.lock().unwrap().flush(),
WsTransportKind::Tls(c) => c.flush(),
WsTransportKind::Shared(s) => s.lock().unwrap().flush(),
}
}
fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
match self {
WsTransportKind::Plain { read, .. } => read.lock().unwrap().set_read_timeout(dur),
WsTransportKind::Tls(c) => c.set_read_timeout(dur),
WsTransportKind::Shared(_) => Ok(()), }
}
}
struct TransportIo<'a>(&'a WsTransport);
impl Read for TransportIo<'_> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.0.read(buf)
}
}
impl Write for TransportIo<'_> {
fn write(&mut self, data: &[u8]) -> io::Result<usize> {
self.0.write_all(data)?;
Ok(data.len())
}
fn flush(&mut self) -> io::Result<()> {
self.0.flush()
}
}
pub struct WebSocket {
reader: WsReader,
writer: WsWriter,
subprotocol: Option<String>,
}
pub struct WsReader {
transport: Arc<WsTransport>,
rxbuf: Vec<u8>,
partial: Option<PartialMessage>,
pmd: Option<Pmd>,
compression: bool,
auto_pong: bool,
send_closed: Arc<AtomicBool>,
recv_closed: Arc<AtomicBool>,
shutdown: Option<ShutdownSock>,
}
struct PartialMessage {
opcode: u8,
compressed: bool,
buf: Vec<u8>,
}
pub struct WsWriter {
transport: Arc<WsTransport>,
compress: bool,
send_closed: Arc<AtomicBool>,
recv_closed: Arc<AtomicBool>,
shutdown: Option<ShutdownSock>,
}
type ShutdownSock = Arc<Mutex<Box<dyn crate::net::NetStream>>>;
#[derive(Clone)]
pub struct WsShutdown {
sock: ShutdownSock,
flag: Arc<AtomicBool>,
}
impl WsShutdown {
pub fn shutdown(&self) -> Result<()> {
self.flag.store(true, Ordering::SeqCst);
match self.sock.lock().unwrap().shutdown(std::net::Shutdown::Both) {
Ok(()) => Ok(()),
Err(e) if e.kind() == std::io::ErrorKind::NotConnected => Ok(()),
Err(e) => Err(Error::Io(e)),
}
}
}
pub struct WsKeepalive {
stop: Arc<AtomicBool>,
handle: Option<thread::JoinHandle<()>>,
}
impl WsKeepalive {
pub fn stop(mut self) {
self.stop_inner();
}
fn stop_inner(&mut self) {
self.stop.store(true, Ordering::SeqCst);
if let Some(h) = self.handle.take() {
let _ = h.join();
}
}
}
impl Drop for WsKeepalive {
fn drop(&mut self) {
self.stop_inner();
}
}
fn spawn_keepalive(
transport: Arc<WsTransport>,
send_closed: Arc<AtomicBool>,
recv_closed: Arc<AtomicBool>,
interval: Duration,
) -> WsKeepalive {
let stop = Arc::new(AtomicBool::new(false));
let stop_thread = Arc::clone(&stop);
let handle = thread::spawn(move || {
let done = |t: &WsTransport| {
stop_thread.load(Ordering::SeqCst)
|| send_closed.load(Ordering::SeqCst)
|| recv_closed.load(Ordering::SeqCst)
|| t.shutdown.load(Ordering::SeqCst)
};
loop {
if done(&transport) {
return;
}
thread::sleep(SHUTDOWN_POLL);
if done(&transport) {
return;
}
if transport.idle() >= interval {
let sent = build_client_frame(OPCODE_PING, &[])
.and_then(|frame| transport.write_unless_closed(&send_closed, &frame));
if sent.is_err() {
return;
}
}
}
});
WsKeepalive {
stop,
handle: Some(handle),
}
}
impl WebSocket {
pub fn connect(url: &str) -> Result<WebSocket> {
crate::net::Client::new().websocket(url)
}
pub fn connect_with_subprotocols(url: &str, subprotocols: &[&str]) -> Result<WebSocket> {
crate::net::Client::new().websocket_with_subprotocols(url, subprotocols)
}
pub(crate) fn open(
url: &Url,
cfg: &crate::net::NetConfig,
read_timeout: Option<Duration>,
subprotocols: &[String],
) -> Result<WebSocket> {
let (transport, shutdown_sock): (Arc<WsTransport>, Option<Box<dyn crate::net::NetStream>>) =
match url.scheme.as_str() {
"ws" => {
let data = cfg.connect(&url.host, url.port)?;
data.set_write_timeout(Some(SEND_TIMEOUT))
.map_err(Error::Io)?;
data.set_read_timeout(Some(SEND_TIMEOUT))
.map_err(Error::Io)?;
match data.try_clone_box() {
Ok(read) => {
read.set_read_timeout(Some(SEND_TIMEOUT))
.map_err(Error::Io)?;
let shutdown = data.try_clone_box().ok();
(
Arc::new(WsTransport::new(WsTransportKind::Plain {
read: Mutex::new(read),
write: Mutex::new(data),
})),
shutdown,
)
}
Err(_) => (
Arc::new(WsTransport::new(WsTransportKind::Shared(Mutex::new(
Box::new(data),
)))),
None,
),
}
}
"wss" => {
let data = cfg.connect(&url.host, url.port)?;
data.set_write_timeout(Some(SEND_TIMEOUT))
.map_err(Error::Io)?;
data.set_read_timeout(Some(SEND_TIMEOUT))
.map_err(Error::Io)?;
let read_clone = data.try_clone_box().ok();
let shutdown = data.try_clone_box().ok();
let tls = cfg.tls_connect(data, &url.host)?;
match read_clone {
Some(read) => {
read.set_read_timeout(Some(SEND_TIMEOUT))
.map_err(Error::Io)?;
(
Arc::new(WsTransport::new(WsTransportKind::Tls(Box::new(
tls.into_concurrent(read),
)))),
shutdown,
)
}
None => (
Arc::new(WsTransport::new(WsTransportKind::Shared(Mutex::new(
Box::new(tls),
)))),
None,
),
}
}
other => return Err(Error::UnsupportedScheme(other.to_string())),
};
let (pmd, subprotocol) = handshake(&mut TransportIo(&transport), url, subprotocols)?;
let _ = transport.set_read_timeout(read_timeout);
let compression = pmd.is_some();
let send_closed = Arc::new(AtomicBool::new(false));
let recv_closed = Arc::new(AtomicBool::new(false));
let shutdown: Option<ShutdownSock> = shutdown_sock.map(|s| Arc::new(Mutex::new(s)));
Ok(WebSocket {
reader: WsReader {
transport: Arc::clone(&transport),
rxbuf: Vec::new(),
partial: None,
pmd,
compression,
auto_pong: true,
send_closed: Arc::clone(&send_closed),
recv_closed: Arc::clone(&recv_closed),
shutdown: shutdown.clone(),
},
writer: WsWriter {
transport,
compress: compression,
send_closed,
recv_closed,
shutdown,
},
subprotocol,
})
}
pub fn split(self) -> (WsReader, WsWriter) {
(self.reader, self.writer)
}
pub fn shutdown_handle(&self) -> Option<WsShutdown> {
self.reader.shutdown_handle()
}
pub fn subprotocol(&self) -> Option<&str> {
self.subprotocol.as_deref()
}
pub fn compression_enabled(&self) -> bool {
self.reader.compression
}
pub fn is_closed(&self) -> bool {
self.writer.is_closed()
}
pub fn set_read_timeout(&self, dur: Option<Duration>) -> Result<()> {
self.reader.set_read_timeout(dur)
}
pub fn send_text(&mut self, text: &str) -> Result<()> {
self.writer.send_text(text)
}
pub fn send_binary(&mut self, data: &[u8]) -> Result<()> {
self.writer.send_binary(data)
}
pub fn send(&mut self, msg: &WsMessage) -> Result<()> {
self.writer.send(msg)
}
pub fn send_stream<R: Read>(&mut self, opcode: WsOpcode, reader: &mut R) -> Result<()> {
self.writer.send_stream(opcode, reader)
}
pub fn keepalive(&self, interval: Duration) -> WsKeepalive {
self.writer.keepalive(interval)
}
pub fn ping(&mut self, payload: &[u8]) -> Result<()> {
self.writer.ping(payload)
}
pub fn recv(&mut self) -> Result<Option<WsMessage>> {
self.reader.recv()
}
pub fn set_auto_pong(&mut self, on: bool) {
self.reader.set_auto_pong(on);
}
pub fn recv_event(&mut self) -> Result<WsEvent> {
self.reader.recv_event()
}
pub fn send_pong(&mut self, payload: &[u8]) -> Result<()> {
self.writer.send_pong(payload)
}
pub fn send_frame(&mut self, fin: bool, opcode: WsOpcode, payload: &[u8]) -> Result<()> {
self.writer.send_frame(fin, opcode, payload)
}
pub fn recv_frame(&mut self) -> Result<WsFrame> {
self.reader.recv_frame()
}
pub fn close(&mut self) -> Result<()> {
self.writer.close()
}
pub fn close_with(&mut self, code: u16, reason: &str) -> Result<()> {
self.writer.close_with(code, reason)
}
}
impl WsReader {
pub fn set_auto_pong(&mut self, on: bool) {
self.auto_pong = on;
}
pub fn set_read_timeout(&self, dur: Option<Duration>) -> Result<()> {
self.transport.set_read_timeout(dur).map_err(Error::Io)
}
pub fn shutdown_handle(&self) -> Option<WsShutdown> {
self.shutdown.as_ref().map(|s| WsShutdown {
sock: Arc::clone(s),
flag: Arc::clone(&self.transport.shutdown),
})
}
pub fn is_closed(&self) -> bool {
self.send_closed.load(Ordering::SeqCst) || self.recv_closed.load(Ordering::SeqCst)
}
pub fn recv(&mut self) -> Result<Option<WsMessage>> {
loop {
match self.next_event(true)? {
WsEvent::Text(s) => return Ok(Some(WsMessage::Text(s))),
WsEvent::Binary(b) => return Ok(Some(WsMessage::Binary(b))),
WsEvent::Ping(_) | WsEvent::Pong(_) => continue,
WsEvent::Close(_) => return Ok(None),
}
}
}
pub fn recv_event(&mut self) -> Result<WsEvent> {
self.next_event(self.auto_pong)
}
fn next_event(&mut self, auto_pong: bool) -> Result<WsEvent> {
if self.recv_closed.load(Ordering::SeqCst) {
return Ok(WsEvent::Close(None));
}
loop {
let frame = self.next_frame()?;
if frame.opcode >= 0x8 {
validate_control_frame(&frame)?;
let mid_message = self.partial.is_some();
match frame.opcode {
OPCODE_PING => {
if auto_pong {
let pong = build_client_frame(OPCODE_PONG, &frame.payload)?;
self.transport.write_all(&pong).map_err(Error::Io)?;
self.transport.flush().map_err(Error::Io)?;
}
if mid_message {
continue;
}
return Ok(WsEvent::Ping(frame.payload));
}
OPCODE_PONG => {
if mid_message {
continue;
}
return Ok(WsEvent::Pong(frame.payload));
}
OPCODE_CLOSE => return self.handle_close(&frame.payload),
other => {
return Err(Error::BadResponse(format!(
"unknown WS control opcode 0x{other:x}"
)));
}
}
}
match frame.opcode {
OPCODE_TEXT | OPCODE_BINARY => {
if self.partial.is_some() {
return Err(Error::BadResponse(
"new data frame began while a fragmented message was in progress"
.into(),
));
}
if frame.rsv1 && self.pmd.is_none() {
return Err(Error::BadResponse(
"RSV1 set on a WS frame but permessage-deflate was not negotiated"
.into(),
));
}
if frame.fin {
return self.finish_event(frame.opcode, frame.payload, frame.rsv1);
}
self.partial = Some(PartialMessage {
opcode: frame.opcode,
compressed: frame.rsv1,
buf: frame.payload,
});
}
OPCODE_CONT => {
let Some(partial) = self.partial.as_mut() else {
return Err(Error::BadResponse(
"continuation frame with no message in progress".into(),
));
};
if frame.rsv1 {
return Err(Error::BadResponse(
"RSV1 set on a WS continuation frame".into(),
));
}
accumulate(&mut partial.buf, &frame.payload)?;
if frame.fin {
let done = self.partial.take().expect("checked above");
return self.finish_event(done.opcode, done.buf, done.compressed);
}
}
other => {
return Err(Error::BadResponse(format!("unknown WS opcode 0x{other:x}")));
}
}
}
}
fn next_frame(&mut self) -> Result<Frame> {
match read_frame(&mut TransportIo(self.transport.as_ref()), &mut self.rxbuf) {
Err(Error::UnexpectedEof) => {
self.recv_closed.store(true, Ordering::SeqCst);
Err(Error::UnexpectedEof)
}
other => other,
}
}
fn handle_close(&mut self, payload: &[u8]) -> Result<WsEvent> {
self.recv_closed.store(true, Ordering::SeqCst);
match parse_close_payload(payload) {
Ok(close) => {
if let Ok(frame) =
build_client_frame(OPCODE_CLOSE, &close_echo_payload(close.as_ref()))
{
let _ = self.transport.write_close(&self.send_closed, &frame);
}
Ok(WsEvent::Close(close))
}
Err(status) => {
self.fail(status);
Err(bad_close_error(status))
}
}
}
fn fail(&self, status: u16) {
if let Ok(frame) =
close_payload(status, "").and_then(|p| build_client_frame(OPCODE_CLOSE, &p))
{
let _ = self.transport.write_close(&self.send_closed, &frame);
}
}
fn finish_event(&mut self, opcode: u8, payload: Vec<u8>, compressed: bool) -> Result<WsEvent> {
let payload = if compressed {
let pmd = self.pmd.as_mut().ok_or_else(|| {
Error::BadResponse(
"compressed WS message without negotiated permessage-deflate".into(),
)
})?;
pmd.inflate_message(&payload)?
} else {
payload
};
if opcode != OPCODE_TEXT {
return Ok(WsEvent::Binary(payload));
}
match String::from_utf8(payload) {
Ok(s) => Ok(WsEvent::Text(s)),
Err(_) => {
self.fail(CLOSE_INVALID_DATA);
Err(Error::BadResponse(
"WS TEXT message payload is not valid UTF-8".into(),
))
}
}
}
pub fn recv_frame(&mut self) -> Result<WsFrame> {
let frame = self.next_frame()?;
if frame.opcode >= 0x8 {
validate_control_frame(&frame)?;
} else if frame.rsv1 && (self.pmd.is_none() || frame.opcode == OPCODE_CONT) {
return Err(Error::BadResponse(
"RSV1 set on a WS frame where permessage-deflate does not allow it".into(),
));
}
Ok(WsFrame {
fin: frame.fin,
opcode: WsOpcode::from_u8(frame.opcode)?,
payload: frame.payload,
})
}
}
impl WsWriter {
pub fn is_closed(&self) -> bool {
self.send_closed.load(Ordering::SeqCst) || self.recv_closed.load(Ordering::SeqCst)
}
pub fn shutdown_handle(&self) -> Option<WsShutdown> {
self.shutdown.as_ref().map(|s| WsShutdown {
sock: Arc::clone(s),
flag: Arc::clone(&self.transport.shutdown),
})
}
fn write_frame(&self, frame: &[u8]) -> Result<()> {
self.transport.write_unless_closed(&self.send_closed, frame)
}
pub fn send_text(&mut self, text: &str) -> Result<()> {
let frame = build_data_frame(OPCODE_TEXT, text.as_bytes(), self.compress)?;
self.write_frame(&frame)
}
pub fn send_binary(&mut self, data: &[u8]) -> Result<()> {
let frame = build_data_frame(OPCODE_BINARY, data, self.compress)?;
self.write_frame(&frame)
}
pub fn send_stream<R: Read>(&mut self, opcode: WsOpcode, reader: &mut R) -> Result<()> {
send_stream_inner(opcode.to_u8(), reader, |frame| self.write_frame(frame))
}
pub fn keepalive(&self, interval: Duration) -> WsKeepalive {
spawn_keepalive(
Arc::clone(&self.transport),
Arc::clone(&self.send_closed),
Arc::clone(&self.recv_closed),
interval,
)
}
pub fn send(&mut self, msg: &WsMessage) -> Result<()> {
match msg {
WsMessage::Text(s) => self.send_text(s),
WsMessage::Binary(b) => self.send_binary(b),
}
}
pub fn ping(&mut self, payload: &[u8]) -> Result<()> {
self.send_control(OPCODE_PING, payload)
}
pub fn send_pong(&mut self, payload: &[u8]) -> Result<()> {
self.send_control(OPCODE_PONG, payload)
}
pub fn send_frame(&mut self, fin: bool, opcode: WsOpcode, payload: &[u8]) -> Result<()> {
if opcode.is_control() {
if !fin {
return Err(Error::BadResponse(
"websocket: control frames cannot be fragmented (fin must be true)".into(),
));
}
check_control_payload(payload)?;
}
let frame = build_client_frame_inner(fin, opcode.to_u8(), payload, false)?;
if opcode == WsOpcode::Close {
return self.send_close_frame(&frame);
}
self.write_frame(&frame)
}
fn send_control(&mut self, opcode: u8, payload: &[u8]) -> Result<()> {
check_control_payload(payload)?;
let frame = build_client_frame(opcode, payload)?;
self.write_frame(&frame)
}
fn send_close_frame(&mut self, frame: &[u8]) -> Result<()> {
self.transport.write_close(&self.send_closed, frame)?;
Ok(())
}
pub fn close(&mut self) -> Result<()> {
let frame = build_client_frame(OPCODE_CLOSE, &[])?;
self.send_close_frame(&frame)
}
pub fn close_with(&mut self, code: u16, reason: &str) -> Result<()> {
let payload = close_payload(code, reason)?;
let frame = build_client_frame(OPCODE_CLOSE, &payload)?;
self.send_close_frame(&frame)
}
}
fn check_control_payload(payload: &[u8]) -> Result<()> {
if payload.len() > MAX_CONTROL_PAYLOAD {
return Err(Error::BadResponse(format!(
"websocket: control frame payload too large: {} bytes (max {MAX_CONTROL_PAYLOAD})",
payload.len()
)));
}
Ok(())
}
fn handshake<S: Read + Write, P: AsRef<str>>(
stream: &mut S,
url: &Url,
subprotocols: &[P],
) -> Result<(Option<Pmd>, Option<String>)> {
let key_b64 = base64_encode(&random_16()?);
let proto_header = subprotocol_header(subprotocols)?;
let req = upgrade_request(url, &key_b64, Some(PMD_OFFER), &proto_header);
stream.write_all(req.as_bytes())?;
stream.flush()?;
let buf = read_handshake_head(stream, HANDSHAKE_DEADLINE)?;
let resp = check_handshake_response(&buf, &key_b64, subprotocols)?;
let pmd = match resp.extensions.as_deref() {
Some(v) => parse_pmd_response(v)?,
None => None,
};
Ok((pmd, resp.subprotocol))
}
pub(crate) fn upgrade_request(
url: &Url,
key_b64: &str,
extensions: Option<&str>,
proto_header: &str,
) -> String {
use std::fmt::Write as _;
let path = if url.path.is_empty() {
"/"
} else {
url.path.as_str()
};
let default_port =
(url.scheme == "ws" && url.port == 80) || (url.scheme == "wss" && url.port == 443);
let mut req = String::with_capacity(256 + path.len() + url.host.len() + proto_header.len());
let _ = write!(req, "GET {path} HTTP/1.1\r\nHost: {}", url.host);
if !default_port {
let _ = write!(req, ":{}", url.port);
}
let _ = write!(
req,
"\r\n\
Upgrade: websocket\r\n\
Connection: Upgrade\r\n\
Sec-WebSocket-Key: {key_b64}\r\n\
Sec-WebSocket-Version: 13\r\n"
);
if let Some(ext) = extensions {
let _ = write!(req, "Sec-WebSocket-Extensions: {ext}\r\n");
}
req.push_str(proto_header);
req.push_str("\r\n");
req
}
pub(crate) fn subprotocol_header<P: AsRef<str>>(subprotocols: &[P]) -> Result<String> {
if subprotocols.is_empty() {
return Ok(String::new());
}
let mut header = String::from("Sec-WebSocket-Protocol: ");
for (i, p) in subprotocols.iter().enumerate() {
let p = p.as_ref();
if !is_http_token(p) {
return Err(Error::InvalidUrl(format!(
"websocket: invalid subprotocol {p:?} (must be an HTTP token)"
)));
}
if i > 0 {
header.push_str(", ");
}
header.push_str(p);
}
header.push_str("\r\n");
Ok(header)
}
fn is_http_token(s: &str) -> bool {
!s.is_empty()
&& s.bytes().all(|b| {
b.is_ascii_alphanumeric()
|| matches!(
b,
b'!' | b'#'
| b'$'
| b'%'
| b'&'
| b'\''
| b'*'
| b'+'
| b'-'
| b'.'
| b'^'
| b'_'
| b'`'
| b'|'
| b'~'
)
})
}
pub(crate) struct HandshakeResponse {
pub(crate) extensions: Option<String>,
pub(crate) subprotocol: Option<String>,
}
pub(crate) fn check_handshake_response<P: AsRef<str>>(
buf: &[u8],
key_b64: &str,
offered: &[P],
) -> Result<HandshakeResponse> {
let head = std::str::from_utf8(buf)
.map_err(|_| Error::BadResponse("non-utf8 handshake response".into()))?;
let mut lines = head.split("\r\n");
let status_line = lines
.next()
.ok_or_else(|| Error::BadResponse("empty handshake response".into()))?;
let mut parts = status_line.split(' ');
if parts.next() != Some("HTTP/1.1") || parts.next() != Some("101") {
return Err(Error::BadResponse(format!(
"expected 101 Switching Protocols, got: {status_line:?}"
)));
}
let mut upgrade_ok = false;
let mut connection_ok = false;
let mut accept_value: Option<&str> = None;
let mut extensions_value: Option<String> = None;
let mut subprotocol_value: Option<String> = None;
for line in lines {
if line.is_empty() {
break;
}
let (k, v) = match line.split_once(':') {
Some((k, v)) => (k.trim(), v.trim()),
None => continue,
};
if k.eq_ignore_ascii_case("upgrade") {
if v.eq_ignore_ascii_case("websocket") {
upgrade_ok = true;
}
} else if k.eq_ignore_ascii_case("connection") {
if v.split(',')
.any(|t| t.trim().eq_ignore_ascii_case("upgrade"))
{
connection_ok = true;
}
} else if k.eq_ignore_ascii_case("sec-websocket-accept") {
accept_value = Some(v);
} else if k.eq_ignore_ascii_case("sec-websocket-protocol") {
if subprotocol_value.replace(v.to_string()).is_some() {
return Err(Error::BadResponse(
"websocket: server sent Sec-WebSocket-Protocol more than once".into(),
));
}
} else if k.eq_ignore_ascii_case("sec-websocket-extensions") {
match &mut extensions_value {
Some(existing) => {
existing.push_str(", ");
existing.push_str(v);
}
None => extensions_value = Some(v.to_string()),
}
}
}
if !upgrade_ok {
return Err(Error::BadResponse(
"missing or wrong Upgrade header in handshake response".into(),
));
}
if !connection_ok {
return Err(Error::BadResponse(
"missing or wrong Connection header in handshake response".into(),
));
}
let accept = accept_value
.ok_or_else(|| Error::BadResponse("missing Sec-WebSocket-Accept header".into()))?;
let expected = derive_accept(key_b64);
if accept != expected {
return Err(Error::BadResponse(format!(
"Sec-WebSocket-Accept mismatch: got {accept:?}, expected {expected:?}"
)));
}
if let Some(sel) = &subprotocol_value {
if !offered.iter().any(|p| p.as_ref() == sel) {
return Err(Error::BadResponse(format!(
"websocket: server selected subprotocol {sel:?}, which was not offered"
)));
}
}
Ok(HandshakeResponse {
extensions: extensions_value.filter(|v| !v.trim().is_empty()),
subprotocol: subprotocol_value,
})
}
fn read_handshake_head<S: Read>(stream: &mut S, deadline: Duration) -> Result<Vec<u8>> {
let mut buf: Vec<u8> = Vec::with_capacity(512);
let start = Instant::now();
loop {
let mut b = [0u8; 1];
let n = stream.read(&mut b)?;
if n == 0 {
return Err(Error::UnexpectedEof);
}
buf.push(b[0]);
if buf.ends_with(b"\r\n\r\n") {
break;
}
if buf.len() > 64 * 1024 {
return Err(Error::BadResponse("handshake response too large".into()));
}
if start.elapsed() > deadline {
return Err(Error::BadResponse(
"handshake response timed out (header read exceeded deadline)".into(),
));
}
}
Ok(buf)
}
fn parse_pmd_response(value: &str) -> Result<Option<Pmd>> {
let bad = |why: &str| {
Error::BadResponse(format!(
"websocket: invalid Sec-WebSocket-Extensions response ({why}): {value:?}"
))
};
let mut pmd: Option<Pmd> = None;
for ext in value.split(',').map(str::trim) {
if ext.is_empty() {
continue;
}
let mut params = ext.split(';').map(str::trim);
let name = params.next().unwrap_or("");
if !name.eq_ignore_ascii_case("permessage-deflate") {
return Err(bad("extension was not offered"));
}
if pmd.is_some() {
return Err(bad("permessage-deflate selected more than once"));
}
let mut client_no_context_takeover = false;
let mut server_no_context_takeover = false;
let mut server_max_window_bits = false;
let mut client_max_window_bits = false;
for param in params {
let (key, val) = match param.split_once('=') {
Some((k, v)) => {
let v = v.trim();
let v = v
.strip_prefix('"')
.and_then(|v| v.strip_suffix('"'))
.unwrap_or(v);
(k.trim(), Some(v))
}
None => (param, None),
};
let flag = |seen: &mut bool| -> Result<()> {
if std::mem::replace(seen, true) {
return Err(bad("duplicate parameter"));
}
Ok(())
};
if key.eq_ignore_ascii_case("client_no_context_takeover") {
flag(&mut client_no_context_takeover)?;
if val.is_some() {
return Err(bad("client_no_context_takeover takes no value"));
}
} else if key.eq_ignore_ascii_case("server_no_context_takeover") {
flag(&mut server_no_context_takeover)?;
if val.is_some() {
return Err(bad("server_no_context_takeover takes no value"));
}
} else if key.eq_ignore_ascii_case("server_max_window_bits") {
flag(&mut server_max_window_bits)?;
if val.and_then(parse_window_bits).is_none() {
return Err(bad("server_max_window_bits must be 8..=15"));
}
} else if key.eq_ignore_ascii_case("client_max_window_bits") {
flag(&mut client_max_window_bits)?;
if val.and_then(parse_window_bits) != Some(15) {
return Err(bad("client_max_window_bits was not offered"));
}
} else {
return Err(bad("unknown parameter"));
}
}
pmd = Some(Pmd {
client_no_context_takeover,
server_no_context_takeover,
decoder: Deflate::decoder(),
});
}
Ok(pmd)
}
fn parse_window_bits(v: &str) -> Option<u8> {
if v.is_empty() || v.len() > 2 || v.starts_with('0') || !v.bytes().all(|b| b.is_ascii_digit()) {
return None;
}
let n: u8 = v.parse().ok()?;
(8..=15).contains(&n).then_some(n)
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum Message {
Data { opcode: u8, payload: Vec<u8> },
Closed,
}
const MAX_CONTROL_PAYLOAD: usize = 125;
fn read_message<S: Read + Write>(stream: &mut S, mut pmd: Option<&mut Pmd>) -> Result<Message> {
let mut frag_opcode: Option<u8> = None;
let mut compressed = false;
let mut buf: Vec<u8> = Vec::new();
let mut rx: Vec<u8> = Vec::new();
loop {
let frame = read_frame(stream, &mut rx)?;
if frame.opcode >= 0x8 {
if frame.rsv1 {
return Err(Error::BadResponse("RSV1 set on a WS control frame".into()));
}
if !frame.fin {
return Err(Error::BadResponse(
"fragmented control frame (FIN=0 on a control opcode)".into(),
));
}
if frame.payload.len() > MAX_CONTROL_PAYLOAD {
return Err(Error::BadResponse(format!(
"control frame payload too large: {} bytes (max {MAX_CONTROL_PAYLOAD})",
frame.payload.len()
)));
}
match frame.opcode {
OPCODE_PING => {
let pong = build_client_frame(OPCODE_PONG, &frame.payload)?;
stream.write_all(&pong)?;
stream.flush()?;
continue;
}
OPCODE_PONG => continue,
OPCODE_CLOSE => {
let (echo, outcome) = match parse_close_payload(&frame.payload) {
Ok(close) => (Ok(close_echo_payload(close.as_ref())), Ok(Message::Closed)),
Err(status) => (close_payload(status, ""), Err(bad_close_error(status))),
};
if let Ok(close) = echo.and_then(|p| build_client_frame(OPCODE_CLOSE, &p)) {
let _ = stream.write_all(&close);
let _ = stream.flush();
}
return outcome;
}
other => {
return Err(Error::BadResponse(format!(
"unknown WS control opcode 0x{other:x}"
)));
}
}
}
match frame.opcode {
OPCODE_TEXT | OPCODE_BINARY => {
if frag_opcode.is_some() {
return Err(Error::BadResponse(
"new data frame began while a fragmented message was in progress".into(),
));
}
if frame.rsv1 {
if pmd.is_none() {
return Err(Error::BadResponse(
"RSV1 set on a WS frame but permessage-deflate was not negotiated"
.into(),
));
}
compressed = true;
}
accumulate(&mut buf, &frame.payload)?;
if frame.fin {
return finish_data_message(frame.opcode, buf, compressed, pmd.as_deref_mut());
}
frag_opcode = Some(frame.opcode);
}
OPCODE_CONT => {
let opcode = frag_opcode.ok_or_else(|| {
Error::BadResponse("continuation frame with no message in progress".into())
})?;
if frame.rsv1 {
return Err(Error::BadResponse(
"RSV1 set on a WS continuation frame".into(),
));
}
accumulate(&mut buf, &frame.payload)?;
if frame.fin {
return finish_data_message(opcode, buf, compressed, pmd.as_deref_mut());
}
}
other => {
return Err(Error::BadResponse(format!("unknown WS opcode 0x{other:x}")));
}
}
}
}
fn finish_data_message(
opcode: u8,
payload: Vec<u8>,
compressed: bool,
pmd: Option<&mut Pmd>,
) -> Result<Message> {
let payload = if compressed {
let pmd = pmd.ok_or_else(|| {
Error::BadResponse("compressed WS message without negotiated permessage-deflate".into())
})?;
pmd.inflate_message(&payload)?
} else {
payload
};
if opcode == OPCODE_TEXT && std::str::from_utf8(&payload).is_err() {
return Err(Error::BadResponse(
"WS TEXT message payload is not valid UTF-8".into(),
));
}
Ok(Message::Data { opcode, payload })
}
pub(crate) fn accumulate(buf: &mut Vec<u8>, chunk: &[u8]) -> Result<()> {
let total = buf.len() as u64 + chunk.len() as u64;
if total > MAX_PAYLOAD_BYTES {
return Err(Error::BadResponse(format!(
"reassembled WS message too large: {total} bytes (max {MAX_PAYLOAD_BYTES})"
)));
}
buf.extend_from_slice(chunk);
Ok(())
}
pub(crate) fn validate_control_frame(frame: &Frame) -> Result<()> {
if frame.rsv1 {
return Err(Error::BadResponse("RSV1 set on a WS control frame".into()));
}
if !frame.fin {
return Err(Error::BadResponse(
"fragmented control frame (FIN=0 on a control opcode)".into(),
));
}
if frame.payload.len() > MAX_CONTROL_PAYLOAD {
return Err(Error::BadResponse(format!(
"control frame payload too large: {} bytes (max {MAX_CONTROL_PAYLOAD})",
frame.payload.len()
)));
}
Ok(())
}
pub(crate) fn is_valid_close_code(code: u16) -> bool {
matches!(code, 1000..=1003 | 1007..=1014 | 3000..=4999)
}
pub(crate) const CLOSE_PROTOCOL_ERROR: u16 = 1002;
pub(crate) const CLOSE_INVALID_DATA: u16 = 1007;
pub(crate) fn parse_close_payload(payload: &[u8]) -> std::result::Result<Option<WsClose>, u16> {
match payload.len() {
0 => return Ok(None),
1 => return Err(CLOSE_PROTOCOL_ERROR),
_ => {}
}
let code = u16::from_be_bytes([payload[0], payload[1]]);
if !is_valid_close_code(code) {
return Err(CLOSE_PROTOCOL_ERROR);
}
let reason = std::str::from_utf8(&payload[2..])
.map_err(|_| CLOSE_INVALID_DATA)?
.to_string();
Ok(Some(WsClose { code, reason }))
}
pub(crate) fn bad_close_error(status: u16) -> Error {
Error::BadResponse(format!(
"websocket: malformed close frame from peer (failed the connection with {status})"
))
}
pub(crate) fn close_echo_payload(close: Option<&WsClose>) -> Vec<u8> {
match close {
Some(c) => c.code.to_be_bytes().to_vec(),
None => Vec::new(),
}
}
fn build_data_frame(opcode: u8, payload: &[u8], compress: bool) -> Result<Vec<u8>> {
if opcode != OPCODE_TEXT && opcode != OPCODE_BINARY {
return Err(Error::BadResponse(format!(
"send_message expects a data opcode (text/binary), got 0x{opcode:x}"
)));
}
if compress {
let compressed = deflate_message(payload)?;
build_client_frame_rsv1(opcode, &compressed)
} else {
build_client_frame(opcode, payload)
}
}
#[cfg(test)]
fn send_message<S: Write>(
stream: &mut S,
opcode: u8,
payload: &[u8],
compress: bool,
) -> Result<()> {
let frame = build_data_frame(opcode, payload, compress)?;
stream.write_all(&frame)?;
stream.flush()?;
Ok(())
}
const STREAM_CHUNK: usize = 64 * 1024;
fn read_chunk<R: Read>(reader: &mut R, buf: &mut Vec<u8>, cap: usize) -> io::Result<()> {
buf.clear();
buf.resize(cap, 0);
let mut filled = 0;
while filled < cap {
match reader.read(&mut buf[filled..]) {
Ok(0) => break,
Ok(n) => filled += n,
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
Err(e) => return Err(e),
}
}
buf.truncate(filled);
Ok(())
}
fn send_stream_inner<R: Read>(
opcode: u8,
reader: &mut R,
mut emit: impl FnMut(&[u8]) -> Result<()>,
) -> Result<()> {
if opcode != OPCODE_TEXT && opcode != OPCODE_BINARY {
return Err(Error::BadResponse(format!(
"send_stream expects a data opcode (text/binary), got 0x{opcode:x}"
)));
}
let mut cur = Vec::new();
let mut next = Vec::new();
read_chunk(reader, &mut cur, STREAM_CHUNK).map_err(Error::Io)?;
let mut first = true;
loop {
read_chunk(reader, &mut next, STREAM_CHUNK).map_err(Error::Io)?;
let fin = next.is_empty();
let op = if first { opcode } else { OPCODE_CONT };
let frame = build_client_frame_inner(fin, op, &cur, false)?;
emit(&frame)?;
first = false;
if fin {
break;
}
std::mem::swap(&mut cur, &mut next);
}
Ok(())
}
fn read_data_and_close<S: Read + Write>(stream: &mut S, pmd: Option<&mut Pmd>) -> Result<Vec<u8>> {
let payload = match read_message(stream, pmd)? {
Message::Data { payload, .. } => payload,
Message::Closed => return Ok(Vec::new()),
};
if let Ok(close) = build_client_frame(OPCODE_CLOSE, &[]) {
let _ = stream.write_all(&close);
let _ = stream.flush();
}
Ok(payload)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct Frame {
pub(crate) fin: bool,
pub(crate) rsv1: bool,
pub(crate) opcode: u8,
pub(crate) payload: Vec<u8>,
}
pub(crate) fn try_parse_frame(buf: &[u8]) -> Result<Option<(Frame, usize)>> {
if buf.len() < 2 {
return Ok(None);
}
let h0 = buf[0];
let h1 = buf[1];
let fin = (h0 & 0x80) != 0;
let rsv1 = (h0 & 0x40) != 0;
if (h0 & 0x30) != 0 {
return Err(Error::BadResponse(
"non-zero RSV2/RSV3 bits on incoming WS frame".into(),
));
}
let opcode = h0 & 0x0F;
if (h1 & 0x80) != 0 {
return Err(Error::BadResponse(
"server-to-client frame is masked".into(),
));
}
let len7 = h1 & 0x7F;
let mut pos = 2usize;
let payload_len: u64 = match len7 {
0..=125 => len7 as u64,
126 => {
if buf.len() < pos + 2 {
return Ok(None);
}
let v = u16::from_be_bytes([buf[pos], buf[pos + 1]]) as u64;
if v < 126 {
return Err(Error::BadResponse(
"non-minimal 16-bit WS payload length".into(),
));
}
pos += 2;
v
}
127 => {
if buf.len() < pos + 8 {
return Ok(None);
}
let v = u64::from_be_bytes(buf[pos..pos + 8].try_into().unwrap());
if v >> 63 != 0 {
return Err(Error::BadResponse(
"WS payload length has its most significant bit set".into(),
));
}
if v <= u16::MAX as u64 {
return Err(Error::BadResponse(
"non-minimal 64-bit WS payload length".into(),
));
}
pos += 8;
v
}
_ => unreachable!(),
};
if opcode >= 0x8 && payload_len > MAX_CONTROL_PAYLOAD as u64 {
return Err(Error::BadResponse(format!(
"control frame payload too large: {payload_len} bytes (max {MAX_CONTROL_PAYLOAD})"
)));
}
if payload_len > MAX_PAYLOAD_BYTES {
return Err(Error::BadResponse(format!(
"WS payload too large: {payload_len} bytes"
)));
}
let plen = payload_len as usize;
if buf.len() < pos + plen {
return Ok(None);
}
let payload = buf[pos..pos + plen].to_vec();
pos += plen;
Ok(Some((
Frame {
fin,
rsv1,
opcode,
payload,
},
pos,
)))
}
fn read_frame<S: Read>(stream: &mut S, rx: &mut Vec<u8>) -> Result<Frame> {
let mut tmp = [0u8; 16 * 1024];
loop {
if let Some((frame, consumed)) = try_parse_frame(rx)? {
rx.drain(..consumed);
return Ok(frame);
}
let n = stream.read(&mut tmp)?;
if n == 0 {
return Err(Error::UnexpectedEof);
}
rx.extend_from_slice(&tmp[..n]);
}
}
pub(crate) fn close_payload(code: u16, reason: &str) -> Result<Vec<u8>> {
if !is_valid_close_code(code) {
return Err(Error::BadResponse(format!(
"websocket: close code {code} may not be sent (use 1000-1003, 1007-1014 or 3000-4999)"
)));
}
let mut payload = Vec::with_capacity(2 + reason.len());
payload.extend_from_slice(&code.to_be_bytes());
payload.extend_from_slice(reason.as_bytes());
if payload.len() > MAX_CONTROL_PAYLOAD {
return Err(Error::BadResponse(format!(
"websocket: close reason too long: {} bytes (max {})",
payload.len(),
MAX_CONTROL_PAYLOAD - 2
)));
}
Ok(payload)
}
pub(crate) fn build_client_frame(opcode: u8, payload: &[u8]) -> Result<Vec<u8>> {
build_client_frame_inner(true, opcode, payload, false)
}
fn build_client_frame_rsv1(opcode: u8, payload: &[u8]) -> Result<Vec<u8>> {
build_client_frame_inner(true, opcode, payload, true)
}
fn build_client_frame_inner(fin: bool, opcode: u8, payload: &[u8], rsv1: bool) -> Result<Vec<u8>> {
let mask: [u8; 4] = {
let r = random_16()?;
[r[0], r[1], r[2], r[3]]
};
let mut out = Vec::with_capacity(2 + 8 + 4 + payload.len());
let fin_bit = if fin { 0x80 } else { 0x00 };
let rsv1_bit = if rsv1 { 0x40 } else { 0x00 };
out.push(fin_bit | rsv1_bit | (opcode & 0x0F)); let n = payload.len();
if n < 126 {
out.push(0x80 | (n as u8));
} else if n <= u16::MAX as usize {
out.push(0x80 | 126);
out.extend_from_slice(&(n as u16).to_be_bytes());
} else {
out.push(0x80 | 127);
out.extend_from_slice(&(n as u64).to_be_bytes());
}
out.extend_from_slice(&mask);
let start = out.len();
out.extend_from_slice(payload);
for (i, b) in out[start..].iter_mut().enumerate() {
*b ^= mask[i & 3];
}
Ok(out)
}
pub(crate) fn derive_accept(key_b64: &str) -> String {
let mut h = Sha1::new();
h.update(key_b64.as_bytes());
h.update(WS_GUID.as_bytes());
let digest = h.finalize();
base64_encode(digest.as_ref())
}
pub(crate) fn random_16() -> Result<[u8; 16]> {
use purecrypto::rng::{OsRng, RngCore};
let mut out = [0u8; 16];
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
OsRng.fill_bytes(&mut out);
}))
.map_err(|_| Error::BadResponse("websocket: no secure entropy source available".into()))?;
Ok(out)
}
pub(crate) fn base64_encode(input: &[u8]) -> String {
const ALPHA: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
let mut out = String::with_capacity(input.len().div_ceil(3) * 4);
let mut i = 0;
while i + 3 <= input.len() {
let b0 = input[i];
let b1 = input[i + 1];
let b2 = input[i + 2];
out.push(ALPHA[(b0 >> 2) as usize] as char);
out.push(ALPHA[(((b0 & 0x03) << 4) | (b1 >> 4)) as usize] as char);
out.push(ALPHA[(((b1 & 0x0F) << 2) | (b2 >> 6)) as usize] as char);
out.push(ALPHA[(b2 & 0x3F) as usize] as char);
i += 3;
}
let rem = input.len() - i;
if rem == 1 {
let b0 = input[i];
out.push(ALPHA[(b0 >> 2) as usize] as char);
out.push(ALPHA[((b0 & 0x03) << 4) as usize] as char);
out.push('=');
out.push('=');
} else if rem == 2 {
let b0 = input[i];
let b1 = input[i + 1];
out.push(ALPHA[(b0 >> 2) as usize] as char);
out.push(ALPHA[(((b0 & 0x03) << 4) | (b1 >> 4)) as usize] as char);
out.push(ALPHA[((b1 & 0x0F) << 2) as usize] as char);
out.push('=');
}
out
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
struct MockStream {
inbound: Cursor<Vec<u8>>,
sent: Vec<u8>,
}
impl MockStream {
fn new(inbound: Vec<u8>) -> Self {
MockStream {
inbound: Cursor::new(inbound),
sent: Vec::new(),
}
}
}
impl Read for MockStream {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
self.inbound.read(buf)
}
}
impl Write for MockStream {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.sent.extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
fn server_frame(fin: bool, opcode: u8, payload: &[u8]) -> Vec<u8> {
let mut out = Vec::new();
let b0 = if fin { 0x80 } else { 0x00 } | (opcode & 0x0F);
out.push(b0);
let n = payload.len();
if n < 126 {
out.push(n as u8);
} else if n <= u16::MAX as usize {
out.push(126);
out.extend_from_slice(&(n as u16).to_be_bytes());
} else {
out.push(127);
out.extend_from_slice(&(n as u64).to_be_bytes());
}
out.extend_from_slice(payload);
out
}
fn server_frame_rsv1(fin: bool, opcode: u8, payload: &[u8]) -> Vec<u8> {
let mut out = server_frame(fin, opcode, payload);
out[0] |= 0x40; out
}
fn pmd_compress(data: &[u8]) -> Vec<u8> {
let mut out = compress_to_vec::<Deflate>(data).expect("deflate encode");
if out.ends_with(&DEFLATE_TAIL) {
out.truncate(out.len() - DEFLATE_TAIL.len());
}
out
}
fn test_pmd() -> Pmd {
Pmd {
client_no_context_takeover: true,
server_no_context_takeover: true,
decoder: Deflate::decoder(),
}
}
fn decode_sent(sent: &[u8]) -> Vec<(u8, Vec<u8>)> {
let mut frames = Vec::new();
let mut i = 0;
while i < sent.len() {
let opcode = sent[i] & 0x0F;
let masked = (sent[i + 1] & 0x80) != 0;
assert!(masked, "client frame must be masked");
let len7 = sent[i + 1] & 0x7F;
i += 2;
let len = match len7 {
0..=125 => len7 as usize,
126 => {
let l = u16::from_be_bytes([sent[i], sent[i + 1]]) as usize;
i += 2;
l
}
127 => {
let mut b = [0u8; 8];
b.copy_from_slice(&sent[i..i + 8]);
i += 8;
u64::from_be_bytes(b) as usize
}
_ => unreachable!(),
};
let mask = [sent[i], sent[i + 1], sent[i + 2], sent[i + 3]];
i += 4;
let mut payload = sent[i..i + len].to_vec();
i += len;
for (j, b) in payload.iter_mut().enumerate() {
*b ^= mask[j & 3];
}
frames.push((opcode, payload));
}
frames
}
struct DripStream {
per_read: Duration,
}
impl Read for DripStream {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
std::thread::sleep(self.per_read);
if buf.is_empty() {
return Ok(0);
}
buf[0] = b'X'; Ok(1)
}
}
#[test]
fn handshake_head_reads_up_to_terminator_without_overreading() {
let head = b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\n\r\n";
let frame = [0x81u8, 0x02, b'h', b'i'];
let mut inbound = head.to_vec();
inbound.extend_from_slice(&frame);
let mut s = MockStream::new(inbound);
let got = read_handshake_head(&mut s, Duration::from_secs(60)).expect("reads header");
assert_eq!(
&got, head,
"must capture exactly the header, no frame bytes"
);
let mut rest = Vec::new();
s.read_to_end(&mut rest).expect("drain remainder");
assert_eq!(rest, frame, "first-frame bytes must not be consumed");
}
#[test]
fn handshake_head_deadline_trips_on_slow_drip() {
let mut s = DripStream {
per_read: Duration::from_millis(5),
};
let err = read_handshake_head(&mut s, Duration::from_millis(20))
.expect_err("slow drip must hit the deadline");
match err {
Error::BadResponse(m) => assert!(m.contains("timed out"), "unexpected message: {m}"),
other => panic!("wrong error: {other:?}"),
}
}
#[test]
fn base64_encode_hello() {
assert_eq!(base64_encode(b"hello"), "aGVsbG8=");
}
#[test]
fn base64_encode_rfc4648_vectors() {
assert_eq!(base64_encode(b""), "");
assert_eq!(base64_encode(b"f"), "Zg==");
assert_eq!(base64_encode(b"fo"), "Zm8=");
assert_eq!(base64_encode(b"foo"), "Zm9v");
assert_eq!(base64_encode(b"foob"), "Zm9vYg==");
assert_eq!(base64_encode(b"fooba"), "Zm9vYmE=");
assert_eq!(base64_encode(b"foobar"), "Zm9vYmFy");
}
#[test]
fn rfc6455_accept_derivation() {
let key = "dGhlIHNhbXBsZSBub25jZQ==";
let expected = "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=";
assert_eq!(derive_accept(key), expected);
}
#[test]
fn parse_short_text_frame() {
let bytes = [0x81, 0x05, b'h', b'e', b'l', b'l', b'o'];
let mut cur = Cursor::new(&bytes[..]);
let f = read_frame(&mut cur, &mut Vec::new()).expect("frame parses");
assert!(f.fin);
assert_eq!(f.opcode, OPCODE_TEXT);
assert_eq!(f.payload, b"hello");
}
#[test]
fn parse_16bit_length_frame() {
let mut bytes: Vec<u8> = vec![0x82, 126, 0x00, 200];
bytes.extend(std::iter::repeat_n(b'A', 200));
let mut cur = Cursor::new(bytes);
let f = read_frame(&mut cur, &mut Vec::new()).expect("frame parses");
assert_eq!(f.opcode, OPCODE_BINARY);
assert_eq!(f.payload.len(), 200);
assert!(f.payload.iter().all(|&b| b == b'A'));
}
#[test]
fn read_frame_rejects_oversized_control_frame_before_alloc() {
let bytes = [0x89u8, 126, 0x00, 126];
let mut cur = Cursor::new(&bytes[..]);
let err = read_frame(&mut cur, &mut Vec::new()).expect_err("must reject");
match err {
Error::BadResponse(msg) => assert!(
msg.contains("control frame payload too large"),
"msg = {msg}"
),
other => panic!("expected BadResponse, got {other:?}"),
}
}
#[test]
fn read_frame_resumes_after_midframe_error() {
struct Flaky {
data: Vec<u8>,
pos: usize,
fail_at: usize,
failed: bool,
}
impl Read for Flaky {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
if !self.failed && self.pos >= self.fail_at {
self.failed = true;
return Err(std::io::Error::new(
std::io::ErrorKind::WouldBlock,
"timeout",
));
}
if self.pos >= self.data.len() {
return Ok(0);
}
buf[0] = self.data[self.pos];
self.pos += 1;
Ok(1)
}
}
let mut r = Flaky {
data: vec![0x81, 0x05, b'h', b'e', b'l', b'l', b'o'],
pos: 0,
fail_at: 3, failed: false,
};
let mut rx = Vec::new();
assert!(read_frame(&mut r, &mut rx).is_err());
let f = read_frame(&mut r, &mut rx).expect("resumes after the error");
assert_eq!(f.opcode, OPCODE_TEXT);
assert_eq!(f.payload, b"hello");
assert!(rx.is_empty(), "buffer fully consumed");
}
#[test]
fn reject_masked_server_frame() {
let bytes = [0x81, 0x80, 0, 0, 0, 0];
let mut cur = Cursor::new(&bytes[..]);
let err = read_frame(&mut cur, &mut Vec::new())
.expect_err("masked server frame must be rejected");
match err {
Error::BadResponse(_) => {}
other => panic!("wrong error: {other:?}"),
}
}
#[test]
fn build_close_frame_short_payload() {
let frame = build_client_frame(OPCODE_CLOSE, &[]).unwrap();
assert_eq!(frame.len(), 6);
assert_eq!(frame[0], 0x88);
assert_eq!(frame[1], 0x80); }
#[test]
fn build_text_frame_masks_payload() {
let payload = b"hi";
let frame = build_client_frame(OPCODE_TEXT, payload).unwrap();
assert_eq!(frame.len(), 8);
assert_eq!(frame[0], 0x81);
assert_eq!(frame[1], 0x82);
let mask = [frame[2], frame[3], frame[4], frame[5]];
let unmasked: Vec<u8> = frame[6..]
.iter()
.enumerate()
.map(|(i, &b)| b ^ mask[i & 3])
.collect();
assert_eq!(unmasked, payload);
}
#[test]
fn build_frame_uses_16bit_length_for_medium_payload() {
let payload = vec![0u8; 200];
let frame = build_client_frame(OPCODE_BINARY, &payload).unwrap();
assert_eq!(frame[0], 0x82);
assert_eq!(frame[1], 0x80 | 126);
let len = u16::from_be_bytes([frame[2], frame[3]]);
assert_eq!(len, 200);
assert_eq!(frame.len(), 208);
}
#[test]
fn build_frame_uses_64bit_length_for_large_payload() {
let payload = vec![0u8; 70_000];
let frame = build_client_frame(OPCODE_BINARY, &payload).unwrap();
assert_eq!(frame[1], 0x80 | 127);
let len = u64::from_be_bytes([
frame[2], frame[3], frame[4], frame[5], frame[6], frame[7], frame[8], frame[9],
]);
assert_eq!(len, 70_000);
}
#[test]
fn random_16_is_nonzero() {
let r = random_16().expect("OS entropy available in the test environment");
assert_ne!(r, [0u8; 16]);
}
#[test]
fn random_16_is_not_constant() {
let a = random_16().unwrap();
let b = random_16().unwrap();
assert_ne!(a, b);
}
#[test]
fn reassembles_fragmented_text_message() {
let mut inbound = server_frame(false, OPCODE_TEXT, b"Hel");
inbound.extend(server_frame(false, OPCODE_CONT, b"lo "));
inbound.extend(server_frame(true, OPCODE_CONT, b"world"));
let mut s = MockStream::new(inbound);
let msg = read_message(&mut s, None).expect("reassembles");
assert_eq!(
msg,
Message::Data {
opcode: OPCODE_TEXT,
payload: b"Hello world".to_vec(),
}
);
}
#[test]
fn invalid_utf8_text_message_is_rejected() {
let inbound = server_frame(true, OPCODE_TEXT, &[0xff, 0xfe]);
let mut s = MockStream::new(inbound);
let err = read_message(&mut s, None).expect_err("invalid utf-8 TEXT must be rejected");
match err {
Error::BadResponse(m) => assert!(m.contains("UTF-8"), "unexpected message: {m}"),
other => panic!("unexpected error: {other:?}"),
}
}
#[test]
fn invalid_utf8_binary_message_is_accepted() {
let inbound = server_frame(true, OPCODE_BINARY, &[0xff, 0xfe]);
let mut s = MockStream::new(inbound);
let msg = read_message(&mut s, None).expect("binary is not utf-8 validated");
assert_eq!(
msg,
Message::Data {
opcode: OPCODE_BINARY,
payload: vec![0xff, 0xfe],
}
);
}
#[test]
fn valid_utf8_text_message_passes() {
let mut inbound = server_frame(false, OPCODE_TEXT, &[0xc3]);
inbound.extend(server_frame(true, OPCODE_CONT, &[0xa9]));
let mut s = MockStream::new(inbound);
let msg = read_message(&mut s, None).expect("valid utf-8 across fragments");
assert_eq!(
msg,
Message::Data {
opcode: OPCODE_TEXT,
payload: "é".as_bytes().to_vec(),
}
);
}
#[test]
fn ping_between_fragments_gets_pong_and_message_completes() {
let mut inbound = server_frame(false, OPCODE_TEXT, b"foo");
inbound.extend(server_frame(true, OPCODE_PING, b"pingdata"));
inbound.extend(server_frame(true, OPCODE_CONT, b"bar"));
let mut s = MockStream::new(inbound);
let msg = read_message(&mut s, None).expect("completes despite ping");
assert_eq!(
msg,
Message::Data {
opcode: OPCODE_TEXT,
payload: b"foobar".to_vec(),
}
);
let sent = decode_sent(&s.sent);
assert_eq!(sent.len(), 1, "exactly one pong expected");
assert_eq!(sent[0].0, OPCODE_PONG);
assert_eq!(sent[0].1, b"pingdata");
}
#[test]
fn close_is_answered_and_returns_closed() {
let inbound = server_frame(true, OPCODE_CLOSE, &[]);
let mut s = MockStream::new(inbound);
let msg = read_message(&mut s, None).expect("handles close");
assert_eq!(msg, Message::Closed);
let sent = decode_sent(&s.sent);
assert_eq!(sent.len(), 1, "exactly one close reply expected");
assert_eq!(sent[0].0, OPCODE_CLOSE);
}
#[test]
fn unsolicited_pong_is_ignored_then_data_returns() {
let mut inbound = server_frame(true, OPCODE_PONG, b"x");
inbound.extend(server_frame(true, OPCODE_TEXT, b"hi"));
let mut s = MockStream::new(inbound);
let msg = read_message(&mut s, None).expect("ignores pong");
assert_eq!(
msg,
Message::Data {
opcode: OPCODE_TEXT,
payload: b"hi".to_vec(),
}
);
assert!(s.sent.is_empty(), "unsolicited pong must not be answered");
}
#[test]
fn send_message_produces_masked_frame() {
let mut s = MockStream::new(Vec::new());
send_message(&mut s, OPCODE_TEXT, b"hello", false).expect("sends");
assert_eq!(s.sent[0], 0x81);
assert_eq!(s.sent[1], 0x80 | 5);
assert_eq!(s.sent.len(), 2 + 4 + 5);
let decoded = decode_sent(&s.sent);
assert_eq!(decoded, vec![(OPCODE_TEXT, b"hello".to_vec())]);
}
#[test]
fn send_message_rejects_control_opcode() {
let mut s = MockStream::new(Vec::new());
let err = send_message(&mut s, OPCODE_PING, b"x", false)
.expect_err("control opcode must be rejected for send_message");
match err {
Error::BadResponse(_) => {}
other => panic!("wrong error: {other:?}"),
}
assert!(s.sent.is_empty());
}
#[test]
fn oversized_cumulative_fragmented_payload_is_rejected() {
let mut buf = vec![0u8; (MAX_PAYLOAD_BYTES - 1) as usize];
accumulate(&mut buf, &[0u8]).expect("exactly at the cap is allowed");
let err = accumulate(&mut buf, &[0u8]).expect_err("over the cap must be rejected");
match err {
Error::BadResponse(_) => {}
other => panic!("wrong error: {other:?}"),
}
}
#[test]
fn fragmented_control_frame_is_rejected() {
let inbound = server_frame(false, OPCODE_PING, b"x");
let mut s = MockStream::new(inbound);
let err = read_message(&mut s, None).expect_err("fragmented control must be rejected");
match err {
Error::BadResponse(_) => {}
other => panic!("wrong error: {other:?}"),
}
}
#[test]
fn oversized_control_frame_is_rejected() {
let inbound = server_frame(true, OPCODE_PING, &[0u8; 126]);
let mut s = MockStream::new(inbound);
let err = read_message(&mut s, None).expect_err("oversized control must be rejected");
match err {
Error::BadResponse(_) => {}
other => panic!("wrong error: {other:?}"),
}
}
#[test]
fn new_data_frame_during_fragmentation_is_rejected() {
let mut inbound = server_frame(false, OPCODE_TEXT, b"a");
inbound.extend(server_frame(true, OPCODE_TEXT, b"b"));
let mut s = MockStream::new(inbound);
let err =
read_message(&mut s, None).expect_err("interleaved new data frame must be rejected");
match err {
Error::BadResponse(_) => {}
other => panic!("wrong error: {other:?}"),
}
}
#[test]
fn lone_continuation_frame_is_rejected() {
let inbound = server_frame(true, OPCODE_CONT, b"x");
let mut s = MockStream::new(inbound);
let err = read_message(&mut s, None).expect_err("lone continuation must be rejected");
match err {
Error::BadResponse(_) => {}
other => panic!("wrong error: {other:?}"),
}
}
#[test]
fn read_data_and_close_returns_reassembled_payload_and_sends_close() {
let mut inbound = server_frame(false, OPCODE_BINARY, &[1, 2, 3]);
inbound.extend(server_frame(true, OPCODE_CONT, &[4, 5]));
let mut s = MockStream::new(inbound);
let payload = read_data_and_close(&mut s, None).expect("reads message");
assert_eq!(payload, vec![1, 2, 3, 4, 5]);
let sent = decode_sent(&s.sent);
assert_eq!(sent.last().map(|f| f.0), Some(OPCODE_CLOSE));
}
fn decode_first_frame_with_rsv1(sent: &[u8]) -> (u8, bool, Vec<u8>) {
let opcode = sent[0] & 0x0F;
let rsv1 = (sent[0] & 0x40) != 0;
let masked = (sent[1] & 0x80) != 0;
assert!(masked, "client frame must be masked");
let len7 = sent[1] & 0x7F;
let mut i = 2;
let len = match len7 {
0..=125 => len7 as usize,
126 => {
let l = u16::from_be_bytes([sent[i], sent[i + 1]]) as usize;
i += 2;
l
}
127 => {
let mut b = [0u8; 8];
b.copy_from_slice(&sent[i..i + 8]);
i += 8;
u64::from_be_bytes(b) as usize
}
_ => unreachable!(),
};
let mask = [sent[i], sent[i + 1], sent[i + 2], sent[i + 3]];
i += 4;
let mut payload = sent[i..i + len].to_vec();
for (j, b) in payload.iter_mut().enumerate() {
*b ^= mask[j & 3];
}
(opcode, rsv1, payload)
}
fn pmd_inflate(compressed: &[u8]) -> Vec<u8> {
let mut input = compressed.to_vec();
input.extend_from_slice(&DEFLATE_TAIL);
let mut dec = Deflate::decoder();
let mut out = Vec::new();
let mut scratch = vec![0u8; 32 * 1024];
let mut consumed = 0usize;
loop {
let before_c = consumed;
let before_w = out.len();
let (p, status) = dec
.decode(&input[consumed..], &mut scratch)
.expect("inflate");
out.extend_from_slice(&scratch[..p.written]);
consumed += p.consumed;
match status {
Status::StreamEnd => break,
Status::OutputFull => continue,
Status::InputEmpty => {
if consumed >= input.len() || (consumed == before_c && out.len() == before_w) {
break;
}
}
}
}
out
}
#[test]
fn handshake_offers_permessage_deflate() {
struct Recorder {
request: Vec<u8>,
response: Cursor<Vec<u8>>,
}
impl Read for Recorder {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
self.response.read(buf)
}
}
impl Write for Recorder {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.request.extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
let resp = b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: wrong\r\n\r\n".to_vec();
let mut rec = Recorder {
request: Vec::new(),
response: Cursor::new(resp),
};
let url = Url::parse("ws://example.com/chat").expect("url");
let _ = handshake(&mut rec, &url, NO_SUBPROTOCOLS); let req = String::from_utf8(rec.request).expect("utf8 request");
assert!(
req.contains("Sec-WebSocket-Extensions: permessage-deflate"),
"request must offer permessage-deflate, got:\n{req}"
);
assert!(req.contains("client_no_context_takeover"));
assert!(req.contains("server_no_context_takeover"));
}
#[test]
fn parse_pmd_response_enables_compression() {
let pmd = parse_pmd_response("permessage-deflate; server_no_context_takeover")
.unwrap()
.expect("permessage-deflate accepted");
assert!(pmd.server_no_context_takeover);
assert!(!pmd.client_no_context_takeover);
let pmd2 = parse_pmd_response(
"permessage-deflate; client_no_context_takeover; server_no_context_takeover",
)
.unwrap()
.expect("accepted with both flags");
assert!(pmd2.client_no_context_takeover);
assert!(pmd2.server_no_context_takeover);
for ok in [
"permessage-deflate; server_max_window_bits=10",
"permessage-deflate; server_max_window_bits=\"8\"",
"permessage-deflate; client_max_window_bits=15",
] {
assert!(parse_pmd_response(ok).unwrap().is_some(), "{ok}");
}
}
#[test]
fn parse_pmd_response_without_extension_is_none() {
assert!(parse_pmd_response("").unwrap().is_none());
assert!(parse_pmd_response(" , ").unwrap().is_none());
}
#[test]
fn parse_pmd_response_rejects_what_was_not_offered() {
for bad in [
"some-other-extension",
"foo; bar=1",
"permessage-deflate, x-webkit-deflate-frame",
"permessage-deflate, permessage-deflate",
"permessage-deflate; foo",
"permessage-deflate; server_no_context_takeover; server_no_context_takeover",
"permessage-deflate; server_no_context_takeover=1",
"permessage-deflate; server_max_window_bits=7",
"permessage-deflate; server_max_window_bits=16",
"permessage-deflate; server_max_window_bits=010",
"permessage-deflate; server_max_window_bits",
"permessage-deflate; client_max_window_bits=10",
"permessage-deflate; client_max_window_bits",
] {
assert!(parse_pmd_response(bad).is_err(), "should reject {bad:?}");
}
}
#[test]
fn handshake_no_longer_offers_client_max_window_bits() {
assert!(!PMD_OFFER.contains("client_max_window_bits"));
}
#[test]
fn upgrade_request_bytes() {
let url = Url::parse("wss://example.com").expect("url");
assert_eq!(
upgrade_request(&url, "KEY", None, ""),
"GET / HTTP/1.1\r\nHost: example.com\r\nUpgrade: websocket\r\n\
Connection: Upgrade\r\nSec-WebSocket-Key: KEY\r\n\
Sec-WebSocket-Version: 13\r\n\r\n"
);
let url = Url::parse("ws://example.com:8080/chat?x=1").expect("url");
assert_eq!(
upgrade_request(&url, "KEY", Some("ext"), "Sec-WebSocket-Protocol: a, b\r\n"),
"GET /chat?x=1 HTTP/1.1\r\nHost: example.com:8080\r\nUpgrade: websocket\r\n\
Connection: Upgrade\r\nSec-WebSocket-Key: KEY\r\n\
Sec-WebSocket-Version: 13\r\nSec-WebSocket-Extensions: ext\r\n\
Sec-WebSocket-Protocol: a, b\r\n\r\n"
);
}
#[test]
fn inflate_compressed_message_decodes_to_original() {
let original = b"the quick brown fox jumps over the lazy dog, the quick brown fox";
let compressed = pmd_compress(original);
assert!(
compressed.len() < original.len(),
"fixture should actually compress"
);
let inbound = server_frame_rsv1(true, OPCODE_TEXT, &compressed);
let mut s = MockStream::new(inbound);
let mut pmd = test_pmd();
let msg = read_message(&mut s, Some(&mut pmd)).expect("decodes compressed message");
assert_eq!(
msg,
Message::Data {
opcode: OPCODE_TEXT,
payload: original.to_vec(),
}
);
}
#[test]
fn send_message_compressed_round_trips() {
let input = b"hello hello hello permessage-deflate round trip";
let mut s = MockStream::new(Vec::new());
send_message(&mut s, OPCODE_TEXT, input, true).expect("sends compressed");
let (opcode, rsv1, payload) = decode_first_frame_with_rsv1(&s.sent);
assert_eq!(opcode, OPCODE_TEXT);
assert!(rsv1, "compressed frame must have RSV1 set");
assert_ne!(payload, input, "payload should be compressed, not raw");
assert_eq!(pmd_inflate(&payload), input);
}
#[test]
fn rsv1_without_negotiation_is_rejected() {
let inbound = server_frame_rsv1(true, OPCODE_TEXT, b"whatever");
let mut s = MockStream::new(inbound);
let err = read_message(&mut s, None).expect_err("RSV1 without PMD must be rejected");
match err {
Error::BadResponse(_) => {}
other => panic!("wrong error: {other:?}"),
}
}
#[test]
fn rsv2_or_rsv3_always_rejected() {
let mut inbound = server_frame(true, OPCODE_TEXT, b"x");
inbound[0] |= 0x20;
let mut s = MockStream::new(inbound);
let mut pmd = test_pmd();
let err = read_message(&mut s, Some(&mut pmd)).expect_err("RSV2 must be rejected");
assert!(matches!(err, Error::BadResponse(_)));
let mut inbound = server_frame(true, OPCODE_TEXT, b"x");
inbound[0] |= 0x10;
let mut s = MockStream::new(inbound);
let mut pmd = test_pmd();
let err = read_message(&mut s, Some(&mut pmd)).expect_err("RSV3 must be rejected");
assert!(matches!(err, Error::BadResponse(_)));
}
#[test]
fn rsv1_on_control_frame_is_rejected() {
let inbound = server_frame_rsv1(true, OPCODE_PING, b"x");
let mut s = MockStream::new(inbound);
let mut pmd = test_pmd();
let err =
read_message(&mut s, Some(&mut pmd)).expect_err("RSV1 on control must be rejected");
assert!(matches!(err, Error::BadResponse(_)));
}
#[test]
fn compressed_bomb_exceeding_cap_is_rejected() {
let huge = vec![0u8; (MAX_PAYLOAD_BYTES + (1 << 20)) as usize];
let compressed = pmd_compress(&huge);
assert!(
(compressed.len() as u64) < MAX_PAYLOAD_BYTES,
"fixture must be much smaller than the cap"
);
let inbound = server_frame_rsv1(true, OPCODE_BINARY, &compressed);
let mut s = MockStream::new(inbound);
let mut pmd = test_pmd();
let err =
read_message(&mut s, Some(&mut pmd)).expect_err("compression bomb must be rejected");
match err {
Error::BadResponse(msg) => {
assert!(msg.contains("permessage-deflate"), "got {msg:?}")
}
other => panic!("wrong error: {other:?}"),
}
}
#[test]
fn fragmented_compressed_message_reassembles_and_inflates() {
let original = b"fragmented compressed payload, split across two frames on the wire";
let compressed = pmd_compress(original);
assert!(compressed.len() >= 4, "need enough bytes to split");
let mid = compressed.len() / 2;
let mut inbound = server_frame_rsv1(false, OPCODE_TEXT, &compressed[..mid]);
inbound.extend(server_frame(true, OPCODE_CONT, &compressed[mid..]));
let mut s = MockStream::new(inbound);
let mut pmd = test_pmd();
let msg = read_message(&mut s, Some(&mut pmd)).expect("reassembles + inflates");
assert_eq!(
msg,
Message::Data {
opcode: OPCODE_TEXT,
payload: original.to_vec(),
}
);
}
#[test]
fn rsv1_on_continuation_frame_is_rejected() {
let original = b"continuation rsv1 should be rejected here";
let compressed = pmd_compress(original);
let mid = compressed.len() / 2;
let mut inbound = server_frame_rsv1(false, OPCODE_TEXT, &compressed[..mid]);
inbound.extend(server_frame_rsv1(true, OPCODE_CONT, &compressed[mid..]));
let mut s = MockStream::new(inbound);
let mut pmd = test_pmd();
let err = read_message(&mut s, Some(&mut pmd))
.expect_err("RSV1 on continuation must be rejected");
assert!(matches!(err, Error::BadResponse(_)));
}
#[test]
fn uncompressed_message_passes_through_when_pmd_negotiated() {
let inbound = server_frame(true, OPCODE_TEXT, b"plain text, no rsv1");
let mut s = MockStream::new(inbound);
let mut pmd = test_pmd();
let msg = read_message(&mut s, Some(&mut pmd)).expect("plain message");
assert_eq!(
msg,
Message::Data {
opcode: OPCODE_TEXT,
payload: b"plain text, no rsv1".to_vec(),
}
);
}
use std::sync::atomic::AtomicBool;
use std::sync::{Arc, Mutex};
#[derive(Clone)]
struct SharedMock {
inbound: Arc<Mutex<Cursor<Vec<u8>>>>,
sent: Arc<Mutex<Vec<u8>>>,
}
impl SharedMock {
fn new(inbound: Vec<u8>) -> Self {
SharedMock {
inbound: Arc::new(Mutex::new(Cursor::new(inbound))),
sent: Arc::new(Mutex::new(Vec::new())),
}
}
fn sent(&self) -> Vec<u8> {
self.sent.lock().unwrap().clone()
}
}
impl Read for SharedMock {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
self.inbound.lock().unwrap().read(buf)
}
}
impl Write for SharedMock {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.sent.lock().unwrap().extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
fn ws_over(mock: SharedMock, pmd: Option<Pmd>) -> WebSocket {
let compression = pmd.is_some();
let transport = Arc::new(WsTransport::new(WsTransportKind::Shared(Mutex::new(
Box::new(mock),
))));
let send_closed = Arc::new(AtomicBool::new(false));
let recv_closed = Arc::new(AtomicBool::new(false));
WebSocket {
reader: WsReader {
transport: Arc::clone(&transport),
rxbuf: Vec::new(),
partial: None,
pmd,
compression,
auto_pong: true,
send_closed: Arc::clone(&send_closed),
recv_closed: Arc::clone(&recv_closed),
shutdown: None,
},
writer: WsWriter {
transport,
compress: compression,
send_closed,
recv_closed,
shutdown: None,
},
subprotocol: None,
}
}
#[derive(Clone, Default)]
struct ScriptedMock {
reads: Arc<Mutex<std::collections::VecDeque<std::io::Result<Vec<u8>>>>>,
sent: Arc<Mutex<Vec<u8>>>,
fail_writes: Arc<AtomicBool>,
}
impl ScriptedMock {
fn new(reads: Vec<std::io::Result<Vec<u8>>>) -> Self {
ScriptedMock {
reads: Arc::new(Mutex::new(reads.into())),
..Default::default()
}
}
fn sent(&self) -> Vec<u8> {
self.sent.lock().unwrap().clone()
}
}
impl Read for ScriptedMock {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
match self.reads.lock().unwrap().pop_front() {
None => Ok(0),
Some(Err(e)) => Err(e),
Some(Ok(chunk)) => {
buf[..chunk.len()].copy_from_slice(&chunk);
Ok(chunk.len())
}
}
}
}
impl Write for ScriptedMock {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
if self.fail_writes.load(Ordering::SeqCst) {
let half = buf.len() / 2;
self.sent.lock().unwrap().extend_from_slice(&buf[..half]);
return Err(std::io::Error::new(std::io::ErrorKind::TimedOut, "stall"));
}
self.sent.lock().unwrap().extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
fn ws_over_scripted(mock: ScriptedMock) -> WebSocket {
let transport = Arc::new(WsTransport::new(WsTransportKind::Shared(Mutex::new(
Box::new(mock),
))));
let send_closed = Arc::new(AtomicBool::new(false));
let recv_closed = Arc::new(AtomicBool::new(false));
WebSocket {
reader: WsReader {
transport: Arc::clone(&transport),
rxbuf: Vec::new(),
partial: None,
pmd: None,
compression: false,
auto_pong: true,
send_closed: Arc::clone(&send_closed),
recv_closed: Arc::clone(&recv_closed),
shutdown: None,
},
writer: WsWriter {
transport,
compress: false,
send_closed,
recv_closed,
shutdown: None,
},
subprotocol: None,
}
}
fn timed_out() -> std::io::Error {
std::io::Error::new(std::io::ErrorKind::WouldBlock, "read timed out")
}
#[test]
fn read_timeout_between_fragments_is_resumable() {
let mock = ScriptedMock::new(vec![
Ok(server_frame(false, OPCODE_TEXT, b"hel")),
Err(timed_out()),
Ok(server_frame(true, OPCODE_CONT, b"lo")),
]);
let mut ws = ws_over_scripted(mock);
assert!(ws.recv().is_err(), "the timeout surfaces");
assert!(!ws.is_closed());
assert_eq!(ws.recv().unwrap(), Some(WsMessage::Text("hello".into())));
}
#[test]
fn peer_close_is_echoed_with_its_code() {
let mut payload = 1001u16.to_be_bytes().to_vec();
payload.extend_from_slice(b"bye");
let mock = SharedMock::new(server_frame(true, OPCODE_CLOSE, &payload));
let mut ws = ws_over(mock.clone(), None);
assert_eq!(ws.recv().unwrap(), None);
assert_eq!(
decode_sent(&mock.sent()),
vec![(OPCODE_CLOSE, 1001u16.to_be_bytes().to_vec())]
);
}
#[test]
fn malformed_close_fails_the_connection() {
for (payload, status) in [
(vec![0x03], 1002u16), (999u16.to_be_bytes().to_vec(), 1002), (1004u16.to_be_bytes().to_vec(), 1002), (1006u16.to_be_bytes().to_vec(), 1002), (1015u16.to_be_bytes().to_vec(), 1002), (2000u16.to_be_bytes().to_vec(), 1002), (vec![0x03, 0xE8, 0xC3, 0x28], 1007), ] {
let mock = SharedMock::new(server_frame(true, OPCODE_CLOSE, &payload));
let mut ws = ws_over(mock.clone(), None);
assert!(ws.recv_event().is_err(), "{payload:?}");
assert!(ws.is_closed());
assert_eq!(
decode_sent(&mock.sent()),
vec![(OPCODE_CLOSE, status.to_be_bytes().to_vec())],
"{payload:?}"
);
}
for code in [1000u16, 1003, 1007, 1011, 1014, 3000, 4999] {
let mock = SharedMock::new(server_frame(true, OPCODE_CLOSE, &code.to_be_bytes()));
let mut ws = ws_over(mock, None);
assert!(ws.recv_event().is_ok(), "{code}");
}
}
#[test]
fn close_with_refuses_codes_that_may_not_be_sent() {
let mock = SharedMock::new(Vec::new());
let mut ws = ws_over(mock.clone(), None);
for code in [0u16, 1004, 1005, 1006, 1015, 2999, 5000] {
assert!(ws.close_with(code, "").is_err(), "{code}");
}
assert!(
!ws.is_closed(),
"a refused close must not mark the socket closed"
);
assert!(mock.sent().is_empty());
}
#[test]
fn invalid_utf8_text_fails_with_1007() {
let mock = SharedMock::new(server_frame(true, OPCODE_TEXT, &[0xC3]));
let mut ws = ws_over(mock.clone(), None);
assert!(ws.recv().is_err());
assert_eq!(
decode_sent(&mock.sent()),
vec![(OPCODE_CLOSE, 1007u16.to_be_bytes().to_vec())]
);
}
#[test]
fn eof_marks_the_socket_closed() {
let mock = SharedMock::new(Vec::new());
let mut ws = ws_over(mock, None);
assert!(matches!(ws.recv(), Err(Error::UnexpectedEof)));
assert!(ws.is_closed());
}
#[test]
fn recv_answers_pings_even_with_auto_pong_off() {
let mut inbound = server_frame(true, OPCODE_PING, b"hi");
inbound.extend(server_frame(true, OPCODE_TEXT, b"x"));
let mock = SharedMock::new(inbound);
let mut ws = ws_over(mock.clone(), None);
ws.set_auto_pong(false);
assert_eq!(ws.recv().unwrap(), Some(WsMessage::Text("x".into())));
assert_eq!(
decode_sent(&mock.sent()),
vec![(OPCODE_PONG, b"hi".to_vec())]
);
}
#[test]
fn failed_partial_write_poisons_the_transport() {
let mock = ScriptedMock::default();
let mut ws = ws_over_scripted(mock.clone());
mock.fail_writes.store(true, Ordering::SeqCst);
assert!(ws.send_text("first").is_err());
let after_failure = mock.sent().len();
mock.fail_writes.store(false, Ordering::SeqCst);
assert!(ws.send_text("second").is_err());
assert_eq!(mock.sent().len(), after_failure, "nothing more written");
}
#[test]
fn send_stream_stops_once_the_connection_closes() {
let mock = SharedMock::new(Vec::new());
let mut ws = ws_over(mock.clone(), None);
let writer = &mut ws.writer;
struct ClosingReader {
send_closed: Arc<AtomicBool>,
served: usize,
}
impl Read for ClosingReader {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
self.served += 1;
if self.served > 2 {
self.send_closed.store(true, Ordering::SeqCst);
}
if self.served > 6 {
return Ok(0);
}
let n = buf.len().min(STREAM_CHUNK);
buf[..n].fill(b'a');
Ok(n)
}
}
let mut src = ClosingReader {
send_closed: Arc::clone(&writer.send_closed),
served: 0,
};
assert!(writer.send_stream(WsOpcode::Binary, &mut src).is_err());
let frames = decode_sent(&mock.sent());
assert_eq!(frames.len(), 1);
assert_eq!(frames[0].0, OPCODE_BINARY);
}
#[test]
fn raw_recv_frame_rejects_unnegotiated_rsv1() {
let mock = SharedMock::new(server_frame_rsv1(true, OPCODE_TEXT, b"x"));
let mut ws = ws_over(mock, None);
assert!(ws.recv_frame().is_err());
let mut inbound = server_frame(false, OPCODE_TEXT, b"a");
inbound.extend(server_frame_rsv1(true, OPCODE_CONT, b"b"));
let mock = SharedMock::new(inbound);
let mut ws = ws_over(mock, Some(test_pmd()));
assert!(ws.recv_frame().is_ok());
assert!(ws.recv_frame().is_err(), "RSV1 on a continuation frame");
}
#[test]
fn non_minimal_length_encodings_are_rejected() {
let mut f = vec![0x82, 126, 0, 5];
f.extend_from_slice(b"hello");
assert!(try_parse_frame(&f).is_err());
let mut f = vec![0x82, 127];
f.extend_from_slice(&300u64.to_be_bytes());
f.extend(vec![0u8; 300]);
assert!(try_parse_frame(&f).is_err());
let mut f = vec![0x82, 127];
f.extend_from_slice(&(1u64 << 63 | 70_000).to_be_bytes());
assert!(try_parse_frame(&f).is_err());
}
#[test]
fn context_takeover_decoder_survives_bfinal_messages() {
let mut pmd = Pmd {
client_no_context_takeover: true,
server_no_context_takeover: false,
decoder: Deflate::decoder(),
};
assert_eq!(
pmd.inflate_message(&pmd_compress(b"first")).unwrap(),
b"first"
);
assert_eq!(
pmd.inflate_message(&pmd_compress(b"second")).unwrap(),
b"second"
);
}
#[test]
fn data_after_the_final_deflate_block_is_rejected() {
let mut pmd = test_pmd();
let mut payload = pmd_compress(b"hello");
payload.extend_from_slice(b"trailing junk");
assert!(pmd.inflate_message(&payload).is_err());
}
fn check(resp: &str, key: &str, offered: &[&str]) -> Result<HandshakeResponse> {
check_handshake_response(resp.as_bytes(), key, offered)
}
fn resp_101(status: &str, key: &str, extra: &str) -> String {
format!(
"{status}\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\
Sec-WebSocket-Accept: {}\r\n{extra}\r\n",
derive_accept(key)
)
}
#[test]
fn handshake_response_subprotocol_must_have_been_offered() {
let key = "dGhlIHNhbXBsZSBub25jZQ==";
let ok = resp_101(
"HTTP/1.1 101 Switching Protocols",
key,
"Sec-WebSocket-Protocol: chat\r\n",
);
assert_eq!(
check(&ok, key, &["chat"]).unwrap().subprotocol.as_deref(),
Some("chat")
);
assert!(check(&ok, key, &["other"]).is_err(), "not offered");
assert!(check(&ok, key, &[]).is_err(), "none offered");
assert!(check(&ok, key, &["Chat"]).is_err(), "exact match only");
let twice = resp_101(
"HTTP/1.1 101 Switching Protocols",
key,
"Sec-WebSocket-Protocol: chat\r\nSec-WebSocket-Protocol: chat\r\n",
);
assert!(check(&twice, key, &["chat"]).is_err());
let none = resp_101("HTTP/1.1 101 Switching Protocols", key, "");
assert!(check(&none, key, &["chat"]).unwrap().subprotocol.is_none());
}
#[test]
fn handshake_response_requires_http11_101() {
let key = "dGhlIHNhbXBsZSBub25jZQ==";
for status in [
"HTTP/1.0 101 Switching Protocols",
"HTTP/1.1 1010 x",
"HTTP/1.1 200 OK",
] {
assert!(
check(&resp_101(status, key, ""), key, &[]).is_err(),
"{status}"
);
}
assert!(check(&resp_101("HTTP/1.1 101", key, ""), key, &[]).is_ok());
}
#[test]
fn subprotocol_values_must_be_tokens() {
assert_eq!(
subprotocol_header(&["chat", "v2.proto"]).unwrap(),
"Sec-WebSocket-Protocol: chat, v2.proto\r\n"
);
for bad in ["", "a b", "a,b", "chat\r\nX-Evil: 1", "caf\u{e9}"] {
assert!(subprotocol_header(&[bad]).is_err(), "{bad:?}");
}
}
#[test]
fn websocket_recv_maps_text_binary_and_close() {
let mut inbound = Vec::new();
inbound.extend(server_frame(true, OPCODE_TEXT, b"hello"));
inbound.extend(server_frame(true, OPCODE_BINARY, &[1, 2, 3]));
inbound.extend(server_frame(true, OPCODE_CLOSE, &[]));
let mut ws = ws_over(SharedMock::new(inbound), None);
assert_eq!(ws.recv().unwrap(), Some(WsMessage::Text("hello".into())));
assert_eq!(ws.recv().unwrap(), Some(WsMessage::Binary(vec![1, 2, 3])));
assert_eq!(ws.recv().unwrap(), None);
assert!(ws.is_closed());
assert_eq!(ws.recv().unwrap(), None);
}
#[test]
fn websocket_send_text_writes_one_masked_text_frame() {
let mock = SharedMock::new(Vec::new());
let mut ws = ws_over(mock.clone(), None);
ws.send_text("hi there").unwrap();
let frames = decode_sent(&mock.sent());
assert_eq!(frames.len(), 1);
assert_eq!(frames[0].0, OPCODE_TEXT);
assert_eq!(frames[0].1, b"hi there");
}
#[test]
fn websocket_close_sends_close_frame_and_blocks_further_sends() {
let mock = SharedMock::new(Vec::new());
let mut ws = ws_over(mock.clone(), None);
ws.close().unwrap();
let frames = decode_sent(&mock.sent());
assert_eq!(frames.len(), 1);
assert_eq!(frames[0].0, OPCODE_CLOSE);
assert!(ws.is_closed());
assert!(ws.send_text("nope").is_err());
ws.close().unwrap();
}
#[test]
fn websocket_recv_autoresponds_ping_with_pong() {
let mut inbound = Vec::new();
inbound.extend(server_frame(true, OPCODE_PING, b"pingdata"));
inbound.extend(server_frame(true, OPCODE_TEXT, b"after ping"));
let mock = SharedMock::new(inbound);
let mut ws = ws_over(mock.clone(), None);
assert_eq!(
ws.recv().unwrap(),
Some(WsMessage::Text("after ping".into()))
);
let frames = decode_sent(&mock.sent());
assert!(
frames
.iter()
.any(|(op, pl)| *op == OPCODE_PONG && pl.as_slice() == b"pingdata"),
"expected an automatic PONG echoing the ping data, got {frames:?}"
);
}
#[test]
fn websocket_recv_inflates_compressed_message() {
let payload = "compress me ".repeat(8);
let inbound = server_frame_rsv1(true, OPCODE_TEXT, &pmd_compress(payload.as_bytes()));
let mut ws = ws_over(SharedMock::new(inbound), Some(test_pmd()));
assert!(ws.compression_enabled());
assert_eq!(ws.recv().unwrap(), Some(WsMessage::Text(payload)));
}
#[test]
fn websocket_recv_event_surfaces_ping_pong_and_close_code() {
let mut inbound = Vec::new();
inbound.extend(server_frame(true, OPCODE_PING, b"pp"));
inbound.extend(server_frame(true, OPCODE_PONG, b"qq"));
let mut close_payload = 1001u16.to_be_bytes().to_vec();
close_payload.extend_from_slice(b"bye");
inbound.extend(server_frame(true, OPCODE_CLOSE, &close_payload));
let mock = SharedMock::new(inbound);
let mut ws = ws_over(mock.clone(), None);
assert_eq!(ws.recv_event().unwrap(), WsEvent::Ping(b"pp".to_vec()));
assert_eq!(ws.recv_event().unwrap(), WsEvent::Pong(b"qq".to_vec()));
assert_eq!(
ws.recv_event().unwrap(),
WsEvent::Close(Some(WsClose {
code: 1001,
reason: "bye".to_string(),
}))
);
assert!(ws.is_closed());
let frames = decode_sent(&mock.sent());
assert!(frames
.iter()
.any(|(op, pl)| *op == OPCODE_PONG && pl == b"pp"));
}
#[test]
fn websocket_no_autopong_when_disabled() {
let inbound = server_frame(true, OPCODE_PING, b"hi");
let mock = SharedMock::new(inbound);
let mut ws = ws_over(mock.clone(), None);
ws.set_auto_pong(false);
assert_eq!(ws.recv_event().unwrap(), WsEvent::Ping(b"hi".to_vec()));
let frames = decode_sent(&mock.sent());
assert!(
!frames.iter().any(|(op, _)| *op == OPCODE_PONG),
"auto-pong was disabled but a PONG was sent: {frames:?}"
);
}
#[test]
fn websocket_send_pong_and_close_with_write_expected_frames() {
let mock = SharedMock::new(Vec::new());
let mut ws = ws_over(mock.clone(), None);
ws.send_pong(b"keepalive").unwrap();
ws.close_with(1000, "done").unwrap();
let frames = decode_sent(&mock.sent());
assert_eq!(frames.len(), 2);
assert_eq!(frames[0].0, OPCODE_PONG);
assert_eq!(frames[0].1, b"keepalive");
assert_eq!(frames[1].0, OPCODE_CLOSE);
let mut expected = 1000u16.to_be_bytes().to_vec();
expected.extend_from_slice(b"done");
assert_eq!(frames[1].1, expected);
assert!(ws.is_closed());
}
#[test]
fn websocket_send_frame_fragments_a_message() {
let mock = SharedMock::new(Vec::new());
let mut ws = ws_over(mock.clone(), None);
ws.send_frame(false, WsOpcode::Text, b"ab").unwrap();
ws.send_frame(true, WsOpcode::Continuation, b"cd").unwrap();
let raw = mock.sent();
assert_eq!(raw[0] & 0x80, 0x00, "first fragment must have FIN=0");
let frames = decode_sent(&raw);
assert_eq!(frames.len(), 2);
assert_eq!(frames[0].0, OPCODE_TEXT);
assert_eq!(frames[0].1, b"ab");
assert_eq!(frames[1].0, OPCODE_CONT);
assert_eq!(frames[1].1, b"cd");
assert!(ws.send_frame(false, WsOpcode::Ping, b"x").is_err());
}
#[test]
fn websocket_recv_frame_returns_raw_frame() {
let inbound = server_frame(true, OPCODE_TEXT, b"raw");
let mut ws = ws_over(SharedMock::new(inbound), None);
let f = ws.recv_frame().unwrap();
assert_eq!(
f,
WsFrame {
fin: true,
opcode: WsOpcode::Text,
payload: b"raw".to_vec(),
}
);
}
fn decode_sent_fin(sent: &[u8]) -> Vec<(bool, u8, Vec<u8>)> {
let mut frames = Vec::new();
let mut i = 0;
while i < sent.len() {
let fin = (sent[i] & 0x80) != 0;
let opcode = sent[i] & 0x0F;
assert!((sent[i + 1] & 0x80) != 0, "client frame must be masked");
let len7 = sent[i + 1] & 0x7F;
i += 2;
let len = match len7 {
0..=125 => len7 as usize,
126 => {
let l = u16::from_be_bytes([sent[i], sent[i + 1]]) as usize;
i += 2;
l
}
127 => {
let mut b = [0u8; 8];
b.copy_from_slice(&sent[i..i + 8]);
i += 8;
u64::from_be_bytes(b) as usize
}
_ => unreachable!(),
};
let mask = [sent[i], sent[i + 1], sent[i + 2], sent[i + 3]];
i += 4;
let mut payload = sent[i..i + len].to_vec();
i += len;
for (j, b) in payload.iter_mut().enumerate() {
*b ^= mask[j & 3];
}
frames.push((fin, opcode, payload));
}
frames
}
#[test]
fn send_stream_fragments_payload_across_frames() {
let data: Vec<u8> = (0..(STREAM_CHUNK * 2 + 100)).map(|n| n as u8).collect();
let mut sink = Vec::new();
send_stream_inner(OPCODE_BINARY, &mut Cursor::new(data.clone()), |f| {
sink.extend_from_slice(f);
Ok(())
})
.unwrap();
let frames = decode_sent_fin(&sink);
assert_eq!(frames.len(), 3, "two 64 KiB chunks + a tail => 3 frames");
assert_eq!((frames[0].0, frames[0].1), (false, OPCODE_BINARY));
assert_eq!((frames[1].0, frames[1].1), (false, OPCODE_CONT));
assert_eq!((frames[2].0, frames[2].1), (true, OPCODE_CONT));
let got: Vec<u8> = frames.iter().flat_map(|f| f.2.clone()).collect();
assert_eq!(got, data);
}
#[test]
fn send_stream_empty_source_sends_single_fin_frame() {
let mut sink = Vec::new();
send_stream_inner(OPCODE_TEXT, &mut Cursor::new(Vec::new()), |f| {
sink.extend_from_slice(f);
Ok(())
})
.unwrap();
let frames = decode_sent_fin(&sink);
assert_eq!(frames.len(), 1);
assert_eq!((frames[0].0, frames[0].1), (true, OPCODE_TEXT));
assert!(frames[0].2.is_empty());
}
#[test]
fn send_stream_rejects_control_opcode() {
let mut sink = Vec::new();
assert!(
send_stream_inner(OPCODE_PING, &mut Cursor::new(vec![1, 2, 3]), |f| {
sink.extend_from_slice(f);
Ok(())
})
.is_err()
);
}
#[test]
fn send_stream_via_public_writer_api() {
let mock = SharedMock::new(vec![]);
let sent = Arc::clone(&mock.sent);
let mut ws = ws_over(mock, None);
let data = vec![7u8; STREAM_CHUNK + 5];
ws.send_stream(WsOpcode::Binary, &mut Cursor::new(data.clone()))
.unwrap();
let frames = decode_sent_fin(&sent.lock().unwrap());
assert_eq!(frames.len(), 2);
assert_eq!(frames[0].1, OPCODE_BINARY);
assert_eq!((frames[1].0, frames[1].1), (true, OPCODE_CONT));
let got: Vec<u8> = frames.iter().flat_map(|f| f.2.clone()).collect();
assert_eq!(got, data);
}
#[test]
fn keepalive_pings_an_idle_connection() {
let mock = SharedMock::new(vec![]);
let sent = Arc::clone(&mock.sent);
let transport = Arc::new(WsTransport::new(WsTransportKind::Shared(Mutex::new(
Box::new(mock),
))));
let ka = spawn_keepalive(
transport,
Arc::new(AtomicBool::new(false)),
Arc::new(AtomicBool::new(false)),
Duration::from_millis(10),
);
thread::sleep(SHUTDOWN_POLL * 2 + Duration::from_millis(100));
ka.stop();
let frames = decode_sent_fin(&sent.lock().unwrap());
assert!(!frames.is_empty(), "expected at least one keepalive ping");
assert!(frames.iter().all(|f| f.1 == OPCODE_PING && f.0));
}
#[test]
fn keepalive_skips_a_busy_connection() {
let mock = SharedMock::new(vec![]);
let sent = Arc::clone(&mock.sent);
let transport = Arc::new(WsTransport::new(WsTransportKind::Shared(Mutex::new(
Box::new(mock),
))));
let ka = spawn_keepalive(
transport,
Arc::new(AtomicBool::new(false)),
Arc::new(AtomicBool::new(false)),
Duration::from_secs(3600),
);
thread::sleep(SHUTDOWN_POLL + Duration::from_millis(100));
ka.stop();
assert!(
sent.lock().unwrap().is_empty(),
"no ping should fire before the idle interval elapses"
);
}
}