#![allow(
clippy::cast_possible_truncation,
clippy::cast_lossless,
clippy::unwrap_used,
clippy::expect_used,
clippy::indexing_slicing,
clippy::module_name_repetitions
)]
use super::close::{
CloseCode, CloseError, encode_close_payload, is_valid_endpoint_sent, parse_close_payload,
};
use super::frame::{FrameError, FrameHeader, MAX_HEADER_LEN, OpCode, encode_header, parse_header};
use super::handshake::{
HandshakeError, UpgradeRequest, encode_upgrade_request, generate_key, verify_upgrade_response,
};
use super::mask::mask_inplace;
use super::parser::{FeedOutcome, FrameEvent, FrameParser};
use crate::cursor_buf::CursorBuf;
use crate::http::{HttpError, parse_response};
use thiserror::Error;
pub const DEFAULT_MAX_MESSAGE_SIZE: usize = 8 * 1024 * 1024;
pub const DEFAULT_MAX_FRAME_PAYLOAD: u64 = 16 * 1024 * 1024;
pub const MASK_POOL_BYTES: usize = 256;
#[derive(Debug, Clone)]
pub struct WsConfig {
pub host: String,
pub path: String,
pub subprotocols: Vec<String>,
pub origin: Option<String>,
pub max_message_size: usize,
pub max_frame_payload: u64,
pub initial_recv_buffer_capacity: Option<usize>,
pub initial_message_buffer_capacity: Option<usize>,
pub initial_tx_buffer_capacity: Option<usize>,
pub auto_pong: bool,
}
impl WsConfig {
#[must_use]
pub fn new(host: impl Into<String>, path: impl Into<String>) -> Self {
Self {
host: host.into(),
path: path.into(),
subprotocols: Vec::new(),
origin: None,
max_message_size: DEFAULT_MAX_MESSAGE_SIZE,
max_frame_payload: DEFAULT_MAX_FRAME_PAYLOAD,
initial_recv_buffer_capacity: None,
initial_message_buffer_capacity: None,
initial_tx_buffer_capacity: None,
auto_pong: true,
}
}
#[must_use]
pub const fn with_max_message_size(mut self, bytes: usize) -> Self {
self.max_message_size = bytes;
self
}
#[must_use]
pub const fn with_max_frame_payload(mut self, bytes: u64) -> Self {
self.max_frame_payload = bytes;
self
}
#[must_use]
pub const fn with_initial_recv_buffer_capacity(mut self, bytes: usize) -> Self {
self.initial_recv_buffer_capacity = Some(bytes);
self
}
#[must_use]
pub const fn with_initial_message_buffer_capacity(mut self, bytes: usize) -> Self {
self.initial_message_buffer_capacity = Some(bytes);
self
}
#[must_use]
pub const fn with_initial_tx_buffer_capacity(mut self, bytes: usize) -> Self {
self.initial_tx_buffer_capacity = Some(bytes);
self
}
#[must_use]
pub const fn with_initial_buffer_capacities(
mut self,
recv_bytes: usize,
message_bytes: usize,
tx_bytes: usize,
) -> Self {
self.initial_recv_buffer_capacity = Some(recv_bytes);
self.initial_message_buffer_capacity = Some(message_bytes);
self.initial_tx_buffer_capacity = Some(tx_bytes);
self
}
#[must_use]
pub const fn with_auto_pong(mut self, on: bool) -> Self {
self.auto_pong = on;
self
}
}
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum ConnState {
Connecting,
Open,
Closing,
Closed,
}
#[derive(Debug)]
pub enum Event<'a> {
HandshakeComplete,
Text(&'a str),
Binary(&'a [u8]),
Ping(&'a [u8]),
Pong(&'a [u8]),
Close { code: u16, reason: &'a str },
}
#[derive(Debug)]
pub enum DataEvent<'a> {
Text(&'a str),
Binary(&'a [u8]),
}
#[derive(Debug, Error)]
pub enum WsError {
#[error("frame error: {0}")]
Frame(#[from] FrameError),
#[error("handshake error: {0}")]
Handshake(#[from] HandshakeError),
#[error("http parse error: {0}")]
Http(#[from] HttpError),
#[error("close payload error: {0}")]
Close(#[from] CloseError),
#[error("protocol error: {0}")]
Protocol(&'static str),
#[error("operation not allowed in websocket state {0:?}")]
InvalidState(ConnState),
#[error("text frame had invalid UTF-8")]
Utf8Invalid,
#[error("assembled message exceeded max_message_size")]
MessageTooLarge,
#[error("internal state error: {0}")]
Internal(&'static str),
}
#[derive(Debug, Clone, Copy)]
enum EmitKind {
HandshakeComplete,
Text,
Binary,
BorrowedText,
BorrowedBinary,
Ping,
Pong,
Close,
}
#[derive(Debug)]
struct BorrowedPayload {
payload_start: usize,
payload_end: usize,
frame_len: usize,
}
#[derive(Debug)]
pub struct WsClient {
state: ConnState,
parser: FrameParser,
recv_buf: CursorBuf,
msg_buf: Vec<u8>,
borrowed_payload: Option<BorrowedPayload>,
msg_opcode: Option<OpCode>,
ctl_buf: [u8; 125],
ctl_len: u8,
cur_opcode: Option<OpCode>,
cur_fin: bool,
tx_buf: Vec<u8>,
tx_head: usize,
config: WsConfig,
client_key: String,
handshake_started: bool,
last_emitted: Option<EmitKind>,
mask_pool: [u8; MASK_POOL_BYTES],
mask_pool_cursor: usize,
mask_rng: ring::rand::SystemRandom,
}
impl WsClient {
pub fn new_client(config: WsConfig) -> Result<Self, WsError> {
use ring::rand::SecureRandom;
let key = generate_key()?;
let default_msg_cap = config.max_message_size;
let default_io_cap = default_msg_cap
.checked_add(MAX_HEADER_LEN)
.ok_or(WsError::Protocol("max_message_size overflow"))?;
let recv_cap = config
.initial_recv_buffer_capacity
.unwrap_or(default_io_cap);
let msg_cap = config
.initial_message_buffer_capacity
.unwrap_or(default_msg_cap);
let tx_cap = config.initial_tx_buffer_capacity.unwrap_or(default_io_cap);
let mask_rng = ring::rand::SystemRandom::new();
let mut mask_pool = [0_u8; MASK_POOL_BYTES];
mask_rng
.fill(&mut mask_pool)
.map_err(|_| WsError::Handshake(HandshakeError::RngFailure))?;
Ok(Self {
state: ConnState::Connecting,
parser: FrameParser::new(),
recv_buf: CursorBuf::with_capacity(recv_cap),
msg_buf: Vec::with_capacity(msg_cap),
borrowed_payload: None,
msg_opcode: None,
ctl_buf: [0; 125],
ctl_len: 0,
cur_opcode: None,
cur_fin: false,
tx_buf: Vec::with_capacity(tx_cap),
tx_head: 0,
config,
client_key: key,
handshake_started: false,
last_emitted: None,
mask_pool,
mask_pool_cursor: 0,
mask_rng,
})
}
pub fn begin_handshake(&mut self) -> Result<(), WsError> {
if self.state != ConnState::Connecting {
return Err(WsError::InvalidState(self.state));
}
if self.handshake_started {
return Err(WsError::Protocol("handshake already begun"));
}
self.compact_tx_if_needed(256);
let subprotos_refs: Vec<&str> = self
.config
.subprotocols
.iter()
.map(String::as_str)
.collect();
let req = UpgradeRequest {
host: &self.config.host,
path: &self.config.path,
key: &self.client_key,
subprotocols: &subprotos_refs,
origin: self.config.origin.as_deref(),
};
encode_upgrade_request(&mut self.tx_buf, &req);
self.handshake_started = true;
Ok(())
}
pub fn feed_recv(&mut self, bytes: &[u8]) {
self.recv_buf.extend_from_slice(bytes);
}
pub fn poll_event(&mut self) -> Option<Result<Event<'_>, WsError>> {
self.clear_after_emit();
loop {
if self.state == ConnState::Closed {
return None;
}
match self.try_emit_borrowed_data() {
Ok(BorrowedResult::Emit(kind)) => {
self.last_emitted = Some(kind);
return Some(Ok(self.build_event(kind)));
}
Ok(BorrowedResult::WaitForMore) => return None,
Ok(BorrowedResult::Fallback) => {}
Err(e) => {
if self.state != ConnState::Closing {
self.state = ConnState::Closed;
}
return Some(Err(e));
}
}
match self.advance() {
Ok(AdvanceResult::Emit(kind)) => {
self.last_emitted = Some(kind);
return Some(Ok(self.build_event(kind)));
}
Ok(AdvanceResult::Progressed) => continue,
Ok(AdvanceResult::NeedMore) => return None,
Err(e) => {
if self.state != ConnState::Closing {
self.state = ConnState::Closed;
}
return Some(Err(e));
}
}
}
}
fn try_emit_borrowed_data(&mut self) -> Result<BorrowedResult, WsError> {
if self.state != ConnState::Open
|| !self.parser.is_idle()
|| self.msg_opcode.is_some()
|| self.recv_buf.is_empty()
{
return Ok(BorrowedResult::Fallback);
}
let Some((header, header_len)) = parse_header(self.recv_buf.as_slice()).map_err(|e| {
self.queue_close(CloseCode::ProtocolError.as_u16(), "frame parse error");
WsError::Frame(e)
})?
else {
return Ok(BorrowedResult::WaitForMore);
};
if header.mask.is_some() {
self.queue_close(
CloseCode::ProtocolError.as_u16(),
"server sent masked frame",
);
return Err(WsError::Frame(FrameError::ServerSentMaskedFrame));
}
if !header.fin || !matches!(header.opcode, OpCode::Text | OpCode::Binary) {
return Ok(BorrowedResult::Fallback);
}
if header.payload_len > self.config.max_frame_payload {
self.queue_close(CloseCode::MessageTooBig.as_u16(), "frame too large");
return Err(WsError::MessageTooLarge);
}
let payload_len = usize::try_from(header.payload_len).map_err(|_| {
self.queue_close(CloseCode::MessageTooBig.as_u16(), "frame too large");
WsError::MessageTooLarge
})?;
if payload_len > self.config.max_message_size {
self.queue_close(
CloseCode::MessageTooBig.as_u16(),
"message exceeds max_message_size",
);
return Err(WsError::MessageTooLarge);
}
let Some(frame_len) = header_len.checked_add(payload_len) else {
self.queue_close(CloseCode::MessageTooBig.as_u16(), "frame too large");
return Err(WsError::MessageTooLarge);
};
if self.recv_buf.len() < frame_len {
return Ok(BorrowedResult::WaitForMore);
}
let payload_end = frame_len;
let payload = &self.recv_buf.as_slice()[header_len..payload_end];
let kind = match header.opcode {
OpCode::Text => {
if std::str::from_utf8(payload).is_err() {
self.queue_close(CloseCode::InvalidPayload.as_u16(), "invalid utf-8");
return Err(WsError::Utf8Invalid);
}
EmitKind::BorrowedText
}
OpCode::Binary => EmitKind::BorrowedBinary,
_ => unreachable!("guarded above"),
};
self.borrowed_payload = Some(BorrowedPayload {
payload_start: header_len,
payload_end,
frame_len,
});
Ok(BorrowedResult::Emit(kind))
}
pub fn drain_data_events<F>(&mut self, mut sink: F) -> Result<usize, WsError>
where
F: FnMut(DataEvent<'_>),
{
let mut events = 0_usize;
loop {
self.clear_after_emit();
let (direct_events, direct_result) = self.try_drain_data_events(&mut sink)?;
events += direct_events;
match direct_result {
DirectDataResult::Drained => continue,
DirectDataResult::WaitForMore => return Ok(events),
DirectDataResult::Fallback => {}
}
let Some(res) = self.poll_event() else {
return Ok(events);
};
let ev = res?;
events += 1;
match ev {
Event::Text(s) => sink(DataEvent::Text(s)),
Event::Binary(bytes) => sink(DataEvent::Binary(bytes)),
Event::HandshakeComplete
| Event::Ping(_)
| Event::Pong(_)
| Event::Close { .. } => {}
}
}
}
fn try_drain_data_events<F>(
&mut self,
sink: &mut F,
) -> Result<(usize, DirectDataResult), WsError>
where
F: FnMut(DataEvent<'_>),
{
if self.state != ConnState::Open
|| !self.parser.is_idle()
|| self.msg_opcode.is_some()
|| self.recv_buf.is_empty()
{
return Ok((0, DirectDataResult::Fallback));
}
let mut consumed = 0_usize;
let mut events = 0_usize;
loop {
let bytes = &self.recv_buf.as_slice()[consumed..];
if bytes.is_empty() {
self.recv_buf.consume(consumed);
return Ok((events, DirectDataResult::Drained));
}
let (header, header_len) = match parse_header(bytes) {
Ok(Some(parsed)) => parsed,
Ok(None) => {
self.recv_buf.consume(consumed);
return Ok((events, DirectDataResult::WaitForMore));
}
Err(e) => {
self.recv_buf.consume(consumed);
self.queue_close(CloseCode::ProtocolError.as_u16(), "frame parse error");
return Err(WsError::Frame(e));
}
};
if header.mask.is_some() {
self.recv_buf.consume(consumed);
self.queue_close(
CloseCode::ProtocolError.as_u16(),
"server sent masked frame",
);
return Err(WsError::Frame(FrameError::ServerSentMaskedFrame));
}
if !header.fin || !matches!(header.opcode, OpCode::Text | OpCode::Binary) {
self.recv_buf.consume(consumed);
return Ok((events, DirectDataResult::Fallback));
}
if header.payload_len > self.config.max_frame_payload {
self.recv_buf.consume(consumed);
self.queue_close(CloseCode::MessageTooBig.as_u16(), "frame too large");
return Err(WsError::MessageTooLarge);
}
let payload_len = if let Ok(payload_len) = usize::try_from(header.payload_len) {
payload_len
} else {
self.recv_buf.consume(consumed);
self.queue_close(CloseCode::MessageTooBig.as_u16(), "frame too large");
return Err(WsError::MessageTooLarge);
};
if payload_len > self.config.max_message_size {
self.recv_buf.consume(consumed);
self.queue_close(
CloseCode::MessageTooBig.as_u16(),
"message exceeds max_message_size",
);
return Err(WsError::MessageTooLarge);
}
let Some(frame_len) = header_len.checked_add(payload_len) else {
self.recv_buf.consume(consumed);
self.queue_close(CloseCode::MessageTooBig.as_u16(), "frame too large");
return Err(WsError::MessageTooLarge);
};
if bytes.len() < frame_len {
self.recv_buf.consume(consumed);
return Ok((events, DirectDataResult::WaitForMore));
}
let payload = &bytes[header_len..frame_len];
match header.opcode {
OpCode::Text => {
let text = if let Ok(text) = std::str::from_utf8(payload) {
text
} else {
self.recv_buf.consume(consumed);
self.queue_close(CloseCode::InvalidPayload.as_u16(), "invalid utf-8");
return Err(WsError::Utf8Invalid);
};
sink(DataEvent::Text(text));
}
OpCode::Binary => sink(DataEvent::Binary(payload)),
_ => unreachable!("guarded above"),
}
consumed += frame_len;
events += 1;
}
}
pub fn send_text(&mut self, payload: &[u8]) -> Result<(), WsError> {
if std::str::from_utf8(payload).is_err() {
return Err(WsError::Utf8Invalid);
}
self.assert_can_send_data()?;
self.write_frame(true, OpCode::Text, payload);
Ok(())
}
pub fn send_binary(&mut self, payload: &[u8]) -> Result<(), WsError> {
self.assert_can_send_data()?;
self.write_frame(true, OpCode::Binary, payload);
Ok(())
}
pub fn send_ping(&mut self, payload: &[u8]) -> Result<(), WsError> {
if payload.len() > 125 {
return Err(WsError::Protocol("ping payload > 125 bytes"));
}
self.assert_can_send_data()?;
self.write_frame(true, OpCode::Ping, payload);
Ok(())
}
pub fn send_pong(&mut self, payload: &[u8]) -> Result<(), WsError> {
if payload.len() > 125 {
return Err(WsError::Protocol("pong payload > 125 bytes"));
}
self.assert_can_send_data()?;
self.write_frame(true, OpCode::Pong, payload);
Ok(())
}
pub fn send_close(&mut self, code: u16, reason: &str) -> Result<(), WsError> {
if matches!(self.state, ConnState::Closing | ConnState::Closed) {
return Ok(());
}
if self.state != ConnState::Open {
return Err(WsError::InvalidState(self.state));
}
if !is_valid_endpoint_sent(code) {
return Err(WsError::Close(CloseError::InvalidCode(code)));
}
if reason.len() > 123 {
return Err(WsError::Protocol("close reason > 123 bytes"));
}
let mut payload = [0_u8; 125];
let n = encode_close_payload(&mut payload, code, reason);
self.write_frame(true, OpCode::Close, &payload[..n]);
self.state = ConnState::Closing;
Ok(())
}
fn assert_can_send_data(&self) -> Result<(), WsError> {
if self.state == ConnState::Open {
Ok(())
} else {
Err(WsError::InvalidState(self.state))
}
}
#[must_use]
pub fn pending_tx(&self) -> &[u8] {
&self.tx_buf[self.tx_head..]
}
pub fn ack_tx(&mut self, n: usize) {
debug_assert!(n <= self.pending_tx().len());
self.tx_head += n;
if self.tx_head == self.tx_buf.len() {
self.tx_buf.clear();
self.tx_head = 0;
}
}
fn compact_tx_if_needed(&mut self, extra: usize) {
if self.tx_head == 0 {
return;
}
if self.tx_buf.len() + extra > self.tx_buf.capacity() {
self.tx_buf.drain(..self.tx_head);
self.tx_head = 0;
}
}
#[inline]
#[must_use]
pub const fn state(&self) -> ConnState {
self.state
}
fn advance(&mut self) -> Result<AdvanceResult, WsError> {
match self.state {
ConnState::Connecting => self.advance_handshake(),
ConnState::Open | ConnState::Closing => self.advance_frame(),
ConnState::Closed => Ok(AdvanceResult::NeedMore),
}
}
fn advance_handshake(&mut self) -> Result<AdvanceResult, WsError> {
let parsed = parse_response(self.recv_buf.as_slice())?;
let (status, end) = match parsed {
Some((r, end)) => {
let offered: Vec<&str> = self
.config
.subprotocols
.iter()
.map(String::as_str)
.collect();
verify_upgrade_response(&r, &self.client_key, &offered)?;
(r.status, end)
}
None => return Ok(AdvanceResult::NeedMore),
};
debug_assert_eq!(status, 101);
self.recv_buf.consume(end);
self.state = ConnState::Open;
Ok(AdvanceResult::Emit(EmitKind::HandshakeComplete))
}
fn advance_frame(&mut self) -> Result<AdvanceResult, WsError> {
let (consumed, action) = {
let outcome = match self.parser.feed_one(self.recv_buf.as_slice()) {
Ok(outcome) => outcome,
Err(e) => {
self.queue_close(CloseCode::ProtocolError.as_u16(), "frame parse error");
return Err(WsError::Frame(e));
}
};
match outcome {
FeedOutcome::NeedMore { consumed } => (consumed, Action::NeedMore),
FeedOutcome::Event {
consumed,
event: FrameEvent::FrameStart(h),
} => (consumed, Action::FrameStart(h)),
FeedOutcome::Event {
consumed,
event: FrameEvent::PayloadChunk(slice),
} => {
if let Some(op) = self.cur_opcode {
if op.is_control() {
let start = self.ctl_len as usize;
let end = start + slice.len();
self.ctl_buf[start..end].copy_from_slice(slice);
self.ctl_len = end as u8;
} else {
self.msg_buf.extend_from_slice(slice);
}
}
(consumed, Action::PayloadChunk)
}
FeedOutcome::Event {
consumed,
event: FrameEvent::FrameEnd,
} => (consumed, Action::FrameEnd),
}
};
self.recv_buf.consume(consumed);
match action {
Action::NeedMore => Ok(AdvanceResult::NeedMore),
Action::FrameStart(h) => {
self.on_frame_start(h)?;
Ok(AdvanceResult::Progressed)
}
Action::PayloadChunk => Ok(AdvanceResult::Progressed),
Action::FrameEnd => match self.on_frame_end()? {
Some(kind) => Ok(AdvanceResult::Emit(kind)),
None => Ok(AdvanceResult::Progressed),
},
}
}
fn on_frame_start(&mut self, h: FrameHeader) -> Result<(), WsError> {
if h.payload_len > self.config.max_frame_payload {
self.queue_close(CloseCode::MessageTooBig.as_u16(), "frame too large");
return Err(WsError::MessageTooLarge);
}
if h.opcode.is_control() {
self.ctl_len = 0;
self.cur_opcode = Some(h.opcode);
self.cur_fin = h.fin;
return Ok(());
}
match h.opcode {
OpCode::Continuation => {
if self.msg_opcode.is_none() {
self.queue_close(
CloseCode::ProtocolError.as_u16(),
"continuation without start",
);
return Err(WsError::Protocol("continuation without start"));
}
}
OpCode::Text | OpCode::Binary => {
if self.msg_opcode.is_some() {
self.queue_close(
CloseCode::ProtocolError.as_u16(),
"new data frame mid-fragmentation",
);
return Err(WsError::Protocol("nested data frame"));
}
self.msg_opcode = Some(h.opcode);
self.msg_buf.clear();
}
_ => return Err(WsError::Internal("non-data non-control opcode")),
}
let prospective = self.msg_buf.len().saturating_add(h.payload_len as usize);
if prospective > self.config.max_message_size {
self.queue_close(
CloseCode::MessageTooBig.as_u16(),
"message exceeds max_message_size",
);
return Err(WsError::MessageTooLarge);
}
self.cur_opcode = Some(h.opcode);
self.cur_fin = h.fin;
Ok(())
}
fn on_frame_end(&mut self) -> Result<Option<EmitKind>, WsError> {
let opcode = self
.cur_opcode
.take()
.ok_or(WsError::Internal("FrameEnd without FrameStart"))?;
let fin = self.cur_fin;
if opcode.is_control() {
return match opcode {
OpCode::Ping => {
if self.config.auto_pong {
let len = self.ctl_len as usize;
let mut payload = [0_u8; 125];
payload[..len].copy_from_slice(&self.ctl_buf[..len]);
self.write_frame(true, OpCode::Pong, &payload[..len]);
}
Ok(Some(EmitKind::Ping))
}
OpCode::Pong => Ok(Some(EmitKind::Pong)),
OpCode::Close => {
self.on_received_close()?;
Ok(Some(EmitKind::Close))
}
_ => Err(WsError::Internal("bad control opcode")),
};
}
if !fin {
return Ok(None);
}
let msg_opcode = self
.msg_opcode
.take()
.ok_or(WsError::Internal("FrameEnd FIN without active message"))?;
match msg_opcode {
OpCode::Text => {
if std::str::from_utf8(&self.msg_buf).is_err() {
self.queue_close(CloseCode::InvalidPayload.as_u16(), "invalid utf-8");
return Err(WsError::Utf8Invalid);
}
Ok(Some(EmitKind::Text))
}
OpCode::Binary => Ok(Some(EmitKind::Binary)),
_ => Err(WsError::Internal("non-data msg_opcode")),
}
}
fn on_received_close(&mut self) -> Result<(), WsError> {
let payload = &self.ctl_buf[..self.ctl_len as usize];
let parsed = parse_close_payload(payload).map_err(WsError::Close);
let (echo_code, _echo_reason) = match parsed {
Ok(Some((code, reason))) => (code, reason),
Ok(None) => (CloseCode::Normal.as_u16(), ""),
Err(e) => {
self.queue_close(CloseCode::ProtocolError.as_u16(), "bad close payload");
self.state = ConnState::Closing;
return Err(e);
}
};
if self.state == ConnState::Open {
self.queue_close(echo_code, "");
}
self.state = ConnState::Closed;
Ok(())
}
fn queue_close(&mut self, code: u16, reason: &str) {
if matches!(self.state, ConnState::Closing | ConnState::Closed) {
return;
}
let mut payload = [0_u8; 125];
let n = encode_close_payload(&mut payload, code, reason);
self.write_frame(true, OpCode::Close, &payload[..n]);
if self.state == ConnState::Open {
self.state = ConnState::Closing;
}
}
fn write_frame(&mut self, fin: bool, opcode: OpCode, payload: &[u8]) {
let mask = self.next_mask();
let mut hdr = [0_u8; MAX_HEADER_LEN];
let n = encode_header(&mut hdr, fin, opcode, Some(mask), payload.len() as u64);
self.compact_tx_if_needed(n + payload.len());
self.tx_buf.extend_from_slice(&hdr[..n]);
let payload_start = self.tx_buf.len();
self.tx_buf.extend_from_slice(payload);
let payload_slice = &mut self.tx_buf[payload_start..];
mask_inplace(payload_slice, mask);
}
fn next_mask(&mut self) -> [u8; 4] {
if self.mask_pool_cursor + 4 > MASK_POOL_BYTES {
use ring::rand::SecureRandom;
self.mask_rng
.fill(&mut self.mask_pool)
.expect("SystemRandom::fill must not fail on supported platforms");
self.mask_pool_cursor = 0;
}
let i = self.mask_pool_cursor;
let key = [
self.mask_pool[i],
self.mask_pool[i + 1],
self.mask_pool[i + 2],
self.mask_pool[i + 3],
];
self.mask_pool_cursor = i + 4;
key
}
fn build_event(&self, kind: EmitKind) -> Event<'_> {
match kind {
EmitKind::HandshakeComplete => Event::HandshakeComplete,
EmitKind::Text => {
let s = std::str::from_utf8(&self.msg_buf).unwrap_or("");
Event::Text(s)
}
EmitKind::Binary => Event::Binary(&self.msg_buf),
EmitKind::BorrowedText => {
let payload = self.borrowed_payload();
let s = std::str::from_utf8(payload)
.unwrap_or_else(|_| unreachable!("borrowed text must retain valid utf-8"));
Event::Text(s)
}
EmitKind::BorrowedBinary => Event::Binary(self.borrowed_payload()),
EmitKind::Ping => Event::Ping(&self.ctl_buf[..self.ctl_len as usize]),
EmitKind::Pong => Event::Pong(&self.ctl_buf[..self.ctl_len as usize]),
EmitKind::Close => {
let payload = &self.ctl_buf[..self.ctl_len as usize];
match parse_close_payload(payload) {
Ok(Some((code, reason))) => Event::Close { code, reason },
Ok(None) => Event::Close {
code: 1005,
reason: "",
},
Err(_) => unreachable!("on_received_close gates Err before EmitKind::Close"),
}
}
}
}
fn clear_after_emit(&mut self) {
match self.last_emitted.take() {
None => {}
Some(EmitKind::HandshakeComplete) => {}
Some(EmitKind::Text | EmitKind::Binary) => {
self.msg_buf.clear();
}
Some(EmitKind::BorrowedText | EmitKind::BorrowedBinary) => {
let borrowed = self
.borrowed_payload
.take()
.expect("borrowed emit must retain payload range");
self.recv_buf.consume(borrowed.frame_len);
}
Some(EmitKind::Ping | EmitKind::Pong | EmitKind::Close) => {
self.ctl_len = 0;
}
}
}
fn borrowed_payload(&self) -> &[u8] {
let borrowed = self
.borrowed_payload
.as_ref()
.expect("borrowed emit must retain payload range");
&self.recv_buf.as_slice()[borrowed.payload_start..borrowed.payload_end]
}
}
#[derive(Debug)]
enum Action {
NeedMore,
FrameStart(FrameHeader),
PayloadChunk,
FrameEnd,
}
#[derive(Debug)]
enum AdvanceResult {
Emit(EmitKind),
Progressed,
NeedMore,
}
enum BorrowedResult {
Emit(EmitKind),
WaitForMore,
Fallback,
}
enum DirectDataResult {
Drained,
WaitForMore,
Fallback,
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::indexing_slicing, clippy::panic)]
mod tests {
use super::*;
fn mk_client() -> WsClient {
WsClient::new_client(WsConfig::new("example.com", "/ws")).unwrap()
}
fn server_text(payload: &[u8]) -> Vec<u8> {
let mut buf = vec![0_u8; MAX_HEADER_LEN];
let n = super::super::frame::encode_header(
&mut buf,
true,
OpCode::Text,
None,
payload.len() as u64,
);
buf.truncate(n);
buf.extend_from_slice(payload);
buf
}
fn server_binary(payload: &[u8]) -> Vec<u8> {
let mut buf = vec![0_u8; MAX_HEADER_LEN];
let n = super::super::frame::encode_header(
&mut buf,
true,
OpCode::Binary,
None,
payload.len() as u64,
);
buf.truncate(n);
buf.extend_from_slice(payload);
buf
}
fn server_close(code: u16, reason: &str) -> Vec<u8> {
let mut payload = [0_u8; 125];
let n = encode_close_payload(&mut payload, code, reason);
let mut buf = vec![0_u8; MAX_HEADER_LEN];
let hn = super::super::frame::encode_header(&mut buf, true, OpCode::Close, None, n as u64);
buf.truncate(hn);
buf.extend_from_slice(&payload[..n]);
buf
}
fn server_control(opcode: OpCode, payload: &[u8]) -> Vec<u8> {
let mut buf = vec![0_u8; MAX_HEADER_LEN];
let hn =
super::super::frame::encode_header(&mut buf, true, opcode, None, payload.len() as u64);
buf.truncate(hn);
buf.extend_from_slice(payload);
buf
}
fn fake_101_response(client_key: &str) -> Vec<u8> {
let accept = super::super::handshake::compute_accept(client_key);
let mut s = String::new();
s.push_str("HTTP/1.1 101 Switching Protocols\r\n");
s.push_str("Upgrade: websocket\r\n");
s.push_str("Connection: Upgrade\r\n");
s.push_str(&format!("Sec-WebSocket-Accept: {accept}\r\n"));
s.push_str("\r\n");
s.into_bytes()
}
#[test]
fn handshake_then_text_then_close() {
let mut c = mk_client();
c.begin_handshake().unwrap();
assert!(!c.pending_tx().is_empty());
let n = c.pending_tx().len();
c.ack_tx(n);
let resp = fake_101_response(&c.client_key);
c.feed_recv(&resp);
match c.poll_event() {
Some(Ok(Event::HandshakeComplete)) => {}
other => panic!("expected HandshakeComplete, got {other:?}"),
}
assert_eq!(c.state(), ConnState::Open);
c.feed_recv(&server_text(b"hi"));
match c.poll_event() {
Some(Ok(Event::Text(s))) => assert_eq!(s, "hi"),
other => panic!("expected Text, got {other:?}"),
}
c.feed_recv(&server_close(1000, "bye"));
match c.poll_event() {
Some(Ok(Event::Close { code, reason })) => {
assert_eq!(code, 1000);
assert_eq!(reason, "bye");
}
other => panic!("expected Close, got {other:?}"),
}
assert_eq!(c.state(), ConnState::Closed);
assert!(!c.pending_tx().is_empty());
}
#[test]
fn begin_handshake_rejects_duplicate_call() {
let mut c = mk_client();
c.begin_handshake().unwrap();
let pending_len = c.pending_tx().len();
let err = c.begin_handshake().unwrap_err();
assert!(matches!(err, WsError::Protocol("handshake already begun")));
assert_eq!(c.pending_tx().len(), pending_len);
}
#[test]
fn complete_text_uses_borrowed_payload_path() {
let mut c = mk_client();
c.begin_handshake().unwrap();
c.ack_tx(c.pending_tx().len());
c.feed_recv(&fake_101_response(&c.client_key));
c.poll_event();
c.feed_recv(&server_text(b"borrowed"));
match c.poll_event() {
Some(Ok(Event::Text(s))) => assert_eq!(s, "borrowed"),
other => panic!("expected borrowed Text, got {other:?}"),
}
assert!(c.msg_buf.is_empty());
assert!(c.borrowed_payload.is_some());
assert!(c.poll_event().is_none());
assert!(c.borrowed_payload.is_none());
assert!(c.recv_buf.is_empty());
}
#[test]
fn split_text_waits_then_uses_borrowed_payload_path() {
let mut c = mk_client();
c.begin_handshake().unwrap();
c.ack_tx(c.pending_tx().len());
c.feed_recv(&fake_101_response(&c.client_key));
c.poll_event();
let frame = server_text(b"split");
c.feed_recv(&frame[..4]);
assert!(c.poll_event().is_none());
assert!(c.parser.is_idle());
assert!(c.msg_buf.is_empty());
c.feed_recv(&frame[4..]);
match c.poll_event() {
Some(Ok(Event::Text(s))) => assert_eq!(s, "split"),
other => panic!("expected borrowed Text, got {other:?}"),
}
assert!(c.msg_buf.is_empty());
assert!(c.borrowed_payload.is_some());
}
#[test]
fn ping_triggers_auto_pong() {
let mut c = mk_client();
c.begin_handshake().unwrap();
c.ack_tx(c.pending_tx().len());
c.feed_recv(&fake_101_response(&c.client_key));
c.poll_event();
let mut frame = vec![0_u8; MAX_HEADER_LEN];
let n = super::super::frame::encode_header(&mut frame, true, OpCode::Ping, None, 4);
frame.truncate(n);
frame.extend_from_slice(b"ping");
c.feed_recv(&frame);
match c.poll_event() {
Some(Ok(Event::Ping(p))) => assert_eq!(p, b"ping"),
other => panic!("{other:?}"),
}
let tx = c.pending_tx();
assert!(!tx.is_empty());
assert_eq!(tx[0], 0x8A);
}
#[test]
fn drain_data_events_handles_control_and_dispatches_text_binary() {
let mut c = mk_client();
c.begin_handshake().unwrap();
c.ack_tx(c.pending_tx().len());
c.feed_recv(&fake_101_response(&c.client_key));
c.poll_event(); c.ack_tx(c.pending_tx().len());
let mut wire = Vec::new();
wire.extend_from_slice(&server_text(b"{\"op\":\"subscribed\"}"));
wire.extend_from_slice(&server_control(OpCode::Ping, b"hb"));
wire.extend_from_slice(&server_binary(b"\x01\x02\x03\x04"));
c.feed_recv(&wire);
let mut data = Vec::new();
let events = c
.drain_data_events(|ev| match ev {
DataEvent::Text(s) => data.push(format!("text:{s}")),
DataEvent::Binary(bytes) => data.push(format!("binary:{}", bytes.len())),
})
.unwrap();
assert_eq!(events, 3);
assert_eq!(
data,
vec![
"text:{\"op\":\"subscribed\"}".to_owned(),
"binary:4".to_owned()
]
);
assert!(
!c.pending_tx().is_empty(),
"Ping should have queued an auto-pong"
);
}
#[test]
fn outgoing_text_is_masked() {
let mut c = mk_client();
c.begin_handshake().unwrap();
c.ack_tx(c.pending_tx().len());
c.feed_recv(&fake_101_response(&c.client_key));
c.poll_event();
assert_eq!(c.state(), ConnState::Open);
c.send_text(b"hello").unwrap();
let tx = c.pending_tx();
assert_eq!(tx[0], 0x81);
assert_eq!(tx[1], 0x85);
assert_eq!(tx.len(), 11);
let key = [tx[2], tx[3], tx[4], tx[5]];
let mut payload: Vec<u8> = tx[6..].to_vec();
super::super::mask::mask_inplace(&mut payload, key);
assert_eq!(&payload, b"hello");
}
#[test]
fn outgoing_pong_is_masked() {
let mut c = mk_client();
c.begin_handshake().unwrap();
c.ack_tx(c.pending_tx().len());
c.feed_recv(&fake_101_response(&c.client_key));
c.poll_event();
assert_eq!(c.state(), ConnState::Open);
c.send_pong(b"pong").unwrap();
let tx = c.pending_tx();
assert_eq!(tx[0], 0x8A);
assert_eq!(tx[1], 0x84);
assert_eq!(tx.len(), 10);
let key = [tx[2], tx[3], tx[4], tx[5]];
let mut payload: Vec<u8> = tx[6..].to_vec();
super::super::mask::mask_inplace(&mut payload, key);
assert_eq!(&payload, b"pong");
}
#[test]
fn send_rejects_invalid_payloads_in_release_path() {
let mut c = mk_client();
assert!(matches!(
c.send_text(b"hello").unwrap_err(),
WsError::InvalidState(ConnState::Connecting)
));
assert!(matches!(
c.send_pong(b"pong").unwrap_err(),
WsError::InvalidState(ConnState::Connecting)
));
c.begin_handshake().unwrap();
c.ack_tx(c.pending_tx().len());
c.feed_recv(&fake_101_response(&c.client_key));
c.poll_event();
assert_eq!(c.state(), ConnState::Open);
assert!(matches!(c.send_text(&[0xFF]), Err(WsError::Utf8Invalid)));
assert!(matches!(
c.send_ping(&[0_u8; 126]),
Err(WsError::Protocol("ping payload > 125 bytes"))
));
assert!(matches!(
c.send_pong(&[0_u8; 126]),
Err(WsError::Protocol("pong payload > 125 bytes"))
));
assert!(matches!(
c.send_close(1006, ""),
Err(WsError::Close(CloseError::InvalidCode(1006)))
));
assert!(matches!(
c.send_close(1000, "x".repeat(124).as_str()),
Err(WsError::Protocol("close reason > 123 bytes"))
));
}
#[test]
fn frame_parse_error_queues_close_before_error() {
let mut c = mk_client();
c.begin_handshake().unwrap();
c.ack_tx(c.pending_tx().len());
c.feed_recv(&fake_101_response(&c.client_key));
c.poll_event();
c.feed_recv(b"\x81\x85\x12\x34\x56\x78Hello");
match c.poll_event() {
Some(Err(WsError::Frame(FrameError::ServerSentMaskedFrame))) => {}
other => panic!("{other:?}"),
}
assert_eq!(c.state(), ConnState::Closing);
assert!(!c.pending_tx().is_empty());
}
#[test]
fn invalid_close_payload_queues_protocol_close() {
let mut c = mk_client();
c.begin_handshake().unwrap();
c.ack_tx(c.pending_tx().len());
c.feed_recv(&fake_101_response(&c.client_key));
c.poll_event();
c.feed_recv(&server_control(OpCode::Close, &[0x03]));
match c.poll_event() {
Some(Err(WsError::Close(CloseError::OneByte))) => {}
other => panic!("{other:?}"),
}
assert_eq!(c.state(), ConnState::Closing);
assert!(!c.pending_tx().is_empty());
}
#[test]
fn fragmented_text_message_assembled() {
let mut c = mk_client();
c.begin_handshake().unwrap();
c.ack_tx(c.pending_tx().len());
c.feed_recv(&fake_101_response(&c.client_key));
c.poll_event();
let mut buf = vec![0_u8; MAX_HEADER_LEN];
let n = super::super::frame::encode_header(&mut buf, false, OpCode::Text, None, 3);
buf.truncate(n);
buf.extend_from_slice(b"hel");
let mut buf2 = vec![0_u8; MAX_HEADER_LEN];
let n2 = super::super::frame::encode_header(&mut buf2, true, OpCode::Continuation, None, 2);
buf2.truncate(n2);
buf2.extend_from_slice(b"lo");
c.feed_recv(&buf);
while let Some(_ev) = c.poll_event() {
}
c.feed_recv(&buf2);
match c.poll_event() {
Some(Ok(Event::Text(s))) => assert_eq!(s, "hello"),
other => panic!("{other:?}"),
}
}
#[test]
fn bad_handshake_status_rejected() {
let mut c = mk_client();
c.begin_handshake().unwrap();
c.feed_recv(b"HTTP/1.1 400 Bad Request\r\n\r\n");
match c.poll_event() {
Some(Err(WsError::Handshake(HandshakeError::BadStatus(400)))) => {}
other => panic!("{other:?}"),
}
assert_eq!(c.state(), ConnState::Closed);
}
}