use std::{
io,
pin::Pin,
sync::{Arc, Mutex as StdMutex},
task::{Context, Poll},
time::Duration,
};
#[cfg(feature = "server")]
use axum::extract::ws::{CloseFrame, Message as AxumMessage, WebSocket};
use futures::channel::mpsc as futures_mpsc;
use futures::{SinkExt, StreamExt};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::sync::mpsc;
use tokio::sync::oneshot;
pub const WS_MESSAGE_CAP: usize = 1024 * 1024;
const READ_SLOTS: usize = 64;
const WRITE_SLOTS: usize = 64;
pub const INBOUND_WS_MESSAGE_CAP: usize = 1024 * 1024;
pub const INBOUND_WS_FRAME_CAP: usize = 1024 * 1024;
pub const PENDING_BUFFER_CAP: usize = MAX_CHUNK_LEN as usize + 8;
pub const WS_PROTOCOL_ERROR: u16 = 1002;
pub const WS_INTERNAL_ERROR: u16 = 1011;
pub const WS_GOING_AWAY: u16 = 1001;
pub const DEFAULT_WS_WRITE_TIMEOUT: Duration = Duration::from_secs(60);
pub const DEFAULT_WS_IDLE_TIMEOUT: Duration = Duration::from_secs(60);
struct InboundChunkProgress {
header: [u8; 8],
header_fill: u8,
payload_remaining: u64,
}
impl InboundChunkProgress {
fn new() -> Self {
Self {
header: [0u8; 8],
header_fill: 0,
payload_remaining: 0,
}
}
fn observe(&mut self, bytes: &[u8]) -> u32 {
let mut completed = 0u32;
for &byte in bytes {
if self.payload_remaining > 0 {
self.payload_remaining -= 1;
if self.payload_remaining == 0 {
completed += 1;
}
continue;
}
self.header[self.header_fill as usize] = byte;
self.header_fill += 1;
if self.header_fill == 8 {
self.header_fill = 0;
let len = u32::from_be_bytes([
self.header[4],
self.header[5],
self.header[6],
self.header[7],
]);
self.payload_remaining = len as u64;
if len == 0 {
completed += 1;
}
}
}
completed
}
}
pub use alkcall::channels::wire::MAX_CHUNK_LEN;
pub(crate) enum WriteMsg {
Bytes(Vec<u8>),
CloseWith(u16, &'static str),
}
type WriteErrorSlot = Arc<StdMutex<Option<oneshot::Sender<&'static str>>>>;
fn make_write_error_slot() -> (WriteErrorSlot, oneshot::Receiver<&'static str>) {
let (tx, rx) = oneshot::channel::<&'static str>();
(Arc::new(StdMutex::new(Some(tx))), rx)
}
fn send_write_error(slot: &WriteErrorSlot, reason: &'static str) {
if let Some(tx) = slot.lock().unwrap_or_else(|e| e.into_inner()).take() {
let _ = tx.send(reason);
}
}
pub(crate) mod close_reason {
pub(crate) const IDLE_READ_TIMEOUT: &str =
"connection made no inbound chunk progress past the read timeout";
pub(crate) const TEXT_NOT_SUPPORTED: &str = "text messages not supported";
pub(crate) const INBOUND_FRAME_REJECTED: &str = "inbound frame rejected (size cap)";
}
trait WsFraming: Sized {
type Bytes: AsRef<[u8]>;
type Msg;
fn binary(msg: &Self::Msg) -> Option<&Self::Bytes>;
fn is_text(msg: &Self::Msg) -> bool;
fn is_close(msg: &Self::Msg) -> bool;
fn close_message(code: u16, reason: &'static str) -> Self::Msg;
fn binary_message(bytes: Vec<u8>) -> Self::Msg;
}
trait IntoWsResult<M: WsFraming> {
fn into_ws_result(self) -> Result<M::Msg, ()>;
}
impl<M, E> IntoWsResult<M> for Result<M::Msg, E>
where
M: WsFraming,
{
fn into_ws_result(self) -> Result<M::Msg, ()> {
self.map_err(|_| ())
}
}
async fn run_read_pump<M, S, F>(
mut ws_stream: S,
read_tx: mpsc::Sender<Vec<u8>>,
write_tx_for_read: futures_mpsc::Sender<WriteMsg>,
write_error_slot_for_read: WriteErrorSlot,
idle_timeout: Option<Duration>,
on_end: F,
) where
S: futures::Stream + Unpin,
M: WsFraming,
S::Item: IntoWsResult<M>,
F: FnOnce(),
{
let mut progress = InboundChunkProgress::new();
let mut last_progress_at = tokio::time::Instant::now();
loop {
let budget = idle_timeout.map(|window| window.saturating_sub(last_progress_at.elapsed()));
let msg = match budget {
None => ws_stream.next().await,
Some(budget) => match tokio::time::timeout(budget, ws_stream.next()).await {
Ok(msg) => msg,
Err(_elapsed) => {
send_write_error(&write_error_slot_for_read, close_reason::IDLE_READ_TIMEOUT);
let _ = write_tx_for_read
.clone()
.send(WriteMsg::CloseWith(
WS_GOING_AWAY,
close_reason::IDLE_READ_TIMEOUT,
))
.await;
break;
}
},
};
let Some(msg) = msg else {
break;
};
match msg.into_ws_result() {
Ok(m) => {
if let Some(b) = M::binary(&m) {
let completed = progress.observe(b.as_ref());
if read_tx.send(b.as_ref().to_vec()).await.is_err() {
break;
}
if completed > 0 {
last_progress_at = tokio::time::Instant::now();
}
} else if M::is_text(&m) {
let _ = write_tx_for_read
.clone()
.send(WriteMsg::CloseWith(
WS_PROTOCOL_ERROR,
close_reason::TEXT_NOT_SUPPORTED,
))
.await;
break;
} else if M::is_close(&m) {
break;
}
}
Err(()) => {
send_write_error(
&write_error_slot_for_read,
close_reason::INBOUND_FRAME_REJECTED,
);
let _ = write_tx_for_read
.clone()
.send(WriteMsg::CloseWith(
WS_INTERNAL_ERROR,
close_reason::INBOUND_FRAME_REJECTED,
))
.await;
break;
}
}
}
on_end();
}
async fn fail_write_pump<M, S>(
ws_sink: &mut futures::stream::SplitSink<S, <M as WsFraming>::Msg>,
slot: &WriteErrorSlot,
reason: &'static str,
) where
M: WsFraming,
S: futures::Sink<<M as WsFraming>::Msg> + Unpin,
{
send_write_error(slot, reason);
let _ = ws_sink
.send(<M as WsFraming>::close_message(WS_INTERNAL_ERROR, reason))
.await;
}
async fn run_write_pump<M, S>(
mut ws_sink: futures::stream::SplitSink<S, <M as WsFraming>::Msg>,
mut write_rx: futures_mpsc::Receiver<WriteMsg>,
write_error_slot: WriteErrorSlot,
write_timeout: Option<Duration>,
) where
S: futures::Sink<<M as WsFraming>::Msg> + Unpin,
M: WsFraming,
{
let mut pending: Vec<u8> = Vec::new();
while let Some(msg) = write_rx.next().await {
match msg {
WriteMsg::CloseWith(code, reason) => {
let close = <M as WsFraming>::close_message(code, reason);
match write_timeout {
None => {
let _ = ws_sink.send(close).await;
}
Some(window) => {
let _ = tokio::time::timeout(window, ws_sink.send(close)).await;
}
}
break;
}
WriteMsg::Bytes(b) => {
debug_assert!(b.len() <= PENDING_BUFFER_CAP);
pending.extend_from_slice(&b);
}
}
loop {
if pending.len() < 8 {
break;
}
let len_bytes = [pending[4], pending[5], pending[6], pending[7]];
let len = u32::from_be_bytes(len_bytes);
if len > MAX_CHUNK_LEN {
fail_write_pump::<M, _>(
&mut ws_sink,
&write_error_slot,
"chunk length exceeds MAX_CHUNK_LEN",
)
.await;
return;
}
let total = 8usize.saturating_add(len as usize);
if pending.len() < total {
break;
}
let chunk: Vec<u8> = pending.drain(..total).collect();
for piece in chunk.chunks(WS_MESSAGE_CAP) {
if !send_bounded::<M, _>(&mut ws_sink, piece, write_timeout, &write_error_slot)
.await
{
return;
}
}
}
if pending.len() > PENDING_BUFFER_CAP {
fail_write_pump::<M, _>(
&mut ws_sink,
&write_error_slot,
"write pending buffer exceeded cap",
)
.await;
return;
}
}
let _ = ws_sink.close().await;
}
async fn send_bounded<M, S>(
ws_sink: &mut futures::stream::SplitSink<S, <M as WsFraming>::Msg>,
piece: &[u8],
write_timeout: Option<Duration>,
slot: &WriteErrorSlot,
) -> bool
where
M: WsFraming,
S: futures::Sink<<M as WsFraming>::Msg> + Unpin,
{
match write_timeout {
None => ws_sink
.send(<M as WsFraming>::binary_message(piece.to_vec()))
.await
.is_ok(),
Some(window) => {
let send = ws_sink.send(<M as WsFraming>::binary_message(piece.to_vec()));
match tokio::time::timeout(window, send).await {
Ok(Ok(())) => true,
Ok(Err(_)) => false,
Err(_elapsed) => {
send_write_error(
slot,
"connection made no write progress past the write timeout",
);
false
}
}
}
}
}
pub struct WsByteStream {
read_rx: mpsc::Receiver<Vec<u8>>,
read_buf: Vec<u8>,
read_pos: usize,
eof: bool,
write_tx: Option<futures_mpsc::Sender<WriteMsg>>,
write_open: bool,
write_error: oneshot::Receiver<&'static str>,
write_failed: bool,
}
pub struct WsPumps {
read_task: tokio::task::JoinHandle<()>,
write_task: tokio::task::JoinHandle<()>,
#[cfg_attr(
all(any(test, feature = "wss"), not(feature = "client")),
allow(dead_code)
)]
#[cfg(any(test, feature = "wss"))]
read_eof: tokio::sync::watch::Sender<bool>,
}
impl WsPumps {
pub fn abort(&self) {
self.read_task.abort();
self.write_task.abort();
}
#[cfg(any(test, feature = "wss"))]
#[cfg_attr(not(feature = "client"), allow(dead_code))]
pub(crate) fn read_eof(&self) -> tokio::sync::watch::Receiver<bool> {
self.read_eof.subscribe()
}
}
#[cfg(feature = "server")]
struct AxumFraming;
#[cfg(feature = "server")]
impl WsFraming for AxumFraming {
type Bytes = axum::body::Bytes;
type Msg = AxumMessage;
fn binary(msg: &AxumMessage) -> Option<&Self::Bytes> {
match msg {
AxumMessage::Binary(b) => Some(b),
_ => None,
}
}
fn is_text(msg: &AxumMessage) -> bool {
matches!(msg, AxumMessage::Text(_))
}
fn is_close(msg: &AxumMessage) -> bool {
matches!(msg, AxumMessage::Close(_))
}
fn close_message(code: u16, reason: &'static str) -> AxumMessage {
AxumMessage::Close(Some(CloseFrame {
code,
reason: reason.into(),
}))
}
fn binary_message(bytes: Vec<u8>) -> AxumMessage {
AxumMessage::Binary(bytes.into())
}
}
#[cfg(feature = "server")]
pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
split_ws_to_bytes_idle(socket, Some(DEFAULT_WS_IDLE_TIMEOUT))
}
#[cfg(feature = "server")]
pub fn split_ws_to_bytes_idle(
socket: WebSocket,
idle_timeout: Option<Duration>,
) -> (WsByteStream, WsPumps) {
split_ws_to_bytes_idle_with_write(socket, idle_timeout, None)
}
#[cfg(feature = "server")]
pub fn split_ws_to_bytes_idle_with_write(
socket: WebSocket,
idle_timeout: Option<Duration>,
write_timeout: Option<Duration>,
) -> (WsByteStream, WsPumps) {
let (ws_sink, ws_stream) = socket.split();
let (read_tx, read_rx) = mpsc::channel::<Vec<u8>>(READ_SLOTS);
let (write_tx, write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
let (write_error_slot, write_error_rx) = make_write_error_slot();
#[cfg(any(test, feature = "wss"))]
let read_eof = tokio::sync::watch::channel(false).0;
let write_tx_for_read = write_tx.clone();
let write_error_slot_for_read = Arc::clone(&write_error_slot);
#[cfg(any(test, feature = "wss"))]
let read_eof_for_task = read_eof.clone();
let read_task = tokio::spawn(run_read_pump::<AxumFraming, _, _>(
ws_stream,
read_tx,
write_tx_for_read,
write_error_slot_for_read,
idle_timeout,
move || {
#[cfg(any(test, feature = "wss"))]
{
let _ = read_eof_for_task.send(true);
}
},
));
let write_task = tokio::spawn(run_write_pump::<AxumFraming, _>(
ws_sink,
write_rx,
write_error_slot,
Some(write_timeout.unwrap_or(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT)),
));
(
WsByteStream {
read_rx,
read_buf: Vec::new(),
read_pos: 0,
eof: false,
write_tx: Some(write_tx),
write_open: true,
write_error: write_error_rx,
write_failed: false,
},
WsPumps {
read_task,
write_task,
#[cfg(any(test, feature = "wss"))]
read_eof,
},
)
}
impl AsyncRead for WsByteStream {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut tokio::io::ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let this = self.get_mut();
loop {
if this.read_pos < this.read_buf.len() {
let n = (this.read_buf.len() - this.read_pos).min(buf.remaining());
let end = this.read_pos + n;
buf.put_slice(&this.read_buf[this.read_pos..end]);
this.read_pos = end;
if this.read_pos == this.read_buf.len() {
this.read_buf.clear();
this.read_pos = 0;
}
return Poll::Ready(Ok(()));
}
if this.eof {
return Poll::Ready(Ok(()));
}
match this.read_rx.poll_recv(cx) {
Poll::Ready(Some(bytes)) => {
this.read_buf = bytes;
this.read_pos = 0;
}
Poll::Ready(None) => {
this.eof = true;
}
Poll::Pending => return Poll::Pending,
}
}
}
}
impl AsyncWrite for WsByteStream {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
let this = self.get_mut();
if !this.write_open {
return Poll::Ready(Err(io::Error::new(
io::ErrorKind::BrokenPipe,
"ws stream shut down",
)));
}
if let Some(err) = this.poll_write_error() {
return Poll::Ready(Err(err));
}
let Some(write_tx) = this.write_tx_ref() else {
return Poll::Ready(Err(io::Error::new(
io::ErrorKind::BrokenPipe,
"ws stream shut down",
)));
};
if buf.len() > PENDING_BUFFER_CAP {
return Poll::Ready(Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"write of {} bytes exceeds the pending buffer cap ({})",
buf.len(),
PENDING_BUFFER_CAP
),
)));
}
match write_tx.poll_ready(cx) {
Poll::Ready(Ok(())) => match write_tx.try_send(WriteMsg::Bytes(buf.to_vec())) {
Ok(()) => Poll::Ready(Ok(buf.len())),
Err(_disconnected_or_full_race) => Poll::Ready(Err(io::Error::new(
io::ErrorKind::BrokenPipe,
"ws writer closed",
))),
},
Poll::Ready(Err(_send_error)) => Poll::Ready(Err(io::Error::new(
io::ErrorKind::BrokenPipe,
"ws writer closed",
))),
Poll::Pending => Poll::Pending,
}
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
let this = self.get_mut();
if this.write_open {
this.write_open = false;
drop(this.write_tx.take());
}
Poll::Ready(Ok(()))
}
}
impl WsByteStream {
fn poll_write_error(&mut self) -> Option<io::Error> {
if self.write_failed {
return Some(io::Error::new(
io::ErrorKind::InvalidData,
"ws write pump failed (protocol violation)",
));
}
match self.write_error.try_recv() {
Ok(reason) => {
self.write_failed = true;
self.write_open = false;
Some(io::Error::new(
io::ErrorKind::InvalidData,
format!("ws write pump failed: {reason}"),
))
}
Err(oneshot::error::TryRecvError::Closed) => {
self.write_failed = true;
self.write_open = false;
Some(io::Error::new(
io::ErrorKind::BrokenPipe,
"ws write pump ended",
))
}
Err(oneshot::error::TryRecvError::Empty) => None,
}
}
fn write_tx_ref(&mut self) -> Option<&mut futures_mpsc::Sender<WriteMsg>> {
self.write_tx.as_mut()
}
}
#[cfg(any(test, feature = "wss"))]
struct TungsteniteFraming;
#[cfg(any(test, feature = "wss"))]
impl WsFraming for TungsteniteFraming {
type Bytes = bytes::Bytes;
type Msg = tokio_tungstenite::tungstenite::Message;
fn binary(msg: &Self::Msg) -> Option<&Self::Bytes> {
match msg {
tokio_tungstenite::tungstenite::Message::Binary(b) => Some(b),
_ => None,
}
}
fn is_text(msg: &Self::Msg) -> bool {
matches!(msg, tokio_tungstenite::tungstenite::Message::Text(_))
}
fn is_close(msg: &Self::Msg) -> bool {
matches!(msg, tokio_tungstenite::tungstenite::Message::Close(_))
}
fn close_message(code: u16, reason: &'static str) -> Self::Msg {
tokio_tungstenite::tungstenite::Message::Close(Some(
tokio_tungstenite::tungstenite::protocol::CloseFrame {
code: code.into(),
reason: reason.into(),
},
))
}
fn binary_message(bytes: Vec<u8>) -> Self::Msg {
tokio_tungstenite::tungstenite::Message::Binary(bytes.into())
}
}
#[cfg(any(test, feature = "wss"))]
pub fn split_tungstenite_to_bytes<S>(
socket: tokio_tungstenite::WebSocketStream<S>,
) -> (WsByteStream, WsPumps)
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
{
split_tungstenite_to_bytes_idle(socket, Some(DEFAULT_WS_IDLE_TIMEOUT))
}
#[cfg(any(test, feature = "wss"))]
pub fn split_tungstenite_to_bytes_idle<S>(
socket: tokio_tungstenite::WebSocketStream<S>,
idle_timeout: Option<Duration>,
) -> (WsByteStream, WsPumps)
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
{
let (ws_sink, ws_stream) = socket.split();
let (read_tx, read_rx) = mpsc::channel::<Vec<u8>>(READ_SLOTS);
let (write_tx, write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
let (write_error_slot, write_error_rx) = make_write_error_slot();
let read_eof = tokio::sync::watch::channel(false).0;
let write_tx_for_read = write_tx.clone();
let write_error_slot_for_read = Arc::clone(&write_error_slot);
let read_eof_for_task = read_eof.clone();
let read_task = tokio::spawn(run_read_pump::<TungsteniteFraming, _, _>(
ws_stream,
read_tx,
write_tx_for_read,
write_error_slot_for_read,
idle_timeout,
move || {
let _ = read_eof_for_task.send(true);
},
));
let write_task = tokio::spawn(run_write_pump::<TungsteniteFraming, _>(
ws_sink,
write_rx,
write_error_slot,
Some(crate::websocket::DEFAULT_WS_WRITE_TIMEOUT),
));
(
WsByteStream {
read_rx,
read_buf: Vec::new(),
read_pos: 0,
eof: false,
write_tx: Some(write_tx),
write_open: true,
write_error: write_error_rx,
write_failed: false,
},
WsPumps {
read_task,
write_task,
#[cfg(any(test, feature = "wss"))]
read_eof,
},
)
}
#[cfg(test)]
pub(crate) fn split_tungstenite_to_bytes_idle_with_write<S>(
socket: tokio_tungstenite::WebSocketStream<S>,
idle_timeout: Option<Duration>,
write_timeout: Option<Duration>,
) -> (WsByteStream, WsPumps)
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
{
let (ws_sink, ws_stream) = socket.split();
let (read_tx, read_rx) = mpsc::channel::<Vec<u8>>(READ_SLOTS);
let (write_tx, write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
let (write_error_slot, write_error_rx) = make_write_error_slot();
let read_eof = tokio::sync::watch::channel(false).0;
let write_tx_for_read = write_tx.clone();
let write_error_slot_for_read = Arc::clone(&write_error_slot);
let read_eof_for_task = read_eof.clone();
let read_task = tokio::spawn(run_read_pump::<TungsteniteFraming, _, _>(
ws_stream,
read_tx,
write_tx_for_read,
write_error_slot_for_read,
idle_timeout,
move || {
let _ = read_eof_for_task.send(true);
},
));
let write_task = tokio::spawn(run_write_pump::<TungsteniteFraming, _>(
ws_sink,
write_rx,
write_error_slot,
write_timeout,
));
(
WsByteStream {
read_rx,
read_buf: Vec::new(),
read_pos: 0,
eof: false,
write_tx: Some(write_tx),
write_open: true,
write_error: write_error_rx,
write_failed: false,
},
WsPumps {
read_task,
write_task,
read_eof,
},
)
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
#[tokio::test]
async fn tungstenite_write_stall_is_evicted_within_the_write_timeout() {
let (client_io, _server_io) = tokio::io::duplex(64);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let knob = std::time::Duration::from_millis(150);
let (mut stream, _pumps) = split_tungstenite_to_bytes_idle_with_write(ws, None, Some(knob));
let mut chunk = vec![0u8; 8];
chunk[4..8].copy_from_slice(&64u32.to_be_bytes());
chunk.extend_from_slice(&[0u8; 64]);
stream.write_all(&chunk).await.expect("chunk written");
stream.flush().await.expect("flush");
let started = tokio::time::Instant::now();
let err = loop {
match stream.write_all(&[0u8; 8]).await {
Err(e) => break e,
Ok(_) => assert!(
started.elapsed() < std::time::Duration::from_secs(5),
"stream never failed while the sink stayed clogged"
),
}
};
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
assert!(
err.to_string().contains("write timeout"),
"error names the violation: {err}"
);
assert!(
started.elapsed() < std::time::Duration::from_secs(5),
"eviction must happen within the knob (+ slack), not hang"
);
}
pub(crate) fn oversize_header() -> Vec<u8> {
let mut header = vec![0u8; 8];
header[4..8].copy_from_slice(&(MAX_CHUNK_LEN + 1).to_be_bytes());
header
}
#[tokio::test]
async fn tungstenite_write_side_rejects_oversized_chunk_header() {
let (client_io, _server_io) = tokio::io::duplex(64);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let (mut stream, _pumps) = split_tungstenite_to_bytes(ws);
stream
.write_all(&oversize_header())
.await
.expect("write accepted");
stream.flush().await.expect("flush");
let err: io::Error;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5);
loop {
match stream.write_all(&[0u8; 8]).await {
Err(e) => {
err = e;
break;
}
Ok(_) => {
assert!(
tokio::time::Instant::now() < deadline,
"stream never failed after the oversized header"
);
}
}
}
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
assert!(
err.to_string().contains("MAX_CHUNK_LEN"),
"error names the violation: {err}"
);
}
#[tokio::test]
async fn tungstenite_write_at_the_cap_is_accepted() {
let (client_io, mut server_io) = tokio::io::duplex(64);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let (mut stream, _pumps) = split_tungstenite_to_bytes(ws);
let mut chunk = vec![0u8; PENDING_BUFFER_CAP];
chunk[4..8].copy_from_slice(&(PENDING_BUFFER_CAP as u32 - 8).to_be_bytes());
stream.write_all(&chunk).await.expect("cap write accepted");
stream.flush().await.expect("flush");
tokio::time::timeout(std::time::Duration::from_secs(30), async {
let mut buf = [0u8; 2];
server_io
.read_exact(&mut buf)
.await
.expect("cap-sized chunk read");
})
.await
.expect("pump emitted the cap-sized chunk");
}
#[tokio::test]
async fn tungstenite_write_one_byte_over_the_cap_is_rejected_with_invalid_data() {
let (client_io, mut server_io) = tokio::io::duplex(64);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let (mut stream, _pumps) = split_tungstenite_to_bytes(ws);
let err = stream
.write_all(&vec![0u8; PENDING_BUFFER_CAP + 1])
.await
.expect_err("the over-cap write must be rejected");
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
assert!(
err.to_string().contains("pending buffer cap"),
"error names the cap: {err}"
);
assert!(
err.to_string().contains(&PENDING_BUFFER_CAP.to_string()),
"error carries the cap value: {err}"
);
let mut chunk = vec![0u8; PENDING_BUFFER_CAP];
chunk[4..8].copy_from_slice(&(PENDING_BUFFER_CAP as u32 - 8).to_be_bytes());
stream
.write_all(&chunk)
.await
.expect("at-cap write after the rejection is accepted");
stream.flush().await.expect("flush");
tokio::time::timeout(std::time::Duration::from_secs(30), async {
let mut buf = [0u8; 2];
server_io.read_exact(&mut buf).await.expect("chunk read")
})
.await
.expect("pump still emits after a rejected over-cap write");
}
#[tokio::test]
async fn tungstenite_write_side_still_emits_well_framed_chunk() {
let (client_io, mut server_io) = tokio::io::duplex(64);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let (mut stream, _pumps) = split_tungstenite_to_bytes(ws);
let mut chunk = vec![0u8; 8];
chunk[4..8].copy_from_slice(&4u32.to_be_bytes());
chunk.extend_from_slice(b"payload");
stream.write_all(&chunk).await.expect("write accepted");
stream.flush().await.expect("flush");
tokio::time::timeout(std::time::Duration::from_secs(5), async {
let mut buf = [0u8; 2];
server_io
.read_exact(&mut buf)
.await
.expect("mask scan read");
})
.await
.expect("pump emitted the framed chunk");
}
#[tokio::test]
async fn tungstenite_read_side_forwards_binary_message_bytes() {
let (client_io, server_io) = tokio::io::duplex(1 << 16);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let (mut stream, _pumps) = split_tungstenite_to_bytes(ws);
let write_side = tokio::spawn(async move {
let peer = tokio_tungstenite::WebSocketStream::from_raw_socket(
server_io,
tokio_tungstenite::tungstenite::protocol::Role::Server,
None,
)
.await;
let (mut sink, _reader) = peer.split();
use futures::SinkExt;
let _ = sink
.send(tokio_tungstenite::tungstenite::Message::Binary(
b"from-peer".to_vec().into(),
))
.await;
});
write_side.await.expect("peer write task completes");
use tokio::io::AsyncReadExt as _;
let mut buf = [0u8; 9];
tokio::time::timeout(
std::time::Duration::from_secs(5),
stream.read_exact(&mut buf),
)
.await
.expect("read within deadline")
.expect("read");
assert_eq!(&buf, b"from-peer");
}
#[tokio::test]
async fn read_pump_breaks_when_the_byte_stream_side_is_dropped() {
let (client_io, server_io) = tokio::io::duplex(1 << 16);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let (_stream, pumps) = split_tungstenite_to_bytes(ws);
drop(_stream);
let peer = tokio::spawn(async move {
let peer = tokio_tungstenite::WebSocketStream::from_raw_socket(
server_io,
tokio_tungstenite::tungstenite::protocol::Role::Server,
None,
)
.await;
let (mut sink, mut reader) = peer.split();
use futures::{SinkExt, StreamExt};
sink.send(tokio_tungstenite::tungstenite::Message::Binary(
b"after-drop".to_vec().into(),
))
.await
.expect("peer send");
let _ = tokio::time::timeout(std::time::Duration::from_secs(10), reader.next()).await;
});
let ended = tokio::time::timeout(std::time::Duration::from_secs(5), pumps.read_task)
.await
.expect("the pump must end on the demux-gone break, not park");
assert!(
ended.is_ok(),
"the pump ends by itself (the break), not by abort: {ended:?}"
);
peer.abort();
}
#[tokio::test]
async fn tungstenite_read_eof_signal_fires_and_is_retained() {
let (client_io, server_io) = tokio::io::duplex(1 << 16);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let (_stream, pumps) = split_tungstenite_to_bytes(ws);
let mut eof_rx = pumps.read_eof();
drop(server_io);
let changed = tokio::time::timeout(std::time::Duration::from_secs(5), eof_rx.changed())
.await
.expect("EOF observed within deadline");
assert!(changed.is_ok() || *eof_rx.borrow(), "EOF flagged");
assert!(*eof_rx.borrow(), "EOF signal retained as true");
let mut late_rx = pumps.read_eof();
assert!(
*late_rx.borrow_and_update(),
"a receiver taken after EOF still observes the signal"
);
}
#[tokio::test]
async fn tungstenite_read_side_text_message_requests_protocol_close() {
let (client_io, server_io) = tokio::io::duplex(1 << 16);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let (_stream, pumps) = split_tungstenite_to_bytes(ws);
let write_side = tokio::spawn(async move {
let peer = tokio_tungstenite::WebSocketStream::from_raw_socket(
server_io,
tokio_tungstenite::tungstenite::protocol::Role::Server,
None,
)
.await;
let (mut sink, mut reader) = peer.split();
use futures::{SinkExt, StreamExt};
let _ = sink
.send(tokio_tungstenite::tungstenite::Message::Text(
"text frame".into(),
))
.await;
let _ = reader.next().await;
});
let mut eof_rx = pumps.read_eof();
let changed = tokio::time::timeout(std::time::Duration::from_secs(5), eof_rx.changed())
.await
.expect("read pump ends after the text frame (close requested)");
let _ = changed;
write_side.await.expect("peer task completes");
}
#[tokio::test]
async fn tungstenite_shutdown_closes_the_write_sink_toward_the_peer() {
let (client_io, server_io) = tokio::io::duplex(1 << 16);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let (mut stream, _pumps) = split_tungstenite_to_bytes(ws);
let peer = tokio_tungstenite::WebSocketStream::from_raw_socket(
server_io,
tokio_tungstenite::tungstenite::protocol::Role::Server,
None,
)
.await;
let (mut peer_sink, mut peer_stream) = peer.split();
use futures::{SinkExt, StreamExt};
let mut chunk = vec![0u8; 8];
chunk[4..8].copy_from_slice(&4u32.to_be_bytes());
chunk.extend_from_slice(b"pay!");
stream.write_all(&chunk).await.expect("write accepted");
stream.flush().await.expect("flush");
tokio::time::timeout(std::time::Duration::from_secs(5), async {
loop {
match peer_stream.next().await {
Some(Ok(tokio_tungstenite::tungstenite::Message::Binary(b))) => {
assert_eq!(&*b, &chunk, "queued bytes drain before the close");
return;
}
Some(Ok(_)) => continue,
other => panic!("peer stream ended early: {other:?}"),
}
}
})
.await
.expect("chunk emitted within deadline");
stream.shutdown().await.expect("shutdown runs");
peer_sink
.send(tokio_tungstenite::tungstenite::Message::Close(None))
.await
.expect("peer sends close");
let saw_close = tokio::time::timeout(std::time::Duration::from_secs(5), async move {
loop {
match peer_stream.next().await {
None => return false,
Some(Ok(tokio_tungstenite::tungstenite::Message::Close(_))) => return true,
Some(Ok(_)) => continue,
Some(Err(_)) => return false,
}
}
})
.await
.expect("peer observes the close reply after our shutdown");
assert!(saw_close, "the shutdown path emitted a WS Close frame");
}
#[tokio::test]
async fn tungstenite_text_close_carries_protocol_reason() {
let (client_io, server_io) = tokio::io::duplex(1 << 16);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let (_stream, pumps) = split_tungstenite_to_bytes(ws);
let peer = tokio_tungstenite::WebSocketStream::from_raw_socket(
server_io,
tokio_tungstenite::tungstenite::protocol::Role::Server,
None,
)
.await;
let (mut peer_sink, mut peer_stream) = peer.split();
use futures::{SinkExt, StreamExt};
let mut eof_rx = pumps.read_eof();
peer_sink
.send(tokio_tungstenite::tungstenite::Message::Text(
"text frame".into(),
))
.await
.expect("text write accepted");
let (close, eof_fired) = tokio::time::timeout(std::time::Duration::from_secs(5), async {
let close: Option<(u16, String)> = loop {
match peer_stream.next().await {
Some(Ok(tokio_tungstenite::tungstenite::Message::Close(cf))) => {
break cf.map(|f| (u16::from(f.code), f.reason.to_string()))
}
Some(Ok(_)) => continue,
Some(Err(_)) | None => break None,
}
};
let _ = eof_rx.changed().await;
(close, *eof_rx.borrow())
})
.await
.expect("read pump ends after the text frame (close requested)");
let (code, reason) = close.expect("close frame with code + reason");
assert_eq!(code, WS_PROTOCOL_ERROR, "text frame closed with 1002");
assert_eq!(
reason, "text messages not supported",
"close reason names the text cause, got {reason:?}"
);
assert!(eof_fired, "EOF signal fired for the from_wss machinery");
}
#[tokio::test]
async fn idle_read_timeout_closes_a_stalled_connection_with_goingaway() {
let (client_io, server_io) = tokio::io::duplex(1 << 16);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let knob = std::time::Duration::from_millis(150);
let (_stream, pumps) = split_tungstenite_to_bytes_idle(ws, Some(knob));
let peer = tokio_tungstenite::WebSocketStream::from_raw_socket(
server_io,
tokio_tungstenite::tungstenite::protocol::Role::Server,
None,
)
.await;
let (_peer_sink, mut peer_stream) = peer.split();
use futures::StreamExt;
let mut eof_rx = pumps.read_eof();
let (close_seen, eof_fired) =
tokio::time::timeout(std::time::Duration::from_secs(5), async {
let close: Option<(u16, String)> = loop {
match peer_stream.next().await {
Some(Ok(tokio_tungstenite::tungstenite::Message::Close(cf))) => {
break cf.map(|f| (u16::from(f.code), f.reason.to_string()))
}
Some(Ok(_)) => continue,
Some(Err(_)) | None => break None,
}
};
let _ = eof_rx.changed().await;
(close, *eof_rx.borrow())
})
.await
.expect("stalled connection must be torn down within the deadline");
let (code, reason) = close_seen.expect("close frame with code + reason");
assert_eq!(
code, WS_GOING_AWAY,
"peer received the 1001 GoingAway close"
);
assert!(
reason.contains("no inbound chunk progress"),
"close reason names the idle-read cause, got {reason:?}"
);
assert!(eof_fired, "EOF signal fired for the from_wss machinery");
}
#[tokio::test]
async fn idle_read_timeout_bounds_a_forever_dribble_with_no_chunk_progress() {
let (client_io, server_io) = tokio::io::duplex(1 << 16);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let knob = std::time::Duration::from_millis(200);
let (_stream, pumps) = split_tungstenite_to_bytes_idle(ws, Some(knob));
let peer = tokio_tungstenite::WebSocketStream::from_raw_socket(
server_io,
tokio_tungstenite::tungstenite::protocol::Role::Server,
None,
)
.await;
let (mut peer_sink, mut peer_stream) = peer.split();
use futures::{SinkExt, StreamExt};
let mut eof_rx = pumps.read_eof();
let dribble = tokio::spawn(async move {
let mut header = vec![0u8; 8];
header[4..8].copy_from_slice(&64u32.to_be_bytes());
peer_sink
.send(tokio_tungstenite::tungstenite::Message::Binary(
header.into(),
))
.await
.expect("header write accepted");
for byte in 0u8..64 {
tokio::time::sleep(knob / 4).await;
if peer_sink
.send(tokio_tungstenite::tungstenite::Message::Binary(
vec![byte].into(),
))
.await
.is_err()
{
return;
}
}
loop {
tokio::time::sleep(knob).await;
if peer_sink
.send(tokio_tungstenite::tungstenite::Message::Binary(
vec![0u8].into(),
))
.await
.is_err()
{
return;
}
}
});
let outcome = tokio::time::timeout(std::time::Duration::from_secs(10), async {
let close: Option<(u16, String)> = loop {
match peer_stream.next().await {
Some(Ok(tokio_tungstenite::tungstenite::Message::Close(cf))) => {
break cf.map(|f| (u16::from(f.code), f.reason.to_string()))
}
Some(Ok(_)) => continue,
Some(Err(_)) | None => break None,
}
};
let _ = eof_rx.changed().await;
close
})
.await
.expect("dribble must hit the progress deadline within 10 s");
dribble.abort();
let (code, reason) = outcome.expect("close frame with code + reason");
assert_eq!(
code, WS_GOING_AWAY,
"the forever-dribble is evicted with 1001 despite arriving messages"
);
assert!(
reason.contains("no inbound chunk progress"),
"close reason names the idle-read cause, got {reason:?}"
);
}
#[tokio::test]
async fn idle_read_timeout_survives_slow_productive_chunks() {
let (client_io, server_io) = tokio::io::duplex(1 << 16);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let knob = std::time::Duration::from_millis(200);
let (mut stream, pumps) = split_tungstenite_to_bytes_idle(ws, Some(knob));
let mut evict_rx = pumps.read_eof();
let peer = tokio_tungstenite::WebSocketStream::from_raw_socket(
server_io,
tokio_tungstenite::tungstenite::protocol::Role::Server,
None,
)
.await;
let (mut peer_sink, _peer_stream) = peer.split();
use futures::{SinkExt, StreamExt};
for round in 0u8..6u8 {
tokio::time::sleep(knob / 2).await;
let mut chunk = vec![0u8; 12];
chunk[0..4].copy_from_slice(&0u32.to_be_bytes());
chunk[4..8].copy_from_slice(&4u32.to_be_bytes());
chunk[8..12].copy_from_slice(&[round; 4]);
peer_sink
.send(tokio_tungstenite::tungstenite::Message::Binary(
chunk.clone().into(),
))
.await
.expect("chunk write accepted");
let mut buf = vec![0u8; 12];
tokio::time::timeout(knob, stream.read_exact(&mut buf))
.await
.expect("connection alive between productive chunks")
.expect("read");
assert_eq!(buf, chunk, "chunk bytes surfaced in order");
}
let evicted_early =
tokio::time::timeout(std::time::Duration::from_millis(0), evict_rx.changed()).await;
assert!(
evicted_early.is_err(),
"no eviction while chunks keep landing inside the window"
);
let mut header = vec![0u8; 8];
header[4..8].copy_from_slice(&0u32.to_be_bytes());
peer_sink
.send(tokio_tungstenite::tungstenite::Message::Binary(
header.into(),
))
.await
.expect("EOF chunk write accepted");
tokio::time::timeout(std::time::Duration::from_secs(5), async {
loop {
match stream.read(&mut [0u8; 1]).await {
Ok(0) => return,
Ok(_) => continue,
Err(e) => panic!("read failed: {e}"),
}
}
})
.await
.expect("demux EOF observed after the final chunk");
}
#[tokio::test]
async fn idle_read_timeout_disabled_when_none() {
let (client_io, server_io) = tokio::io::duplex(1 << 16);
let ws = tokio_tungstenite::WebSocketStream::from_raw_socket(
client_io,
tokio_tungstenite::tungstenite::protocol::Role::Client,
None,
)
.await;
let (_stream, pumps) = split_tungstenite_to_bytes_idle(ws, None);
let peer = tokio_tungstenite::WebSocketStream::from_raw_socket(
server_io,
tokio_tungstenite::tungstenite::protocol::Role::Server,
None,
)
.await;
let (mut peer_sink, _peer_stream) = peer.split();
use futures::SinkExt;
let eof_rx = pumps.read_eof();
for _ in 0..4 {
tokio::time::sleep(std::time::Duration::from_millis(150)).await;
peer_sink
.send(tokio_tungstenite::tungstenite::Message::Binary(
vec![0u8].into(),
))
.await
.expect("dribble write accepted");
}
assert!(
!*eof_rx.borrow(),
"no idle timer ran with the knob disabled: arriving messages evict nothing"
);
peer_sink
.send(tokio_tungstenite::tungstenite::Message::Close(None))
.await
.expect("peer close accepted");
}
}
#[cfg(all(test, feature = "server"))]
mod axum_framing_tests {
use super::*;
use futures::channel::mpsc as fut_mpsc;
struct AxumFakeSocket {
sink_tx: futures_mpsc::Sender<AxumMessage>,
stream_rx:
std::pin::Pin<Box<dyn futures::Stream<Item = Result<AxumMessage, AxumMessage>> + Send>>,
}
impl futures::Stream for AxumFakeSocket {
type Item = Result<AxumMessage, AxumMessage>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.stream_rx.as_mut().poll_next(cx)
}
}
impl futures::Sink<AxumMessage> for AxumFakeSocket {
type Error = ();
fn poll_ready(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.sink_tx).poll_ready(cx).map_err(|_| ())
}
fn start_send(mut self: Pin<&mut Self>, item: AxumMessage) -> Result<(), Self::Error> {
self.sink_tx.start_send(item).map_err(|_| ())
}
fn poll_flush(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
futures::Sink::poll_flush(Pin::new(&mut self.sink_tx), cx).map_err(|_| ())
}
fn poll_close(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
futures::Sink::poll_close(Pin::new(&mut self.sink_tx), cx).map_err(|_| ())
}
}
type AxumPairParts = (
futures::stream::SplitSink<AxumFakeSocket, AxumMessage>,
futures::stream::SplitStream<AxumFakeSocket>,
futures_mpsc::Receiver<AxumMessage>,
fut_mpsc::Sender<Result<AxumMessage, AxumMessage>>,
);
fn axum_pair() -> AxumPairParts {
let (outbound_tx, outbound_rx) = fut_mpsc::channel::<AxumMessage>(4);
let (inbound_tx, inbound_rx) = fut_mpsc::channel::<Result<AxumMessage, AxumMessage>>(4);
let socket = AxumFakeSocket {
sink_tx: outbound_tx,
stream_rx: Box::pin(inbound_rx),
};
let (sink, stream) = socket.split();
(sink, stream, outbound_rx, inbound_tx)
}
#[tokio::test]
async fn axum_framing_text_message_requests_protocol_close_with_reason() {
let (_sink, stream, _outbound_rx, mut inbound_tx) = axum_pair();
let (write_tx_for_read, mut write_rx_for_read) = fut_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
let read_task = tokio::spawn(run_read_pump::<AxumFraming, _, _>(
stream,
mpsc::channel::<Vec<u8>>(READ_SLOTS).0,
write_tx_for_read,
make_write_error_slot().0,
None,
|| {},
));
inbound_tx
.send(Ok(AxumMessage::Text("text frame".into())))
.await
.expect("inbound message accepted");
let close = tokio::time::timeout(std::time::Duration::from_secs(5), async {
while let Some(msg) = write_rx_for_read.next().await {
if let WriteMsg::CloseWith(code, reason) = msg {
return <AxumFraming as WsFraming>::close_message(code, reason);
}
}
AxumMessage::Text("queue closed".into())
})
.await
.expect("close requested");
let AxumMessage::Close(Some(frame)) = close else {
panic!("expected a close frame, got {close:?}");
};
assert_eq!(frame.code, WS_PROTOCOL_ERROR, "text closes with 1002");
assert_eq!(
frame.reason, "text messages not supported",
"close reason names the text cause"
);
read_task.abort();
}
#[tokio::test]
async fn axum_framing_cap_trip_fails_the_pump_with_the_internal_close() {
let (sink, _stream, mut outbound_rx, _inbound_tx) = axum_pair();
let (slot, _rx) = make_write_error_slot();
let write_task = tokio::spawn(run_write_pump::<AxumFraming, _>(
sink,
{
let (mut tx, rx) = fut_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
tx.send(WriteMsg::Bytes(tests::oversize_header()))
.await
.expect("queue write accepted");
rx
},
slot,
None,
));
let close = tokio::time::timeout(std::time::Duration::from_secs(5), async {
loop {
match outbound_rx.next().await {
Some(m @ AxumMessage::Close(_)) => return m,
Some(_) => continue,
None => return AxumMessage::Text("stream ended".into()),
}
}
})
.await
.expect("close observed");
let AxumMessage::Close(Some(frame)) = close else {
panic!("expected a close frame, got {close:?}");
};
assert_eq!(frame.code, WS_INTERNAL_ERROR, "cap trip closes with 1011");
assert!(
frame.reason.contains("MAX_CHUNK_LEN"),
"close reason names the violation: {}",
frame.reason
);
write_task.abort();
}
}