refactor(websocket): one generic pump implementation (WS-11, COV-03)
This commit is contained in:
+432
-223
@@ -38,7 +38,9 @@
|
||||
//! bytes drain (the mux's EOF sentinels ride the same queue).
|
||||
//!
|
||||
//! Shared with the `from_wss` consumer path (ADR-070): one
|
||||
//! implementation, both directions.
|
||||
//! implementation, both directions ([`WsFraming`] + the generic
|
||||
//! [`run_read_pump`] / [`run_write_pump`] instantiated once per socket
|
||||
//! flavor — WS-11/COV-03).
|
||||
|
||||
use std::{
|
||||
io,
|
||||
@@ -137,6 +139,201 @@ fn send_write_error(slot: &WriteErrorSlot, reason: &'static str) {
|
||||
}
|
||||
}
|
||||
|
||||
/// WS-protocol adapter for one socket flavor: the only places the axum
|
||||
/// and tungstenite message types differ. The generic pump
|
||||
/// ([`run_read_pump`] / [`run_write_pump`]) is written against this
|
||||
/// trait so both paths share one implementation (WS-11, "one
|
||||
/// implementation, both directions").
|
||||
trait WsFraming: Sized {
|
||||
/// The WS library's binary-message byte type.
|
||||
type Bytes: AsRef<[u8]>;
|
||||
/// The sink/stream message type.
|
||||
type Msg;
|
||||
|
||||
/// `Some(bytes)` for a binary message (the byte stream's carrier),
|
||||
/// `None` for anything else.
|
||||
fn binary(msg: &Self::Msg) -> Option<&Self::Bytes>;
|
||||
|
||||
/// `true` for a text message (a protocol error: close 1002).
|
||||
fn is_text(msg: &Self::Msg) -> bool;
|
||||
|
||||
/// `true` for a Close message from the peer.
|
||||
fn is_close(msg: &Self::Msg) -> bool;
|
||||
|
||||
/// A Close message with the given code and reason.
|
||||
fn close_message(code: u16, reason: &'static str) -> Self::Msg;
|
||||
|
||||
/// A binary message carrying `bytes` (chunk pieces up to
|
||||
/// `WS_MESSAGE_CAP`).
|
||||
fn binary_message(bytes: Vec<u8>) -> Self::Msg;
|
||||
}
|
||||
|
||||
/// The read stream's item for a flavor: `Result<Msg, Error>` (both
|
||||
/// flavors surface read failures this way; the error type is flavor
|
||||
/// specific and only ever inspected as "failed").
|
||||
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(|_| ())
|
||||
}
|
||||
}
|
||||
|
||||
/// The read pump: forwards binary-message bytes into `read_tx`, turns
|
||||
/// text into a 1002 close request, and treats peer close / read error
|
||||
/// as the connection end. `on_end` runs after the loop (the `wss`
|
||||
/// feature fires the lossless EOF watch signal through it — the
|
||||
/// from_wss drop monitor's input; its `Sender` semantics are preserved
|
||||
/// EXACTLY: single `send` after the read loop terminates).
|
||||
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,
|
||||
on_end: F,
|
||||
) where
|
||||
S: futures::Stream + Unpin,
|
||||
M: WsFraming,
|
||||
S::Item: IntoWsResult<M>,
|
||||
F: FnOnce(),
|
||||
{
|
||||
while let Some(msg) = ws_stream.next().await {
|
||||
match msg.into_ws_result() {
|
||||
Ok(m) => {
|
||||
if let Some(b) = M::binary(&m) {
|
||||
if read_tx.send(b.as_ref().to_vec()).await.is_err() {
|
||||
break;
|
||||
}
|
||||
} else if M::is_text(&m) {
|
||||
let _ = write_tx_for_read
|
||||
.clone()
|
||||
.send(WriteMsg::CloseWith(WS_PROTOCOL_ERROR))
|
||||
.await;
|
||||
break;
|
||||
} else if M::is_close(&m) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(()) => {
|
||||
send_write_error(
|
||||
&write_error_slot_for_read,
|
||||
"inbound frame rejected (size cap)",
|
||||
);
|
||||
let _ = write_tx_for_read
|
||||
.clone()
|
||||
.send(WriteMsg::CloseWith(WS_INTERNAL_ERROR))
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
on_end();
|
||||
}
|
||||
|
||||
/// Emit the pump's fatal close + error, then end the write task
|
||||
/// (WS-04/HY-09, WS-05: fail loudly instead of accumulating/hanging).
|
||||
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;
|
||||
}
|
||||
|
||||
/// The write pump: drains the queued byte spans, parses the pending
|
||||
/// bytes for complete chunks (8-byte header, length validated against
|
||||
/// `MAX_CHUNK_LEN` — WS-04/HY-09), byte-caps the accumulator
|
||||
/// (`PENDING_BUFFER_CAP` — WS-05), and emits one WS binary message per
|
||||
/// `WS_MESSAGE_CAP` piece. Ends with a WS Close after the queue
|
||||
/// drains; `CloseWith` requests (`WriteMsg::CloseWith`) map to a close
|
||||
/// frame with that code.
|
||||
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,
|
||||
) 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) => {
|
||||
let _ = ws_sink
|
||||
.send(<M as WsFraming>::close_message(
|
||||
code,
|
||||
"text messages not supported",
|
||||
))
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
WriteMsg::Bytes(b) => {
|
||||
if b.len() > PENDING_BUFFER_CAP {
|
||||
fail_write_pump::<M, _>(
|
||||
&mut ws_sink,
|
||||
&write_error_slot,
|
||||
"write pending buffer exceeded cap",
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
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 ws_sink
|
||||
.send(<M as WsFraming>::binary_message(piece.to_vec()))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
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;
|
||||
}
|
||||
|
||||
/// The `AsyncRead + AsyncWrite` view of a `WebSocket` handed to
|
||||
/// alkcall's channels machinery (demux/mux).
|
||||
pub struct WsByteStream {
|
||||
@@ -156,7 +353,7 @@ pub struct WsByteStream {
|
||||
pub struct WsPumps {
|
||||
read_task: tokio::task::JoinHandle<()>,
|
||||
write_task: tokio::task::JoinHandle<()>,
|
||||
#[cfg(feature = "wss")]
|
||||
#[cfg(any(test, feature = "wss"))]
|
||||
read_eof: tokio::sync::watch::Sender<bool>,
|
||||
}
|
||||
|
||||
@@ -172,132 +369,80 @@ impl WsPumps {
|
||||
/// taken or the observer starts awaiting — is still observed, and
|
||||
/// may be observed repeatedly. Used by `from_wss`'s
|
||||
/// connection-drop monitor (ADR-070).
|
||||
#[cfg(feature = "wss")]
|
||||
#[cfg(any(test, feature = "wss"))]
|
||||
pub(crate) fn read_eof(&self) -> tokio::sync::watch::Receiver<bool> {
|
||||
self.read_eof.subscribe()
|
||||
}
|
||||
}
|
||||
|
||||
/// axum's message/CloseFrame types as a [`WsFraming`] flavor.
|
||||
struct AxumFraming;
|
||||
|
||||
impl WsFraming for AxumFraming {
|
||||
type Bytes = bytes::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())
|
||||
}
|
||||
}
|
||||
|
||||
/// Split a `WebSocket` into the byte stream + the pump tasks. The
|
||||
/// adapter is the single seam between axum's WS and alkcall's
|
||||
/// byte-oriented channels machinery; shared with `from_wss`.
|
||||
pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
|
||||
let (mut ws_sink, mut ws_stream) = socket.split();
|
||||
let (ws_sink, ws_stream) = socket.split();
|
||||
let (read_tx, read_rx) = mpsc::channel::<Vec<u8>>(READ_SLOTS);
|
||||
let (write_tx, mut write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
|
||||
let (write_tx, write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
|
||||
let (write_error_slot, write_error_rx) = make_write_error_slot();
|
||||
|
||||
#[cfg(feature = "wss")]
|
||||
#[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(feature = "wss")]
|
||||
#[cfg(any(test, feature = "wss"))]
|
||||
let read_eof_for_task = read_eof.clone();
|
||||
let read_task = tokio::spawn(async move {
|
||||
while let Some(msg) = ws_stream.next().await {
|
||||
match msg {
|
||||
Ok(AxumMessage::Binary(b)) => {
|
||||
if read_tx.send(b.to_vec()).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(AxumMessage::Text(_)) => {
|
||||
let _ = write_tx_for_read
|
||||
.clone()
|
||||
.send(WriteMsg::CloseWith(WS_PROTOCOL_ERROR))
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
Ok(AxumMessage::Close(_)) => break,
|
||||
Err(_) => {
|
||||
send_write_error(&write_error_slot_for_read, "inbound frame rejected (size cap)");
|
||||
let _ = write_tx_for_read
|
||||
.clone()
|
||||
.send(WriteMsg::CloseWith(WS_INTERNAL_ERROR))
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
Ok(_) => {}
|
||||
}
|
||||
}
|
||||
#[cfg(feature = "wss")]
|
||||
let read_task = tokio::spawn(run_read_pump::<AxumFraming, _, _>(
|
||||
ws_stream,
|
||||
read_tx,
|
||||
write_tx_for_read,
|
||||
write_error_slot_for_read,
|
||||
move || {
|
||||
#[cfg(any(test, feature = "wss"))]
|
||||
{
|
||||
let _ = read_eof_for_task.send(true);
|
||||
}
|
||||
});
|
||||
},
|
||||
));
|
||||
|
||||
let write_task = tokio::spawn(async move {
|
||||
let mut pending: Vec<u8> = Vec::new();
|
||||
while let Some(msg) = write_rx.next().await {
|
||||
match msg {
|
||||
WriteMsg::CloseWith(code) => {
|
||||
let _ = ws_sink
|
||||
.send(AxumMessage::Close(Some(CloseFrame {
|
||||
code,
|
||||
reason: "text messages not supported".into(),
|
||||
})))
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
WriteMsg::Bytes(b) => {
|
||||
if b.len() > PENDING_BUFFER_CAP {
|
||||
send_write_error(&write_error_slot, "write pending buffer exceeded cap");
|
||||
let _ = ws_sink
|
||||
.send(AxumMessage::Close(Some(CloseFrame {
|
||||
code: WS_INTERNAL_ERROR,
|
||||
reason: "write pending buffer exceeded cap".into(),
|
||||
})))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
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 {
|
||||
send_write_error(&write_error_slot, "chunk length exceeds MAX_CHUNK_LEN");
|
||||
let _ = ws_sink
|
||||
.send(AxumMessage::Close(Some(CloseFrame {
|
||||
code: WS_INTERNAL_ERROR,
|
||||
reason: "chunk length exceeds MAX_CHUNK_LEN".into(),
|
||||
})))
|
||||
.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 ws_sink
|
||||
.send(AxumMessage::Binary(piece.to_vec().into()))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
if pending.len() > PENDING_BUFFER_CAP {
|
||||
send_write_error(&write_error_slot, "write pending buffer exceeded cap");
|
||||
let _ = ws_sink
|
||||
.send(AxumMessage::Close(Some(CloseFrame {
|
||||
code: WS_INTERNAL_ERROR,
|
||||
reason: "write pending buffer exceeded cap".into(),
|
||||
})))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
let _ = ws_sink.close().await;
|
||||
});
|
||||
let write_task = tokio::spawn(run_write_pump::<AxumFraming, _>(
|
||||
ws_sink,
|
||||
write_rx,
|
||||
write_error_slot,
|
||||
));
|
||||
|
||||
(
|
||||
WsByteStream {
|
||||
@@ -313,7 +458,7 @@ pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
|
||||
WsPumps {
|
||||
read_task,
|
||||
write_task,
|
||||
#[cfg(feature = "wss")]
|
||||
#[cfg(any(test, feature = "wss"))]
|
||||
read_eof,
|
||||
},
|
||||
)
|
||||
@@ -435,12 +580,51 @@ impl WsByteStream {
|
||||
}
|
||||
}
|
||||
|
||||
/// tokio-tungstenite's message types as a [`WsFraming`] flavor.
|
||||
#[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())
|
||||
}
|
||||
}
|
||||
|
||||
/// Client-side split for a tokio-tungstenite `WebSocketStream`
|
||||
/// (`from_wss`, ADR-070): the same WS↔byte-stream seam as
|
||||
/// [`split_ws_to_bytes`], over the tungstenite socket instead of axum's
|
||||
/// server-side `WebSocket`. Message semantics are identical: binary
|
||||
/// messages carry the byte stream, text is a protocol error (close
|
||||
/// 1002), close → read EOF.
|
||||
/// 1002), close → read EOF. Both paths run the same generic pumps
|
||||
/// (WS-11); the socket flavor differs only in the [`WsFraming`] impl.
|
||||
#[cfg(any(test, feature = "wss"))]
|
||||
pub fn split_tungstenite_to_bytes<S>(
|
||||
socket: tokio_tungstenite::WebSocketStream<S>,
|
||||
@@ -448,9 +632,9 @@ pub fn split_tungstenite_to_bytes<S>(
|
||||
where
|
||||
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
|
||||
{
|
||||
let (mut ws_sink, mut ws_stream) = socket.split();
|
||||
let (ws_sink, ws_stream) = socket.split();
|
||||
let (read_tx, read_rx) = mpsc::channel::<Vec<u8>>(READ_SLOTS);
|
||||
let (write_tx, mut write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_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;
|
||||
@@ -458,117 +642,21 @@ where
|
||||
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(async move {
|
||||
while let Some(msg) = ws_stream.next().await {
|
||||
match msg {
|
||||
Ok(tokio_tungstenite::tungstenite::Message::Binary(b)) => {
|
||||
if read_tx.send(b.to_vec()).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(tokio_tungstenite::tungstenite::Message::Text(_)) => {
|
||||
let _ = write_tx_for_read
|
||||
.clone()
|
||||
.send(WriteMsg::CloseWith(WS_PROTOCOL_ERROR))
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
Ok(tokio_tungstenite::tungstenite::Message::Close(_)) => break,
|
||||
Err(_) => {
|
||||
send_write_error(&write_error_slot_for_read, "inbound frame rejected (size cap)");
|
||||
let _ = write_tx_for_read
|
||||
.clone()
|
||||
.send(WriteMsg::CloseWith(WS_INTERNAL_ERROR))
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
Ok(_) => {}
|
||||
}
|
||||
}
|
||||
let read_task = tokio::spawn(run_read_pump::<TungsteniteFraming, _, _>(
|
||||
ws_stream,
|
||||
read_tx,
|
||||
write_tx_for_read,
|
||||
write_error_slot_for_read,
|
||||
move || {
|
||||
let _ = read_eof_for_task.send(true);
|
||||
});
|
||||
},
|
||||
));
|
||||
|
||||
let write_task = tokio::spawn(async move {
|
||||
let mut pending: Vec<u8> = Vec::new();
|
||||
while let Some(msg) = write_rx.next().await {
|
||||
match msg {
|
||||
WriteMsg::CloseWith(code) => {
|
||||
let _ = ws_sink
|
||||
.send(tokio_tungstenite::tungstenite::Message::Close(Some(
|
||||
tokio_tungstenite::tungstenite::protocol::CloseFrame {
|
||||
code: code.into(),
|
||||
reason: "text messages not supported".into(),
|
||||
},
|
||||
)))
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
WriteMsg::Bytes(b) => {
|
||||
if b.len() > PENDING_BUFFER_CAP {
|
||||
send_write_error(&write_error_slot, "write pending buffer exceeded cap");
|
||||
let _ = ws_sink
|
||||
.send(tokio_tungstenite::tungstenite::Message::Close(Some(
|
||||
tokio_tungstenite::tungstenite::protocol::CloseFrame {
|
||||
code: WS_INTERNAL_ERROR.into(),
|
||||
reason: "write pending buffer exceeded cap".into(),
|
||||
},
|
||||
)))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
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 {
|
||||
send_write_error(&write_error_slot, "chunk length exceeds MAX_CHUNK_LEN");
|
||||
let _ = ws_sink
|
||||
.send(tokio_tungstenite::tungstenite::Message::Close(Some(
|
||||
tokio_tungstenite::tungstenite::protocol::CloseFrame {
|
||||
code: WS_INTERNAL_ERROR.into(),
|
||||
reason: "chunk length exceeds MAX_CHUNK_LEN".into(),
|
||||
},
|
||||
)))
|
||||
.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 ws_sink
|
||||
.send(tokio_tungstenite::tungstenite::Message::Binary(
|
||||
piece.to_vec().into(),
|
||||
))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
if pending.len() > PENDING_BUFFER_CAP {
|
||||
send_write_error(&write_error_slot, "write pending buffer exceeded cap");
|
||||
let _ = ws_sink
|
||||
.send(tokio_tungstenite::tungstenite::Message::Close(Some(
|
||||
tokio_tungstenite::tungstenite::protocol::CloseFrame {
|
||||
code: WS_INTERNAL_ERROR.into(),
|
||||
reason: "write pending buffer exceeded cap".into(),
|
||||
},
|
||||
)))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
let _ = ws_sink.close().await;
|
||||
});
|
||||
let write_task = tokio::spawn(run_write_pump::<TungsteniteFraming, _>(
|
||||
ws_sink,
|
||||
write_rx,
|
||||
write_error_slot,
|
||||
));
|
||||
|
||||
(
|
||||
WsByteStream {
|
||||
@@ -584,7 +672,7 @@ where
|
||||
WsPumps {
|
||||
read_task,
|
||||
write_task,
|
||||
#[cfg(feature = "wss")]
|
||||
#[cfg(any(test, feature = "wss"))]
|
||||
read_eof,
|
||||
},
|
||||
)
|
||||
@@ -615,7 +703,10 @@ mod tests {
|
||||
.await;
|
||||
let (mut stream, _pumps) = split_tungstenite_to_bytes(ws);
|
||||
|
||||
stream.write_all(&oversize_header()).await.expect("write accepted");
|
||||
stream
|
||||
.write_all(&oversize_header())
|
||||
.await
|
||||
.expect("write accepted");
|
||||
stream.flush().await.expect("flush");
|
||||
|
||||
let err: io::Error;
|
||||
@@ -708,9 +799,127 @@ mod tests {
|
||||
|
||||
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");
|
||||
server_io
|
||||
.read_exact(&mut buf)
|
||||
.await
|
||||
.expect("mask scan read");
|
||||
})
|
||||
.await
|
||||
.expect("pump emitted the framed chunk");
|
||||
}
|
||||
|
||||
/// Read-pump binary passthrough on the tungstenite path (COV-03):
|
||||
/// a peer binary byte-message surfaces on the `AsyncRead` half of
|
||||
/// the adapter over a raw duplex pair (no network), with a server
|
||||
/// role peer writing the frame.
|
||||
#[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");
|
||||
}
|
||||
|
||||
/// The lossless read-EOF signal on the tungstenite path (COV-03,
|
||||
/// the WS-11 constraint): the peer disappearing surfaces on
|
||||
/// `pumps.read_eof()` as `watch` = true — observable both before
|
||||
/// and after the fact, the from_wss drop-monitor contract.
|
||||
#[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"
|
||||
);
|
||||
}
|
||||
|
||||
/// Text frames over the tungstenite path are a protocol error: the
|
||||
/// pump bridges text → 1002 close on the socket (COV-03 read-pump
|
||||
/// coverage; pitcher on a raw duplex pair).
|
||||
#[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");
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user