use std::collections::VecDeque;
use std::marker::PhantomData;
use std::pin::Pin;
use std::task::ready;
use futures_util::stream::{FuturesOrdered, SplitSink};
use futures_util::{SinkExt, Stream, StreamExt};
use std::future::Future;
use std::task::{Context, Poll};
use tokio::sync::mpsc;
use tokio_tungstenite::tungstenite::http::{HeaderMap, HeaderName, HeaderValue};
use tokio_tungstenite::tungstenite::Message as WsMessage;
use tokio_tungstenite::MaybeTlsStream;
use tokio_tungstenite::{tungstenite::protocol::WebSocketConfig, WebSocketStream};
use chromiumoxide_cdp::cdp::browser_protocol::target::SessionId;
use chromiumoxide_types::{CallId, EventMessage, Message, MethodCall, MethodId};
use crate::error::CdpError;
use crate::error::Result;
type ConnectStream = MaybeTlsStream<tokio::net::TcpStream>;
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct Connection<T: EventMessage> {
pending_commands: VecDeque<MethodCall>,
ws: WebSocketStream<ConnectStream>,
next_id: usize,
needs_flush: bool,
_marker: PhantomData<T>,
}
lazy_static::lazy_static! {
static ref DISABLE_NAGLE: bool = match std::env::var("DISABLE_NAGLE") {
Ok(disable_nagle) => disable_nagle == "true",
_ => true
};
static ref WEBSOCKET_DEFAULTS: bool = match std::env::var("WEBSOCKET_DEFAULTS") {
Ok(d) => d == "true",
_ => false
};
}
pub const DEFAULT_CONNECTION_RETRIES: u32 = 4;
const INITIAL_BACKOFF_MS: u64 = 50;
pub(crate) const MAX_BACKOFF_MS: u64 = 2_000;
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum ConnectHeaderError {
#[error("invalid connect header name: {0}")]
InvalidName(String),
#[error("connect header `{0}` must be visible ASCII")]
NonAscii(String),
#[error("connect header `{0}` is owned by the WebSocket handshake")]
Reserved(String),
#[error("connect headers would total {total} bytes, over the {max} byte limit")]
TooLarge {
total: usize,
max: usize,
},
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ConnectHeaders {
headers: Vec<(HeaderName, HeaderValue)>,
total: usize,
}
impl ConnectHeaders {
pub const MAX_TOTAL_BYTES: usize = 8 * 1024;
const RESERVED: [&'static str; 3] = ["host", "connection", "upgrade"];
pub fn new() -> Self {
Self::default()
}
pub fn insert<N, V>(&mut self, name: N, value: V) -> std::result::Result<(), ConnectHeaderError>
where
N: TryInto<HeaderName>,
N::Error: std::fmt::Display,
V: TryInto<HeaderValue>,
V::Error: std::fmt::Display,
{
let name: HeaderName = name
.try_into()
.map_err(|e| ConnectHeaderError::InvalidName(e.to_string()))?;
let value: HeaderValue = value
.try_into()
.map_err(|_| ConnectHeaderError::NonAscii(name.as_str().to_string()))?;
let lower = name.as_str();
if Self::RESERVED.contains(&lower) || lower.starts_with("sec-websocket-") {
return Err(ConnectHeaderError::Reserved(lower.to_string()));
}
if value.to_str().is_err() {
return Err(ConnectHeaderError::NonAscii(lower.to_string()));
}
let existing = self.headers.iter().position(|(n, _)| *n == name);
let replaced = existing.map_or(0, |i| Self::wire_len(&self.headers[i]));
let total = self.total - replaced + Self::wire_len(&(name.clone(), value.clone()));
if total > Self::MAX_TOTAL_BYTES {
return Err(ConnectHeaderError::TooLarge {
total,
max: Self::MAX_TOTAL_BYTES,
});
}
match existing {
Some(i) => self.headers[i] = (name, value),
None => self.headers.push((name, value)),
}
self.total = total;
Ok(())
}
pub fn with<N, V>(mut self, name: N, value: V) -> std::result::Result<Self, ConnectHeaderError>
where
N: TryInto<HeaderName>,
N::Error: std::fmt::Display,
V: TryInto<HeaderValue>,
V::Error: std::fmt::Display,
{
self.insert(name, value)?;
Ok(self)
}
pub fn is_empty(&self) -> bool {
self.headers.is_empty()
}
pub fn len(&self) -> usize {
self.headers.len()
}
pub fn total_bytes(&self) -> usize {
self.total
}
pub fn iter(&self) -> impl Iterator<Item = (&HeaderName, &HeaderValue)> {
self.headers.iter().map(|(n, v)| (n, v))
}
fn wire_len((name, value): &(HeaderName, HeaderValue)) -> usize {
name.as_str().len() + 2 + value.len() + 2
}
pub(crate) fn apply(
&self,
request: &mut tokio_tungstenite::tungstenite::handshake::client::Request,
) {
let headers = request.headers_mut();
for (name, value) in &self.headers {
headers.insert(name.clone(), value.clone());
}
}
}
impl TryFrom<HeaderMap> for ConnectHeaders {
type Error = ConnectHeaderError;
fn try_from(map: HeaderMap) -> std::result::Result<Self, Self::Error> {
let mut out = Self::new();
for (name, value) in map.iter() {
out.insert(name.clone(), value.clone())?;
}
Ok(out)
}
}
impl TryFrom<&HeaderMap> for ConnectHeaders {
type Error = ConnectHeaderError;
fn try_from(map: &HeaderMap) -> std::result::Result<Self, Self::Error> {
Self::try_from(map.clone())
}
}
impl<T: EventMessage + Unpin> Connection<T> {
pub async fn connect(debug_ws_url: impl AsRef<str>) -> Result<Self> {
Self::connect_with_retries(debug_ws_url, DEFAULT_CONNECTION_RETRIES).await
}
pub async fn connect_with_retries(debug_ws_url: impl AsRef<str>, retries: u32) -> Result<Self> {
Self::connect_with_retries_and_headers(debug_ws_url, retries, &ConnectHeaders::new()).await
}
pub async fn connect_with_headers(
debug_ws_url: impl AsRef<str>,
headers: &ConnectHeaders,
) -> Result<Self> {
Self::connect_with_retries_and_headers(debug_ws_url, DEFAULT_CONNECTION_RETRIES, headers)
.await
}
pub async fn connect_with_retries_and_headers(
debug_ws_url: impl AsRef<str>,
retries: u32,
headers: &ConnectHeaders,
) -> Result<Self> {
let mut config = WebSocketConfig::default();
config.max_write_buffer_size = 4 * 1024 * 1024;
if !*WEBSOCKET_DEFAULTS {
config.max_message_size = None;
config.max_frame_size = None;
}
let url = debug_ws_url.as_ref();
let use_uring = crate::uring_fs::is_enabled();
let mut last_err = None;
for attempt in 0..=retries {
let result = if use_uring {
Self::connect_uring(url, config, headers).await
} else {
Self::connect_default(url, config, headers).await
};
match result {
Ok(ws) => {
return Ok(Self {
pending_commands: Default::default(),
ws,
next_id: 0,
needs_flush: false,
_marker: Default::default(),
});
}
Err(e) => {
let should_retry = match &e {
CdpError::Io(io_err)
if io_err.kind() == std::io::ErrorKind::ConnectionRefused =>
{
false
}
CdpError::Ws(tungstenite_err) => !matches!(
tungstenite_err,
tokio_tungstenite::tungstenite::Error::Http(_)
| tokio_tungstenite::tungstenite::Error::HttpFormat(_)
),
_ => true,
};
last_err = Some(e);
if !should_retry {
break;
}
if attempt < retries {
let backoff_ms =
(INITIAL_BACKOFF_MS * 3u64.saturating_pow(attempt)).min(MAX_BACKOFF_MS);
tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await;
}
}
}
}
Err(last_err.unwrap_or_else(|| CdpError::msg("connection failed")))
}
async fn connect_default(
url: &str,
config: WebSocketConfig,
headers: &ConnectHeaders,
) -> Result<WebSocketStream<ConnectStream>> {
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
if let Ok(mut request) = url.into_client_request() {
headers.apply(&mut request);
let uri = request.uri();
let is_plain_ws = uri.scheme_str() == Some("ws");
let host = uri.host().map(str::to_string);
let port = uri.port_u16().unwrap_or(9222);
if is_plain_ws && crate::dns::enabled() {
if let Some(h) = host {
if !crate::dns::is_ip_literal(&h) {
if let Ok(addrs) = crate::dns::resolve(&h, port).await {
match tokio::net::TcpStream::connect(&addrs[..]).await {
Ok(stream) => {
if *DISABLE_NAGLE {
let _ = stream.set_nodelay(true);
}
let (ws, _) = tokio_tungstenite::client_async_with_config(
request,
MaybeTlsStream::Plain(stream),
Some(config),
)
.await?;
return Ok(ws);
}
Err(e) => {
crate::dns::invalidate(&h, port);
return Err(CdpError::Io(e));
}
}
}
}
}
}
}
if headers.is_empty() {
let (ws, _) =
tokio_tungstenite::connect_async_with_config(url, Some(config), *DISABLE_NAGLE)
.await?;
return Ok(ws);
}
let mut request = url.into_client_request()?;
headers.apply(&mut request);
let (ws, _) =
tokio_tungstenite::connect_async_with_config(request, Some(config), *DISABLE_NAGLE)
.await?;
Ok(ws)
}
async fn connect_uring(
url: &str,
config: WebSocketConfig,
headers: &ConnectHeaders,
) -> Result<WebSocketStream<ConnectStream>> {
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
let mut request = url.into_client_request()?;
headers.apply(&mut request);
let host = request
.uri()
.host()
.ok_or_else(|| CdpError::msg("no host in CDP WebSocket URL"))?;
let port = request.uri().port_u16().unwrap_or(9222);
let addr_str = format!("{}:{}", host, port);
let addr: std::net::SocketAddr = match addr_str.parse() {
Ok(a) => a,
Err(_) => {
return Self::connect_default(url, config, headers).await;
}
};
let std_stream = crate::uring_fs::tcp_connect(addr)
.await
.map_err(CdpError::Io)?;
std_stream.set_nonblocking(true).map_err(CdpError::Io)?;
if *DISABLE_NAGLE {
let _ = std_stream.set_nodelay(true);
}
let tokio_stream = tokio::net::TcpStream::from_std(std_stream).map_err(CdpError::Io)?;
let (ws, _) = tokio_tungstenite::client_async_with_config(
request,
MaybeTlsStream::Plain(tokio_stream),
Some(config),
)
.await?;
Ok(ws)
}
}
impl<T: EventMessage> Connection<T> {
fn next_call_id(&mut self) -> CallId {
let id = CallId::new(self.next_id);
self.next_id = self.next_id.wrapping_add(1);
id
}
pub fn submit_command(
&mut self,
method: MethodId,
session_id: Option<SessionId>,
params: serde_json::Value,
) -> serde_json::Result<CallId> {
let id = self.next_call_id();
let call = MethodCall {
id,
method,
session_id: session_id.map(Into::into),
params,
};
self.pending_commands.push_back(call);
Ok(id)
}
fn start_send_next(&mut self, cx: &mut Context<'_>) -> Result<()> {
if self.needs_flush {
match self.ws.poll_flush_unpin(cx) {
Poll::Ready(Ok(())) => self.needs_flush = false,
Poll::Ready(Err(e)) => return Err(e.into()),
Poll::Pending => return Ok(()),
}
}
let mut sent_any = false;
while !self.pending_commands.is_empty() {
match self.ws.poll_ready_unpin(cx) {
Poll::Ready(Ok(())) => {
let Some(cmd) = self.pending_commands.pop_front() else {
break;
};
tracing::trace!("Sending {:?}", cmd);
let msg = serde_json::to_string(&cmd)?;
self.ws.start_send_unpin(msg.into())?;
sent_any = true;
}
_ => break,
}
}
if sent_any {
match self.ws.poll_flush_unpin(cx) {
Poll::Ready(Ok(())) => {}
Poll::Ready(Err(e)) => return Err(e.into()),
Poll::Pending => self.needs_flush = true,
}
}
Ok(())
}
}
const WS_CMD_CHANNEL_CAPACITY: usize = 2048;
const WS_READ_CHANNEL_CAPACITY: usize = 1024;
const MAX_IN_FLIGHT_DECODES: usize = 32;
const LARGE_FRAME_THRESHOLD: usize = 256 * 1024;
#[derive(Debug)]
pub struct AsyncConnection<T: EventMessage> {
pub reader: WsReader<T>,
pub cmd_tx: mpsc::Sender<MethodCall>,
pub writer_handle: tokio::task::JoinHandle<Result<()>>,
pub reader_handle: tokio::task::JoinHandle<()>,
pub next_id: usize,
}
impl<T: EventMessage + Unpin + Send + 'static> Connection<T> {
pub fn into_async(self) -> AsyncConnection<T> {
let (ws_sink, ws_stream) = self.ws.split();
let (cmd_tx, cmd_rx) = mpsc::channel(WS_CMD_CHANNEL_CAPACITY);
let (msg_tx, msg_rx) = mpsc::channel::<Result<Box<Message<T>>>>(WS_READ_CHANNEL_CAPACITY);
for call in self.pending_commands {
let _ = cmd_tx.try_send(call);
}
let writer_handle = tokio::spawn(ws_write_loop(ws_sink, cmd_rx));
let reader_handle = tokio::spawn(ws_read_loop::<T, _>(ws_stream, msg_tx));
let reader = WsReader {
rx: msg_rx,
_marker: PhantomData,
};
AsyncConnection {
reader,
cmd_tx,
writer_handle,
reader_handle,
next_id: self.next_id,
}
}
}
enum InFlightDecode<T: EventMessage + Send + 'static> {
Ready(Option<Result<Box<Message<T>>>>),
Blocking(tokio::task::JoinHandle<Result<Box<Message<T>>>>),
}
impl<T: EventMessage + Send + 'static> Future for InFlightDecode<T> {
type Output = Result<Box<Message<T>>>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match self.get_mut() {
InFlightDecode::Ready(slot) => Poll::Ready(
slot.take()
.expect("InFlightDecode::Ready polled after completion"),
),
InFlightDecode::Blocking(handle) => match Pin::new(handle).poll(cx) {
Poll::Ready(Ok(res)) => Poll::Ready(res),
Poll::Ready(Err(join_err)) => Poll::Ready(Err(CdpError::msg(format!(
"WS decode blocking task join error: {join_err}"
)))),
Poll::Pending => Poll::Pending,
},
}
}
}
async fn emit_decoded<T>(
tx: &mpsc::Sender<Result<Box<Message<T>>>>,
res: Result<Box<Message<T>>>,
) -> bool
where
T: EventMessage + Send + 'static,
{
match res {
Ok(msg) => tx.send(Ok(msg)).await.is_ok(),
Err(err) => {
tracing::debug!(
target: "chromiumoxide::conn::raw_ws::parse_errors",
"Dropping malformed WS frame: {err}",
);
true
}
}
}
async fn ws_read_loop<T, S>(mut stream: S, tx: mpsc::Sender<Result<Box<Message<T>>>>)
where
T: EventMessage + Send + 'static,
S: Stream<Item = std::result::Result<WsMessage, tokio_tungstenite::tungstenite::Error>> + Unpin,
{
let mut in_flight: FuturesOrdered<InFlightDecode<T>> = FuturesOrdered::new();
let mut stream_terminated = false;
let mut pending_err: Option<CdpError> = None;
loop {
tokio::select! {
biased;
Some(res) = in_flight.next(), if !in_flight.is_empty() => {
if !emit_decoded(&tx, res).await {
return;
}
}
maybe_frame = stream.next(),
if !stream_terminated && in_flight.len() < MAX_IN_FLIGHT_DECODES =>
{
match maybe_frame {
Some(Ok(WsMessage::Text(text))) => {
if text.len() >= LARGE_FRAME_THRESHOLD {
in_flight.push_back(InFlightDecode::Blocking(
tokio::task::spawn_blocking(move || {
decode_message::<T>(text.as_bytes(), None)
}),
));
} else {
let res = decode_message::<T>(text.as_bytes(), Some(&text));
in_flight.push_back(InFlightDecode::Ready(Some(res)));
}
}
Some(Ok(WsMessage::Binary(buf))) => {
if buf.len() >= LARGE_FRAME_THRESHOLD {
in_flight.push_back(InFlightDecode::Blocking(
tokio::task::spawn_blocking(move || {
decode_message::<T>(&buf, None)
}),
));
} else {
let res = decode_message::<T>(&buf, None);
in_flight.push_back(InFlightDecode::Ready(Some(res)));
}
}
Some(Ok(WsMessage::Close(_))) => {
stream_terminated = true;
}
Some(Ok(WsMessage::Ping(_))) | Some(Ok(WsMessage::Pong(_))) => {}
Some(Ok(msg)) => {
tracing::debug!(
target: "chromiumoxide::conn::raw_ws::parse_errors",
"Unexpected WS message type: {:?}",
msg
);
}
Some(Err(err)) => {
stream_terminated = true;
pending_err = Some(CdpError::Ws(err));
}
None => {
stream_terminated = true;
}
}
}
else => {
break;
}
}
}
if let Some(err) = pending_err {
let _ = tx.send(Err(err)).await;
}
}
async fn ws_write_loop(
mut sink: SplitSink<WebSocketStream<ConnectStream>, WsMessage>,
mut rx: mpsc::Receiver<MethodCall>,
) -> Result<()> {
while let Some(call) = rx.recv().await {
let msg = crate::serde_json::to_string(&call)?;
sink.feed(WsMessage::Text(msg.into()))
.await
.map_err(CdpError::Ws)?;
while let Ok(call) = rx.try_recv() {
let msg = crate::serde_json::to_string(&call)?;
sink.feed(WsMessage::Text(msg.into()))
.await
.map_err(CdpError::Ws)?;
}
sink.flush().await.map_err(CdpError::Ws)?;
}
let _ = sink.close().await;
Ok(())
}
#[derive(Debug)]
pub struct WsReader<T: EventMessage> {
rx: mpsc::Receiver<Result<Box<Message<T>>>>,
_marker: PhantomData<T>,
}
impl<T: EventMessage + Unpin> WsReader<T> {
pub async fn next_message(&mut self) -> Option<Result<Box<Message<T>>>> {
self.rx.recv().await
}
}
impl<T: EventMessage + Unpin> Stream for Connection<T> {
type Item = Result<Box<Message<T>>>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let pin = self.get_mut();
if let Err(err) = pin.start_send_next(cx) {
return Poll::Ready(Some(Err(err)));
}
const MAX_SKIPS_PER_POLL: u32 = 16;
let mut skips: u32 = 0;
loop {
match ready!(pin.ws.poll_next_unpin(cx)) {
Some(Ok(WsMessage::Text(text))) => {
match decode_message::<T>(text.as_bytes(), Some(&text)) {
Ok(msg) => return Poll::Ready(Some(Ok(msg))),
Err(err) => {
tracing::debug!(
target: "chromiumoxide::conn::raw_ws::parse_errors",
"Dropping malformed text WS frame: {err}",
);
skips += 1;
}
}
}
Some(Ok(WsMessage::Binary(buf))) => match decode_message::<T>(&buf, None) {
Ok(msg) => return Poll::Ready(Some(Ok(msg))),
Err(err) => {
tracing::debug!(
target: "chromiumoxide::conn::raw_ws::parse_errors",
"Dropping malformed binary WS frame: {err}",
);
skips += 1;
}
},
Some(Ok(WsMessage::Close(_))) => return Poll::Ready(None),
Some(Ok(WsMessage::Ping(_))) | Some(Ok(WsMessage::Pong(_))) => {
skips += 1;
}
Some(Ok(msg)) => {
tracing::debug!(
target: "chromiumoxide::conn::raw_ws::parse_errors",
"Unexpected WS message type: {:?}",
msg
);
skips += 1;
}
Some(Err(err)) => return Poll::Ready(Some(Err(CdpError::Ws(err)))),
None => return Poll::Ready(None),
}
if skips >= MAX_SKIPS_PER_POLL {
cx.waker().wake_by_ref();
return Poll::Pending;
}
}
}
}
#[cfg(not(feature = "serde_stacker"))]
fn decode_message<T: EventMessage>(
bytes: &[u8],
raw_text_for_logging: Option<&str>,
) -> Result<Box<Message<T>>> {
match serde_json::from_slice::<Box<Message<T>>>(bytes) {
Ok(msg) => {
tracing::trace!("Received {:?}", msg);
Ok(msg)
}
Err(err) => {
if let Some(txt) = raw_text_for_logging {
let preview = &txt[..txt.len().min(512)];
tracing::debug!(
target: "chromiumoxide::conn::raw_ws::parse_errors",
msg_len = txt.len(),
"Skipping unrecognized WS message {err} preview={preview}",
);
} else {
tracing::debug!(
target: "chromiumoxide::conn::raw_ws::parse_errors",
"Skipping unrecognized binary WS message {err}",
);
}
Err(err.into())
}
}
}
#[cfg(feature = "serde_stacker")]
fn decode_message<T: EventMessage>(
bytes: &[u8],
raw_text_for_logging: Option<&str>,
) -> Result<Box<Message<T>>> {
use serde::Deserialize;
let mut de = serde_json::Deserializer::from_slice(bytes);
de.disable_recursion_limit();
let de = serde_stacker::Deserializer::new(&mut de);
match Box::<Message<T>>::deserialize(de) {
Ok(msg) => {
tracing::trace!("Received {:?}", msg);
Ok(msg)
}
Err(err) => {
if let Some(txt) = raw_text_for_logging {
let preview = &txt[..txt.len().min(512)];
tracing::debug!(
target: "chromiumoxide::conn::raw_ws::parse_errors",
msg_len = txt.len(),
"Skipping unrecognized WS message {err} preview={preview}",
);
} else {
tracing::debug!(
target: "chromiumoxide::conn::raw_ws::parse_errors",
"Skipping unrecognized binary WS message {err}",
);
}
Err(err.into())
}
}
}
#[cfg(test)]
mod ws_read_loop_tests {
use super::*;
use chromiumoxide_cdp::cdp::CdpEventMessage;
use chromiumoxide_types::CallId;
use futures_util::stream;
use tokio::sync::mpsc;
use tokio_tungstenite::tungstenite::Message as WsMessage;
fn response_frame(id: u64) -> WsMessage {
WsMessage::Text(
format!(r#"{{"id":{id},"result":{{"ok":true}}}}"#)
.to_string()
.into(),
)
}
fn large_response_frame(id: u64, blob_bytes: usize) -> WsMessage {
let blob = "x".repeat(blob_bytes);
WsMessage::Text(
format!(r#"{{"id":{id},"result":{{"blob":"{blob}"}}}}"#)
.to_string()
.into(),
)
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn forwards_messages_in_stream_order() {
let frames = vec![
Ok(response_frame(1)),
Ok(response_frame(2)),
Ok(response_frame(3)),
];
let stream = stream::iter(frames);
let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(8);
let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
for expected in [1u64, 2, 3] {
let msg = rx.recv().await.expect("msg").expect("decode ok");
if let Message::Response(resp) = *msg {
assert_eq!(resp.id, CallId::new(expected as usize));
} else {
panic!("expected Response");
}
}
assert!(rx.recv().await.is_none(), "channel must close on EOF");
task.await.expect("reader task join");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn pings_and_pongs_never_reach_the_handler() {
let frames = vec![
Ok(WsMessage::Ping(vec![1, 2, 3].into())),
Ok(response_frame(7)),
Ok(WsMessage::Pong(vec![].into())),
Ok(response_frame(8)),
];
let stream = stream::iter(frames);
let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(8);
let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
for expected in [7u64, 8] {
let msg = rx.recv().await.expect("msg").expect("decode ok");
if let Message::Response(resp) = *msg {
assert_eq!(resp.id, CallId::new(expected as usize));
}
}
assert!(rx.recv().await.is_none());
task.await.expect("reader task join");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn malformed_frames_do_not_block_subsequent_valid_frames() {
let frames = vec![
Ok(WsMessage::Text("{not valid json".to_string().into())),
Ok(response_frame(42)),
];
let stream = stream::iter(frames);
let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(8);
let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
let msg = rx.recv().await.expect("msg").expect("decode ok");
if let Message::Response(resp) = *msg {
assert_eq!(resp.id, CallId::new(42));
}
assert!(rx.recv().await.is_none());
task.await.expect("reader task join");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn close_frame_terminates_the_reader() {
let frames = vec![
Ok(response_frame(1)),
Ok(WsMessage::Close(None)),
Ok(response_frame(2)), ];
let stream = stream::iter(frames);
let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(8);
let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
let msg = rx.recv().await.expect("msg").expect("decode ok");
if let Message::Response(resp) = *msg {
assert_eq!(resp.id, CallId::new(1));
}
assert!(
rx.recv().await.is_none(),
"reader must exit on Close; frames after Close must not appear"
);
task.await.expect("reader task join");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn transport_error_is_forwarded_once_then_reader_exits() {
let frames = vec![
Ok(response_frame(1)),
Err(tokio_tungstenite::tungstenite::Error::ConnectionClosed),
Ok(response_frame(2)),
];
let stream = stream::iter(frames);
let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(8);
let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
let msg = rx.recv().await.expect("msg").expect("ok");
assert!(matches!(*msg, Message::Response(_)));
match rx.recv().await {
Some(Err(CdpError::Ws(_))) => {}
other => panic!("expected forwarded Ws error, got {other:?}"),
}
assert!(rx.recv().await.is_none());
task.await.expect("reader task join");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn bounded_channel_does_not_deadlock_under_backpressure() {
const N: u64 = 512;
let frames: Vec<_> = (1..=N).map(|id| Ok(response_frame(id))).collect();
let stream = stream::iter(frames);
let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(1);
let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
let deadline = std::time::Duration::from_secs(5);
let collected = tokio::time::timeout(deadline, async {
let mut seen = 0u64;
while let Some(frame) = rx.recv().await {
let msg = frame.expect("decode ok");
if let Message::Response(resp) = *msg {
seen += 1;
assert_eq!(
resp.id,
CallId::new(seen as usize),
"back-pressure must preserve FIFO order"
);
}
}
seen
})
.await
.expect("reader must make forward progress despite cap-1 back-pressure");
assert_eq!(collected, N, "all frames must arrive");
task.await.expect("reader task join");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn large_message_decodes_without_corruption() {
let big = 2 * 1024 * 1024; let frames = vec![Ok(large_response_frame(100, big)), Ok(response_frame(101))];
let stream = stream::iter(frames);
let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(4);
let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
let first = rx.recv().await.expect("msg").expect("ok");
if let Message::Response(resp) = *first {
assert_eq!(resp.id, CallId::new(100));
}
let second = rx.recv().await.expect("msg").expect("ok");
if let Message::Response(resp) = *second {
assert_eq!(resp.id, CallId::new(101));
}
assert!(rx.recv().await.is_none());
task.await.expect("reader task join");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn pipelined_large_and_small_frames_keep_fifo_order() {
let big = 2 * 1024 * 1024; let frames = vec![
Ok(large_response_frame(1, big)),
Ok(response_frame(2)),
Ok(response_frame(3)),
Ok(large_response_frame(4, big)),
Ok(response_frame(5)),
Ok(large_response_frame(6, big)),
Ok(response_frame(7)),
Ok(response_frame(8)),
];
let expected: Vec<usize> = (1..=8).collect();
let stream = stream::iter(frames);
let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(16);
let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
let deadline = std::time::Duration::from_secs(10);
let observed = tokio::time::timeout(deadline, async {
let mut ids = Vec::with_capacity(expected.len());
while let Some(frame) = rx.recv().await {
let msg = frame.expect("decode ok");
if let Message::Response(resp) = *msg {
ids.push(CallId::new(ids.len() + 1));
assert_eq!(
resp.id,
*ids.last().unwrap(),
"pipelined reader must emit frames in strict arrival order \
regardless of per-frame decode latency"
);
}
}
ids
})
.await
.expect("pipelined reader should make forward progress within 10s");
assert_eq!(
observed.len(),
expected.len(),
"all {} frames must reach the Handler",
expected.len()
);
task.await.expect("reader task join");
}
}
#[cfg(test)]
mod connect_headers_tests {
use super::ConnectHeaderError;
use super::ConnectHeaders;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderMap;
#[test]
fn empty_apply_is_a_no_op() {
let url = "ws://127.0.0.1:9222/devtools/browser/x";
let mut req = url.into_client_request().unwrap();
let before = req.headers().clone();
ConnectHeaders::new().apply(&mut req);
assert_eq!(&before, req.headers());
}
#[test]
fn apply_adds_and_replaces_by_name() {
let url = "ws://127.0.0.1:9222/devtools/browser/x";
let mut req = url.into_client_request().unwrap();
let set = ConnectHeaders::new()
.with("x-a", "1")
.unwrap()
.with("x-a", "2")
.unwrap()
.with("x-b", "3")
.unwrap();
assert_eq!(set.len(), 2);
set.apply(&mut req);
assert_eq!(req.headers().get_all("x-a").iter().count(), 1);
assert_eq!(req.headers()["x-a"], "2");
assert_eq!(req.headers()["x-b"], "3");
}
#[test]
fn total_bytes_tracks_replacement() {
let mut set = ConnectHeaders::new();
set.insert("x-a", "12345").unwrap();
assert_eq!(set.total_bytes(), 3 + 2 + 5 + 2);
set.insert("x-a", "1").unwrap();
assert_eq!(set.total_bytes(), 3 + 2 + 1 + 2);
set.insert("x-b", "").unwrap();
assert_eq!(set.total_bytes(), 8 + 7);
}
#[test]
fn limit_counts_the_whole_set() {
let half = "x".repeat(ConnectHeaders::MAX_TOTAL_BYTES / 2);
let mut set = ConnectHeaders::new();
set.insert("x-a", half.as_str()).unwrap();
let err = set.insert("x-b", half.as_str()).unwrap_err();
assert!(matches!(err, ConnectHeaderError::TooLarge { .. }), "{err}");
assert_eq!(set.len(), 1);
set.insert("x-a", "1").unwrap();
set.insert("x-b", half.as_str()).unwrap();
}
#[test]
fn reserved_names_are_case_insensitive() {
for name in ["Host", "CONNECTION", "Upgrade", "Sec-WebSocket-Protocol"] {
let err = ConnectHeaders::new().with(name, "x").unwrap_err();
assert!(
matches!(err, ConnectHeaderError::Reserved(_)),
"{name}: {err}"
);
}
}
#[test]
fn obs_text_bytes_are_refused() {
let value = tokio_tungstenite::tungstenite::http::HeaderValue::from_bytes(&[0xE9]).unwrap();
let err = ConnectHeaders::new().with("x-a", value).unwrap_err();
assert!(matches!(err, ConnectHeaderError::NonAscii(_)), "{err}");
}
#[test]
fn header_map_conversion_validates_every_entry() {
let mut map = HeaderMap::new();
map.insert("x-a", "1".parse().unwrap());
map.insert("host", "h".parse().unwrap());
let err = ConnectHeaders::try_from(map).unwrap_err();
assert!(matches!(err, ConnectHeaderError::Reserved(_)), "{err}");
}
}