chromiumoxide/conn.rs
1use std::collections::VecDeque;
2use std::marker::PhantomData;
3use std::pin::Pin;
4use std::task::ready;
5
6use futures_util::stream::{FuturesOrdered, SplitSink};
7use futures_util::{SinkExt, Stream, StreamExt};
8use std::future::Future;
9use std::task::{Context, Poll};
10use tokio::sync::mpsc;
11use tokio_tungstenite::tungstenite::http::{HeaderMap, HeaderName, HeaderValue};
12use tokio_tungstenite::tungstenite::Message as WsMessage;
13use tokio_tungstenite::MaybeTlsStream;
14use tokio_tungstenite::{tungstenite::protocol::WebSocketConfig, WebSocketStream};
15
16use chromiumoxide_cdp::cdp::browser_protocol::target::SessionId;
17use chromiumoxide_types::{CallId, EventMessage, Message, MethodCall, MethodId};
18
19use crate::error::CdpError;
20use crate::error::Result;
21
22type ConnectStream = MaybeTlsStream<tokio::net::TcpStream>;
23
24/// Exchanges the messages with the websocket
25#[must_use = "streams do nothing unless polled"]
26#[derive(Debug)]
27pub struct Connection<T: EventMessage> {
28 /// Queue of commands to send.
29 pending_commands: VecDeque<MethodCall>,
30 /// The websocket of the chromium instance
31 ws: WebSocketStream<ConnectStream>,
32 /// The identifier for a specific command
33 next_id: usize,
34 /// Whether the write buffer has unsent data that needs flushing.
35 needs_flush: bool,
36 /// The phantom marker.
37 _marker: PhantomData<T>,
38}
39
40lazy_static::lazy_static! {
41 /// Nagle's algorithm disabled?
42 static ref DISABLE_NAGLE: bool = match std::env::var("DISABLE_NAGLE") {
43 Ok(disable_nagle) => disable_nagle == "true",
44 _ => true
45 };
46 /// Websocket config defaults
47 static ref WEBSOCKET_DEFAULTS: bool = match std::env::var("WEBSOCKET_DEFAULTS") {
48 Ok(d) => d == "true",
49 _ => false
50 };
51}
52
53/// Default number of WebSocket connection retry attempts.
54pub const DEFAULT_CONNECTION_RETRIES: u32 = 4;
55
56/// Initial backoff delay between connection retries (in milliseconds).
57const INITIAL_BACKOFF_MS: u64 = 50;
58
59/// Maximum backoff delay between connection retries (in milliseconds).
60pub(crate) const MAX_BACKOFF_MS: u64 = 2_000;
61
62/// Why a header was refused by [`ConnectHeaders`].
63#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
64pub enum ConnectHeaderError {
65 /// The name is not a valid HTTP header name.
66 #[error("invalid connect header name: {0}")]
67 InvalidName(String),
68 /// The value is not visible ASCII. The upgrade request is written as text,
69 /// so a value tungstenite cannot render as a `str` would fail the handshake
70 /// at connect time instead of here.
71 #[error("connect header `{0}` must be visible ASCII")]
72 NonAscii(String),
73 /// The name belongs to the WebSocket handshake itself (`host`,
74 /// `connection`, `upgrade`, `sec-websocket-*`). tungstenite writes those
75 /// and rejects duplicates, so setting one can only break the upgrade.
76 #[error("connect header `{0}` is owned by the WebSocket handshake")]
77 Reserved(String),
78 /// Adding the header would push the set past
79 /// [`ConnectHeaders::MAX_TOTAL_BYTES`].
80 #[error("connect headers would total {total} bytes, over the {max} byte limit")]
81 TooLarge {
82 /// Wire size the set would have after the insert.
83 total: usize,
84 /// The limit, [`ConnectHeaders::MAX_TOTAL_BYTES`].
85 max: usize,
86 },
87}
88
89/// Headers added to the CDP WebSocket upgrade request (the HTTP `GET` that
90/// opens the socket), as opposed to `extra_headers`, which are applied to the
91/// page's own network requests via `Network.setExtraHTTPHeaders`.
92///
93/// A gateway in front of the browser sees these on the upgrade head and can
94/// use them to tag or route the connection. Each name appears once; inserting
95/// a name that is already present replaces its value. The set is validated at
96/// insert time, so a `ConnectHeaders` that exists can always be sent. An empty
97/// set leaves the upgrade request byte for byte as it was without one.
98#[derive(Debug, Clone, Default, PartialEq, Eq)]
99pub struct ConnectHeaders {
100 headers: Vec<(HeaderName, HeaderValue)>,
101 /// Wire bytes the current set would occupy, see [`Self::wire_len`].
102 total: usize,
103}
104
105impl ConnectHeaders {
106 /// Upper bound on the wire size of the whole set, counted as
107 /// `name: value\r\n` per header. Keeps a misconfigured caller from
108 /// shipping a request head a proxy will reject.
109 pub const MAX_TOTAL_BYTES: usize = 8 * 1024;
110
111 const RESERVED: [&'static str; 3] = ["host", "connection", "upgrade"];
112
113 /// An empty set.
114 pub fn new() -> Self {
115 Self::default()
116 }
117
118 /// Add one header, replacing any existing header of the same name.
119 ///
120 /// `name` and `value` accept anything the `http` crate can convert
121 /// (`&str`, `String`, or the typed `HeaderName` / `HeaderValue`).
122 pub fn insert<N, V>(&mut self, name: N, value: V) -> std::result::Result<(), ConnectHeaderError>
123 where
124 N: TryInto<HeaderName>,
125 N::Error: std::fmt::Display,
126 V: TryInto<HeaderValue>,
127 V::Error: std::fmt::Display,
128 {
129 let name: HeaderName = name
130 .try_into()
131 .map_err(|e| ConnectHeaderError::InvalidName(e.to_string()))?;
132 let value: HeaderValue = value
133 .try_into()
134 .map_err(|_| ConnectHeaderError::NonAscii(name.as_str().to_string()))?;
135
136 let lower = name.as_str();
137 if Self::RESERVED.contains(&lower) || lower.starts_with("sec-websocket-") {
138 return Err(ConnectHeaderError::Reserved(lower.to_string()));
139 }
140 // `HeaderValue` allows obs-text (0x80..=0xFF); the handshake writer
141 // calls `to_str`, which does not.
142 if value.to_str().is_err() {
143 return Err(ConnectHeaderError::NonAscii(lower.to_string()));
144 }
145
146 let existing = self.headers.iter().position(|(n, _)| *n == name);
147 let replaced = existing.map_or(0, |i| Self::wire_len(&self.headers[i]));
148 let total = self.total - replaced + Self::wire_len(&(name.clone(), value.clone()));
149 if total > Self::MAX_TOTAL_BYTES {
150 return Err(ConnectHeaderError::TooLarge {
151 total,
152 max: Self::MAX_TOTAL_BYTES,
153 });
154 }
155
156 match existing {
157 Some(i) => self.headers[i] = (name, value),
158 None => self.headers.push((name, value)),
159 }
160 self.total = total;
161 Ok(())
162 }
163
164 /// Builder form of [`Self::insert`].
165 pub fn with<N, V>(mut self, name: N, value: V) -> std::result::Result<Self, ConnectHeaderError>
166 where
167 N: TryInto<HeaderName>,
168 N::Error: std::fmt::Display,
169 V: TryInto<HeaderValue>,
170 V::Error: std::fmt::Display,
171 {
172 self.insert(name, value)?;
173 Ok(self)
174 }
175
176 /// Whether the set holds no headers.
177 pub fn is_empty(&self) -> bool {
178 self.headers.is_empty()
179 }
180
181 /// Number of headers in the set.
182 pub fn len(&self) -> usize {
183 self.headers.len()
184 }
185
186 /// Wire bytes the set occupies on the upgrade request.
187 pub fn total_bytes(&self) -> usize {
188 self.total
189 }
190
191 /// The headers in insertion order.
192 pub fn iter(&self) -> impl Iterator<Item = (&HeaderName, &HeaderValue)> {
193 self.headers.iter().map(|(n, v)| (n, v))
194 }
195
196 /// Bytes `name: value\r\n` takes on the wire.
197 fn wire_len((name, value): &(HeaderName, HeaderValue)) -> usize {
198 name.as_str().len() + 2 + value.len() + 2
199 }
200
201 /// Put the set on an upgrade request. A no-op when empty, so the
202 /// request tungstenite builds from the URL is untouched.
203 pub(crate) fn apply(
204 &self,
205 request: &mut tokio_tungstenite::tungstenite::handshake::client::Request,
206 ) {
207 let headers = request.headers_mut();
208 for (name, value) in &self.headers {
209 headers.insert(name.clone(), value.clone());
210 }
211 }
212}
213
214impl TryFrom<HeaderMap> for ConnectHeaders {
215 type Error = ConnectHeaderError;
216
217 /// Every entry is validated as by [`ConnectHeaders::insert`]. A name that
218 /// repeats in the map keeps its last value.
219 fn try_from(map: HeaderMap) -> std::result::Result<Self, Self::Error> {
220 let mut out = Self::new();
221 for (name, value) in map.iter() {
222 out.insert(name.clone(), value.clone())?;
223 }
224 Ok(out)
225 }
226}
227
228impl TryFrom<&HeaderMap> for ConnectHeaders {
229 type Error = ConnectHeaderError;
230
231 fn try_from(map: &HeaderMap) -> std::result::Result<Self, Self::Error> {
232 Self::try_from(map.clone())
233 }
234}
235
236impl<T: EventMessage + Unpin> Connection<T> {
237 pub async fn connect(debug_ws_url: impl AsRef<str>) -> Result<Self> {
238 Self::connect_with_retries(debug_ws_url, DEFAULT_CONNECTION_RETRIES).await
239 }
240
241 pub async fn connect_with_retries(debug_ws_url: impl AsRef<str>, retries: u32) -> Result<Self> {
242 Self::connect_with_retries_and_headers(debug_ws_url, retries, &ConnectHeaders::new()).await
243 }
244
245 /// Like [`Self::connect`], with `headers` added to the WebSocket upgrade
246 /// request. Every attempt, including retries, carries them.
247 pub async fn connect_with_headers(
248 debug_ws_url: impl AsRef<str>,
249 headers: &ConnectHeaders,
250 ) -> Result<Self> {
251 Self::connect_with_retries_and_headers(debug_ws_url, DEFAULT_CONNECTION_RETRIES, headers)
252 .await
253 }
254
255 /// Like [`Self::connect_with_retries`], with `headers` added to the
256 /// WebSocket upgrade request on every attempt.
257 pub async fn connect_with_retries_and_headers(
258 debug_ws_url: impl AsRef<str>,
259 retries: u32,
260 headers: &ConnectHeaders,
261 ) -> Result<Self> {
262 let mut config = WebSocketConfig::default();
263
264 // Cap the internal write buffer so a slow receiver cannot cause
265 // unbounded memory growth (default is usize::MAX).
266 config.max_write_buffer_size = 4 * 1024 * 1024;
267
268 if !*WEBSOCKET_DEFAULTS {
269 config.max_message_size = None;
270 config.max_frame_size = None;
271 }
272
273 let url = debug_ws_url.as_ref();
274 let use_uring = crate::uring_fs::is_enabled();
275 let mut last_err = None;
276
277 for attempt in 0..=retries {
278 let result = if use_uring {
279 Self::connect_uring(url, config, headers).await
280 } else {
281 Self::connect_default(url, config, headers).await
282 };
283
284 match result {
285 Ok(ws) => {
286 return Ok(Self {
287 pending_commands: Default::default(),
288 ws,
289 next_id: 0,
290 needs_flush: false,
291 _marker: Default::default(),
292 });
293 }
294 Err(e) => {
295 // Detect non-retriable errors early to avoid wasting time
296 // on connections that will never succeed.
297 let should_retry = match &e {
298 // Connection refused — nothing is listening on this port.
299 CdpError::Io(io_err)
300 if io_err.kind() == std::io::ErrorKind::ConnectionRefused =>
301 {
302 false
303 }
304 // HTTP response to a WebSocket upgrade (e.g. wrong path
305 // returns 404 / redirect) — retrying the same URL won't help.
306 CdpError::Ws(tungstenite_err) => !matches!(
307 tungstenite_err,
308 tokio_tungstenite::tungstenite::Error::Http(_)
309 | tokio_tungstenite::tungstenite::Error::HttpFormat(_)
310 ),
311 _ => true,
312 };
313
314 last_err = Some(e);
315
316 if !should_retry {
317 break;
318 }
319
320 if attempt < retries {
321 let backoff_ms =
322 (INITIAL_BACKOFF_MS * 3u64.saturating_pow(attempt)).min(MAX_BACKOFF_MS);
323 tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await;
324 }
325 }
326 }
327 }
328
329 Err(last_err.unwrap_or_else(|| CdpError::msg("connection failed")))
330 }
331
332 /// Default path: let tokio-tungstenite handle TCP connect + WS handshake.
333 ///
334 /// For a plaintext `ws://` endpoint addressed by **hostname**, the resolution
335 /// is served from [`crate::dns`] (a short-TTL cache) and the handshake runs
336 /// over a pre-connected socket — the same shape the io_uring path uses — so
337 /// repeated (re)connects skip `getaddrinfo`. The original request is passed
338 /// to `client_async_with_config`, preserving the `Host` header. Everything
339 /// else — `wss://` (TLS), IP-literal hosts (no DNS), or any URL parse issue —
340 /// falls through to `connect_async_with_config` exactly as before, so TLS
341 /// handling and behavior are unchanged. The cache only maps host→IP; the port
342 /// and path always come from `url`, so it can never retarget a connection.
343 ///
344 /// `headers` go on the upgrade request before the handshake on both
345 /// branches. When the set is empty the fallback still hands the bare `url`
346 /// to tungstenite, so the request it writes is the one it always wrote.
347 async fn connect_default(
348 url: &str,
349 config: WebSocketConfig,
350 headers: &ConnectHeaders,
351 ) -> Result<WebSocketStream<ConnectStream>> {
352 use tokio_tungstenite::tungstenite::client::IntoClientRequest;
353
354 if let Ok(mut request) = url.into_client_request() {
355 headers.apply(&mut request);
356 let uri = request.uri();
357 let is_plain_ws = uri.scheme_str() == Some("ws");
358 let host = uri.host().map(str::to_string);
359 let port = uri.port_u16().unwrap_or(9222);
360
361 if is_plain_ws && crate::dns::enabled() {
362 if let Some(h) = host {
363 // Hostname only — IP literals need no DNS and are left to the
364 // default path below.
365 if !crate::dns::is_ip_literal(&h) {
366 if let Ok(addrs) = crate::dns::resolve(&h, port).await {
367 match tokio::net::TcpStream::connect(&addrs[..]).await {
368 Ok(stream) => {
369 if *DISABLE_NAGLE {
370 let _ = stream.set_nodelay(true);
371 }
372 let (ws, _) = tokio_tungstenite::client_async_with_config(
373 request,
374 MaybeTlsStream::Plain(stream),
375 Some(config),
376 )
377 .await?;
378 return Ok(ws);
379 }
380 Err(e) => {
381 // A cached address that won't connect may be
382 // stale — drop it so the retry re-resolves.
383 crate::dns::invalidate(&h, port);
384 return Err(CdpError::Io(e));
385 }
386 }
387 }
388 // resolve failed → fall through to the default path, which
389 // surfaces the real connection error.
390 }
391 }
392 }
393 }
394
395 if headers.is_empty() {
396 let (ws, _) =
397 tokio_tungstenite::connect_async_with_config(url, Some(config), *DISABLE_NAGLE)
398 .await?;
399 return Ok(ws);
400 }
401
402 let mut request = url.into_client_request()?;
403 headers.apply(&mut request);
404 let (ws, _) =
405 tokio_tungstenite::connect_async_with_config(request, Some(config), *DISABLE_NAGLE)
406 .await?;
407 Ok(ws)
408 }
409
410 /// io_uring path: pre-connect the TCP socket via io_uring, then do WS
411 /// handshake over the pre-connected stream.
412 ///
413 /// `headers` go on the upgrade request before the handshake, and travel
414 /// with the fallback to [`Self::connect_default`].
415 async fn connect_uring(
416 url: &str,
417 config: WebSocketConfig,
418 headers: &ConnectHeaders,
419 ) -> Result<WebSocketStream<ConnectStream>> {
420 use tokio_tungstenite::tungstenite::client::IntoClientRequest;
421
422 let mut request = url.into_client_request()?;
423 headers.apply(&mut request);
424 let host = request
425 .uri()
426 .host()
427 .ok_or_else(|| CdpError::msg("no host in CDP WebSocket URL"))?;
428 let port = request.uri().port_u16().unwrap_or(9222);
429
430 // Resolve host → SocketAddr (CDP is always localhost, so this is fast).
431 let addr_str = format!("{}:{}", host, port);
432 let addr: std::net::SocketAddr = match addr_str.parse() {
433 Ok(a) => a,
434 Err(_) => {
435 // Hostname needs DNS — fall back to default path.
436 return Self::connect_default(url, config, headers).await;
437 }
438 };
439
440 // TCP connect via io_uring.
441 let std_stream = crate::uring_fs::tcp_connect(addr)
442 .await
443 .map_err(CdpError::Io)?;
444
445 // Set non-blocking + Nagle.
446 std_stream.set_nonblocking(true).map_err(CdpError::Io)?;
447 if *DISABLE_NAGLE {
448 let _ = std_stream.set_nodelay(true);
449 }
450
451 // Wrap in tokio TcpStream.
452 let tokio_stream = tokio::net::TcpStream::from_std(std_stream).map_err(CdpError::Io)?;
453
454 // WebSocket handshake over the pre-connected stream.
455 let (ws, _) = tokio_tungstenite::client_async_with_config(
456 request,
457 MaybeTlsStream::Plain(tokio_stream),
458 Some(config),
459 )
460 .await?;
461
462 Ok(ws)
463 }
464}
465
466impl<T: EventMessage> Connection<T> {
467 fn next_call_id(&mut self) -> CallId {
468 let id = CallId::new(self.next_id);
469 self.next_id = self.next_id.wrapping_add(1);
470 id
471 }
472
473 /// Queue in the command to send over the socket and return the id for this
474 /// command
475 pub fn submit_command(
476 &mut self,
477 method: MethodId,
478 session_id: Option<SessionId>,
479 params: serde_json::Value,
480 ) -> serde_json::Result<CallId> {
481 let id = self.next_call_id();
482 let call = MethodCall {
483 id,
484 method,
485 session_id: session_id.map(Into::into),
486 params,
487 };
488 self.pending_commands.push_back(call);
489 Ok(id)
490 }
491
492 /// Buffer all queued commands into the WebSocket sink, then flush once.
493 ///
494 /// This batches multiple CDP commands into a single TCP write instead of
495 /// flushing after every individual message.
496 fn start_send_next(&mut self, cx: &mut Context<'_>) -> Result<()> {
497 // Complete any pending flush from a previous poll first.
498 if self.needs_flush {
499 match self.ws.poll_flush_unpin(cx) {
500 Poll::Ready(Ok(())) => self.needs_flush = false,
501 Poll::Ready(Err(e)) => return Err(e.into()),
502 Poll::Pending => return Ok(()),
503 }
504 }
505
506 // Buffer as many queued commands as the sink will accept.
507 let mut sent_any = false;
508 while !self.pending_commands.is_empty() {
509 match self.ws.poll_ready_unpin(cx) {
510 Poll::Ready(Ok(())) => {
511 let Some(cmd) = self.pending_commands.pop_front() else {
512 break;
513 };
514 tracing::trace!("Sending {:?}", cmd);
515 let msg = serde_json::to_string(&cmd)?;
516 self.ws.start_send_unpin(msg.into())?;
517 sent_any = true;
518 }
519 _ => break,
520 }
521 }
522
523 // Flush the entire batch in one write.
524 if sent_any {
525 match self.ws.poll_flush_unpin(cx) {
526 Poll::Ready(Ok(())) => {}
527 Poll::Ready(Err(e)) => return Err(e.into()),
528 Poll::Pending => self.needs_flush = true,
529 }
530 }
531
532 Ok(())
533 }
534}
535
536/// Capacity of the bounded channel feeding the background WS writer task.
537/// Large enough that bursts of CDP commands never block the handler, small
538/// enough to apply back-pressure before memory grows without bound.
539const WS_CMD_CHANNEL_CAPACITY: usize = 2048;
540
541/// Capacity of the bounded channel from the background WS reader task to
542/// the Handler. Keeps decoded CDP messages buffered so the reader task
543/// can keep reading the socket while the Handler processes a backlog;
544/// applies TCP-level back-pressure on Chrome when the Handler is slow
545/// (the reader awaits channel capacity, stops draining the socket).
546const WS_READ_CHANNEL_CAPACITY: usize = 1024;
547
548/// Maximum number of in-flight decodes the reader pipeline holds at
549/// once. While any of these is still running on the blocking pool,
550/// the reader can keep draining the socket and starting new decodes,
551/// up to this cap. Applies per-connection; the resulting decoded
552/// messages are emitted to the Handler in strict WS arrival order
553/// via a `FuturesOrdered` queue — no behavior change versus the
554/// serial loop, just concurrent execution of independent decodes.
555const MAX_IN_FLIGHT_DECODES: usize = 32;
556
557/// Payload size at/above which `decode_message` runs via
558/// `tokio::task::spawn_blocking` instead of inline on the reader task.
559///
560/// `serde_json::from_slice` is CPU-bound with no `.await` points, so
561/// a multi-MB payload can occupy one tokio worker thread for tens of
562/// milliseconds. Offloading to the blocking thread pool above a
563/// threshold keeps the reader task cooperatively yielding — critical
564/// on single-threaded runtimes where the reader shares its worker
565/// with the Handler, user tasks, and timers.
566///
567/// The threshold is chosen so that typical CDP traffic (events,
568/// responses, small evaluates) stays on the inline fast path and
569/// doesn't pay the ~10-30 µs `spawn_blocking` hand-off cost, while
570/// screenshot payloads, wide network events, and huge console
571/// payloads take the offloaded path.
572const LARGE_FRAME_THRESHOLD: usize = 256 * 1024; // 256 KiB
573
574/// Split parts returned by [`Connection::into_async`].
575#[derive(Debug)]
576pub struct AsyncConnection<T: EventMessage> {
577 /// Receive half for decoded CDP messages. Backed by a bounded mpsc
578 /// fed by a dedicated background reader task — decode runs on that
579 /// task, never on the Handler task, so large CDP responses (multi-MB
580 /// screenshots, huge event payloads) cannot stall the Handler's
581 /// event loop.
582 pub reader: WsReader<T>,
583 /// Sender half for submitting outgoing CDP commands.
584 pub cmd_tx: mpsc::Sender<MethodCall>,
585 /// Handle to the background writer task.
586 pub writer_handle: tokio::task::JoinHandle<Result<()>>,
587 /// Handle to the background reader task (reads + decodes WS frames).
588 pub reader_handle: tokio::task::JoinHandle<()>,
589 /// Next command-call-id counter (continue numbering from where Connection left off).
590 pub next_id: usize,
591}
592
593impl<T: EventMessage + Unpin + Send + 'static> Connection<T> {
594 /// Consume the connection and split into a background reader + writer
595 /// pair, exposing the Handler-facing ends via `AsyncConnection`.
596 ///
597 /// Two `tokio::spawn`'d tasks are created:
598 ///
599 /// * `ws_write_loop` — batches outgoing commands and flushes them in
600 /// one write per wakeup.
601 /// * `ws_read_loop` — reads WS frames, decodes them to typed
602 /// `Message<T>`, and forwards them via a bounded mpsc to the
603 /// Handler. Ping/pong/malformed frames are skipped on this task
604 /// and never reach the Handler. Large-message decode (SerDe CPU
605 /// work) runs here, **not** on the Handler task, so the Handler's
606 /// poll loop never stalls for tens of milliseconds on a 10 MB
607 /// screenshot response.
608 ///
609 /// The design uses only `tokio::spawn` (cooperative async) — no
610 /// `spawn_blocking` or blocking thread-pool — so it scales with the
611 /// tokio runtime's worker threads on multi-threaded runtimes, and
612 /// interleaves cleanly with the Handler task on single-threaded
613 /// runtimes.
614 pub fn into_async(self) -> AsyncConnection<T> {
615 let (ws_sink, ws_stream) = self.ws.split();
616 let (cmd_tx, cmd_rx) = mpsc::channel(WS_CMD_CHANNEL_CAPACITY);
617 let (msg_tx, msg_rx) = mpsc::channel::<Result<Box<Message<T>>>>(WS_READ_CHANNEL_CAPACITY);
618
619 // Replay any commands queued via `submit_command` before the
620 // split — most notably the boot `Target.setDiscoverTargets`
621 // pushed by `Handler::new`. Without this, real Chrome never
622 // emits `Target.targetCreated` and `new_page` hangs forever.
623 // Capacity is `WS_CMD_CHANNEL_CAPACITY`, so the boot batch fits
624 // easily — `try_send` would only fail in a pathological case
625 // and we'd lose those commands either way.
626 for call in self.pending_commands {
627 let _ = cmd_tx.try_send(call);
628 }
629
630 let writer_handle = tokio::spawn(ws_write_loop(ws_sink, cmd_rx));
631 let reader_handle = tokio::spawn(ws_read_loop::<T, _>(ws_stream, msg_tx));
632
633 let reader = WsReader {
634 rx: msg_rx,
635 _marker: PhantomData,
636 };
637
638 AsyncConnection {
639 reader,
640 cmd_tx,
641 writer_handle,
642 reader_handle,
643 next_id: self.next_id,
644 }
645 }
646}
647
648/// An entry in the reader's decode pipeline.
649///
650/// Small frames have been decoded inline on the reader task and sit
651/// in `Ready(Some(result))` waiting their turn to emit — zero
652/// allocation beyond the `Option`. Large frames were offloaded to
653/// `tokio::task::spawn_blocking`, so their entry is the
654/// corresponding `JoinHandle`.
655///
656/// A single concrete enum means `FuturesOrdered<InFlightDecode<T>>`
657/// can hold either kind without `Box<dyn Future>`, keeping the
658/// pipeline cost-proportional to the workload.
659enum InFlightDecode<T: EventMessage + Send + 'static> {
660 /// Small-frame fast path: already decoded inline. `take()`'d
661 /// exactly once when `FuturesOrdered` first polls it to Ready.
662 Ready(Option<Result<Box<Message<T>>>>),
663 /// Large-frame path: decoding on the blocking thread pool.
664 Blocking(tokio::task::JoinHandle<Result<Box<Message<T>>>>),
665}
666
667impl<T: EventMessage + Send + 'static> Future for InFlightDecode<T> {
668 type Output = Result<Box<Message<T>>>;
669
670 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
671 // Safety: both variants are structurally pin-agnostic —
672 // `Option<Result<..>>` is `Unpin`, and `tokio::task::JoinHandle`
673 // is documented as `Unpin`. So we can project out a `&mut`
674 // without unsafe.
675 match self.get_mut() {
676 InFlightDecode::Ready(slot) => Poll::Ready(
677 slot.take()
678 .expect("InFlightDecode::Ready polled after completion"),
679 ),
680 InFlightDecode::Blocking(handle) => match Pin::new(handle).poll(cx) {
681 Poll::Ready(Ok(res)) => Poll::Ready(res),
682 Poll::Ready(Err(join_err)) => Poll::Ready(Err(CdpError::msg(format!(
683 "WS decode blocking task join error: {join_err}"
684 )))),
685 Poll::Pending => Poll::Pending,
686 },
687 }
688 }
689}
690
691/// Emit a single decoded-frame result to the Handler, logging parse
692/// errors. Returns `true` if the channel is still open, `false` if
693/// the Handler has dropped the receiver (caller should exit).
694async fn emit_decoded<T>(
695 tx: &mpsc::Sender<Result<Box<Message<T>>>>,
696 res: Result<Box<Message<T>>>,
697) -> bool
698where
699 T: EventMessage + Send + 'static,
700{
701 match res {
702 Ok(msg) => tx.send(Ok(msg)).await.is_ok(),
703 Err(err) => {
704 tracing::debug!(
705 target: "chromiumoxide::conn::raw_ws::parse_errors",
706 "Dropping malformed WS frame: {err}",
707 );
708 true
709 }
710 }
711}
712
713/// Background task that reads frames from the WebSocket, decodes them to
714/// typed CDP `Message<T>`, and forwards them to the Handler over a
715/// bounded mpsc.
716///
717/// Runs on a `tokio::spawn`'d task. Small-to-medium frames are
718/// decoded inline (fast path); payloads at or above
719/// [`LARGE_FRAME_THRESHOLD`] are offloaded to `spawn_blocking` so
720/// multi-MB deserialization doesn't monopolise a tokio worker
721/// thread — especially important on single-threaded runtimes where
722/// the reader, Handler, and user tasks share the same worker.
723///
724/// Flow per frame:
725///
726/// * `Text` / `Binary` → [`decode_ws_frame`]; decoded `Ok(msg)` is
727/// sent to the Handler. Decode errors are logged and the frame is
728/// dropped (same behavior as the legacy inline decode path).
729/// * `Close` → loop exits cleanly, dropping `tx`. The Handler's
730/// `next_message().await` returns `None` on the next call.
731/// * `Ping` / `Pong` / unexpected frame types → skipped silently; they
732/// never cross the channel to the Handler.
733/// * Transport error → forwarded as `Err(CdpError::Ws(..))`, then the
734/// loop exits (the WS half is considered dead after an error).
735///
736/// Back-pressure: the outbound `tx` is bounded. If the Handler is busy
737/// and the channel fills, `tx.send(..).await` parks this task, which
738/// stops draining the WS socket. TCP flow control then applies
739/// back-pressure to Chrome instead of letting memory grow without bound.
740async fn ws_read_loop<T, S>(mut stream: S, tx: mpsc::Sender<Result<Box<Message<T>>>>)
741where
742 T: EventMessage + Send + 'static,
743 S: Stream<Item = std::result::Result<WsMessage, tokio_tungstenite::tungstenite::Error>> + Unpin,
744{
745 // Pipeline of decodes in strict arrival order. Small-frame decodes
746 // are produced inline (zero allocation, borrowing the frame body);
747 // large-frame decodes are offloaded to `spawn_blocking`. Both
748 // variants share a single concrete `InFlightDecode<T>` so the
749 // queue avoids `Box<dyn Future>` overhead.
750 let mut in_flight: FuturesOrdered<InFlightDecode<T>> = FuturesOrdered::new();
751
752 // Shutdown state. When the stream signals `Close`, transport
753 // error, or end-of-stream, we stop reading new frames but keep
754 // running the select loop so the emit arm can flush any still
755 // in-flight decodes *interleaved with* whatever else the runtime
756 // is doing. A pending transport error is surfaced to the Handler
757 // only after the in-order flush completes.
758 let mut stream_terminated = false;
759 let mut pending_err: Option<CdpError> = None;
760
761 loop {
762 tokio::select! {
763 // Bias: emit already-ready decodes before reading more
764 // frames. Keeps the pipeline small in the steady state
765 // while still allowing concurrency under burst, and —
766 // critically during shutdown — drains the pipeline one
767 // ready item at a time inside the select loop instead
768 // of blocking in a dedicated drain helper.
769 biased;
770
771 // Emit the head of the pipeline as soon as it is ready.
772 // `FuturesOrdered::next` preserves submit order, so
773 // downstream delivery is byte-identical to the serial
774 // loop's ordering guarantee.
775 Some(res) = in_flight.next(), if !in_flight.is_empty() => {
776 if !emit_decoded(&tx, res).await {
777 return;
778 }
779 }
780
781 // Read the next frame if the pipeline has capacity and
782 // the stream hasn't terminated. Disabled once the stream
783 // signals end (Close / None / Err) so subsequent loop
784 // iterations only do emit work.
785 maybe_frame = stream.next(),
786 if !stream_terminated && in_flight.len() < MAX_IN_FLIGHT_DECODES =>
787 {
788 match maybe_frame {
789 Some(Ok(WsMessage::Text(text))) => {
790 // Zero-copy enqueue. The small-frame fast
791 // path decodes inline *now* (borrowing
792 // `text`, keeping the `raw_text_for_logging`
793 // preview); the large-frame path moves the
794 // `Utf8Bytes` (`Send + 'static`) directly
795 // into `spawn_blocking` without an
796 // intermediate allocation.
797 if text.len() >= LARGE_FRAME_THRESHOLD {
798 in_flight.push_back(InFlightDecode::Blocking(
799 tokio::task::spawn_blocking(move || {
800 decode_message::<T>(text.as_bytes(), None)
801 }),
802 ));
803 } else {
804 let res = decode_message::<T>(text.as_bytes(), Some(&text));
805 in_flight.push_back(InFlightDecode::Ready(Some(res)));
806 }
807 }
808 Some(Ok(WsMessage::Binary(buf))) => {
809 // Same shape as Text: move `Bytes`
810 // (`Send + 'static`) into `spawn_blocking`
811 // for large payloads, decode inline for
812 // small ones.
813 if buf.len() >= LARGE_FRAME_THRESHOLD {
814 in_flight.push_back(InFlightDecode::Blocking(
815 tokio::task::spawn_blocking(move || {
816 decode_message::<T>(&buf, None)
817 }),
818 ));
819 } else {
820 let res = decode_message::<T>(&buf, None);
821 in_flight.push_back(InFlightDecode::Ready(Some(res)));
822 }
823 }
824 Some(Ok(WsMessage::Close(_))) => {
825 stream_terminated = true;
826 }
827 Some(Ok(WsMessage::Ping(_))) | Some(Ok(WsMessage::Pong(_))) => {}
828 Some(Ok(msg)) => {
829 tracing::debug!(
830 target: "chromiumoxide::conn::raw_ws::parse_errors",
831 "Unexpected WS message type: {:?}",
832 msg
833 );
834 }
835 Some(Err(err)) => {
836 // Defer the error until after the already
837 // in-flight decodes have emitted — preserves
838 // the ordering contract that callers see
839 // frames up to the failure point before the
840 // error itself.
841 stream_terminated = true;
842 pending_err = Some(CdpError::Ws(err));
843 }
844 None => {
845 // Stream ended (connection closed without a
846 // `Close` frame). No more input, but
847 // in_flight may still hold pending decodes.
848 stream_terminated = true;
849 }
850 }
851 }
852
853 // Both arms disabled: `in_flight` is empty AND
854 // `stream_terminated`. We have nothing more to do.
855 else => {
856 break;
857 }
858 }
859 }
860
861 if let Some(err) = pending_err {
862 let _ = tx.send(Err(err)).await;
863 }
864}
865
866/// Background task that batches and flushes outgoing CDP commands.
867async fn ws_write_loop(
868 mut sink: SplitSink<WebSocketStream<ConnectStream>, WsMessage>,
869 mut rx: mpsc::Receiver<MethodCall>,
870) -> Result<()> {
871 while let Some(call) = rx.recv().await {
872 let msg = crate::serde_json::to_string(&call)?;
873 sink.feed(WsMessage::Text(msg.into()))
874 .await
875 .map_err(CdpError::Ws)?;
876
877 // Batch: drain all buffered commands without waiting.
878 while let Ok(call) = rx.try_recv() {
879 let msg = crate::serde_json::to_string(&call)?;
880 sink.feed(WsMessage::Text(msg.into()))
881 .await
882 .map_err(CdpError::Ws)?;
883 }
884
885 // Flush the entire batch in one write.
886 sink.flush().await.map_err(CdpError::Ws)?;
887 }
888
889 // Cmd channel closed → the Handler is shutting down. Send a graceful
890 // WebSocket Close frame so the remote endpoint (esp. for
891 // `Browser::connect()` to a remote DevTools URL, where there is no
892 // child process whose exit closes the socket) tears the connection
893 // down promptly instead of waiting for an idle timeout. Errors are
894 // expected during shutdown (e.g. `AlreadyClosed` if Chrome closed
895 // first) and are intentionally ignored.
896 let _ = sink.close().await;
897 Ok(())
898}
899
900/// Handler-facing read half of the split WebSocket connection.
901///
902/// Decoded CDP messages are produced by a dedicated background task
903/// (see [`ws_read_loop`]) and forwarded over a bounded mpsc. `WsReader`
904/// itself is a thin `Receiver` wrapper — calling `next_message()` does
905/// a single `rx.recv().await` with no per-message decoding work on the
906/// caller's task. This keeps the Handler's poll loop free of CPU-bound
907/// deserialize time, which matters for large (multi-MB) CDP responses
908/// such as screenshots and wide-header network events.
909#[derive(Debug)]
910pub struct WsReader<T: EventMessage> {
911 rx: mpsc::Receiver<Result<Box<Message<T>>>>,
912 _marker: PhantomData<T>,
913}
914
915impl<T: EventMessage + Unpin> WsReader<T> {
916 /// Read the next CDP message from the WebSocket.
917 ///
918 /// Returns `None` when the background reader task has exited
919 /// (connection closed or sender dropped). This call does only a
920 /// channel `recv` — the actual WS read + JSON decode happens on
921 /// the background `ws_read_loop` task.
922 pub async fn next_message(&mut self) -> Option<Result<Box<Message<T>>>> {
923 self.rx.recv().await
924 }
925}
926
927impl<T: EventMessage + Unpin> Stream for Connection<T> {
928 type Item = Result<Box<Message<T>>>;
929
930 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
931 let pin = self.get_mut();
932
933 // Send and flush outgoing messages
934 if let Err(err) = pin.start_send_next(cx) {
935 return Poll::Ready(Some(Err(err)));
936 }
937
938 // Read from the websocket, skipping non-data frames (pings,
939 // pongs, malformed messages) without yielding back to the
940 // executor. This avoids a full round-trip per skipped frame.
941 //
942 // Cap consecutive skips so a flood of non-data frames (many
943 // pings, malformed/unexpected types) cannot starve the
944 // runtime — yield Pending after `MAX_SKIPS_PER_POLL` and
945 // self-wake so we resume on the next tick.
946 const MAX_SKIPS_PER_POLL: u32 = 16;
947 let mut skips: u32 = 0;
948 loop {
949 match ready!(pin.ws.poll_next_unpin(cx)) {
950 Some(Ok(WsMessage::Text(text))) => {
951 match decode_message::<T>(text.as_bytes(), Some(&text)) {
952 Ok(msg) => return Poll::Ready(Some(Ok(msg))),
953 Err(err) => {
954 tracing::debug!(
955 target: "chromiumoxide::conn::raw_ws::parse_errors",
956 "Dropping malformed text WS frame: {err}",
957 );
958 skips += 1;
959 }
960 }
961 }
962 Some(Ok(WsMessage::Binary(buf))) => match decode_message::<T>(&buf, None) {
963 Ok(msg) => return Poll::Ready(Some(Ok(msg))),
964 Err(err) => {
965 tracing::debug!(
966 target: "chromiumoxide::conn::raw_ws::parse_errors",
967 "Dropping malformed binary WS frame: {err}",
968 );
969 skips += 1;
970 }
971 },
972 Some(Ok(WsMessage::Close(_))) => return Poll::Ready(None),
973 Some(Ok(WsMessage::Ping(_))) | Some(Ok(WsMessage::Pong(_))) => {
974 skips += 1;
975 }
976 Some(Ok(msg)) => {
977 tracing::debug!(
978 target: "chromiumoxide::conn::raw_ws::parse_errors",
979 "Unexpected WS message type: {:?}",
980 msg
981 );
982 skips += 1;
983 }
984 Some(Err(err)) => return Poll::Ready(Some(Err(CdpError::Ws(err)))),
985 None => return Poll::Ready(None),
986 }
987
988 if skips >= MAX_SKIPS_PER_POLL {
989 cx.waker().wake_by_ref();
990 return Poll::Pending;
991 }
992 }
993 }
994}
995
996/// Shared decode path for both text and binary WS frames.
997/// `raw_text_for_logging` is only provided for textual frames so we can log the original
998/// payload on parse failure if desired.
999#[cfg(not(feature = "serde_stacker"))]
1000fn decode_message<T: EventMessage>(
1001 bytes: &[u8],
1002 raw_text_for_logging: Option<&str>,
1003) -> Result<Box<Message<T>>> {
1004 match serde_json::from_slice::<Box<Message<T>>>(bytes) {
1005 Ok(msg) => {
1006 tracing::trace!("Received {:?}", msg);
1007 Ok(msg)
1008 }
1009 Err(err) => {
1010 if let Some(txt) = raw_text_for_logging {
1011 let preview = &txt[..txt.len().min(512)];
1012 tracing::debug!(
1013 target: "chromiumoxide::conn::raw_ws::parse_errors",
1014 msg_len = txt.len(),
1015 "Skipping unrecognized WS message {err} preview={preview}",
1016 );
1017 } else {
1018 tracing::debug!(
1019 target: "chromiumoxide::conn::raw_ws::parse_errors",
1020 "Skipping unrecognized binary WS message {err}",
1021 );
1022 }
1023 Err(err.into())
1024 }
1025 }
1026}
1027
1028/// Shared decode path for both text and binary WS frames.
1029/// `raw_text_for_logging` is only provided for textual frames so we can log the original
1030/// payload on parse failure if desired.
1031#[cfg(feature = "serde_stacker")]
1032fn decode_message<T: EventMessage>(
1033 bytes: &[u8],
1034 raw_text_for_logging: Option<&str>,
1035) -> Result<Box<Message<T>>> {
1036 use serde::Deserialize;
1037 let mut de = serde_json::Deserializer::from_slice(bytes);
1038
1039 de.disable_recursion_limit();
1040
1041 let de = serde_stacker::Deserializer::new(&mut de);
1042
1043 match Box::<Message<T>>::deserialize(de) {
1044 Ok(msg) => {
1045 tracing::trace!("Received {:?}", msg);
1046 Ok(msg)
1047 }
1048 Err(err) => {
1049 if let Some(txt) = raw_text_for_logging {
1050 let preview = &txt[..txt.len().min(512)];
1051 tracing::debug!(
1052 target: "chromiumoxide::conn::raw_ws::parse_errors",
1053 msg_len = txt.len(),
1054 "Skipping unrecognized WS message {err} preview={preview}",
1055 );
1056 } else {
1057 tracing::debug!(
1058 target: "chromiumoxide::conn::raw_ws::parse_errors",
1059 "Skipping unrecognized binary WS message {err}",
1060 );
1061 }
1062 Err(err.into())
1063 }
1064 }
1065}
1066
1067#[cfg(test)]
1068mod ws_read_loop_tests {
1069 //! Unit tests for the `ws_read_loop` background reader task.
1070 //!
1071 //! These tests feed a synthetic `Stream<Item = Result<WsMessage, _>>`
1072 //! into `ws_read_loop` — no real WebSocket, no Chrome — and observe
1073 //! what comes out the other side of the mpsc channel.
1074 //!
1075 //! The properties under test are the ones that make the reader-task
1076 //! decoupling safe: FIFO ordering, no-deadlock on a bounded channel
1077 //! under back-pressure, silent drop of non-data frames, graceful
1078 //! transport-error propagation, and clean exit on `Close`.
1079 //!
1080 //! The typed events are `chromiumoxide_cdp::cdp::CdpEventMessage` —
1081 //! the same instantiation the real Handler uses — so these tests
1082 //! exercise the actual decode path (`serde_json::from_slice`), not
1083 //! a simplified fake.
1084 use super::*;
1085 use chromiumoxide_cdp::cdp::CdpEventMessage;
1086 use chromiumoxide_types::CallId;
1087 use futures_util::stream;
1088 use tokio::sync::mpsc;
1089 use tokio_tungstenite::tungstenite::Message as WsMessage;
1090
1091 /// Build a CDP `Response` WS frame as text — the smallest valid CDP
1092 /// message. `id` tags the frame for ordering assertions.
1093 fn response_frame(id: u64) -> WsMessage {
1094 WsMessage::Text(
1095 format!(r#"{{"id":{id},"result":{{"ok":true}}}}"#)
1096 .to_string()
1097 .into(),
1098 )
1099 }
1100
1101 /// Build a frame far larger than a typical socket chunk, to exercise
1102 /// the "large message" path that motivated this refactor. The blob
1103 /// field pushes serde_json through a big allocation even though the
1104 /// envelope is tiny.
1105 fn large_response_frame(id: u64, blob_bytes: usize) -> WsMessage {
1106 let blob = "x".repeat(blob_bytes);
1107 WsMessage::Text(
1108 format!(r#"{{"id":{id},"result":{{"blob":"{blob}"}}}}"#)
1109 .to_string()
1110 .into(),
1111 )
1112 }
1113
1114 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1115 async fn forwards_messages_in_stream_order() {
1116 let frames = vec![
1117 Ok(response_frame(1)),
1118 Ok(response_frame(2)),
1119 Ok(response_frame(3)),
1120 ];
1121 let stream = stream::iter(frames);
1122 let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(8);
1123 let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
1124
1125 for expected in [1u64, 2, 3] {
1126 let msg = rx.recv().await.expect("msg").expect("decode ok");
1127 if let Message::Response(resp) = *msg {
1128 assert_eq!(resp.id, CallId::new(expected as usize));
1129 } else {
1130 panic!("expected Response");
1131 }
1132 }
1133 assert!(rx.recv().await.is_none(), "channel must close on EOF");
1134 task.await.expect("reader task join");
1135 }
1136
1137 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1138 async fn pings_and_pongs_never_reach_the_handler() {
1139 let frames = vec![
1140 Ok(WsMessage::Ping(vec![1, 2, 3].into())),
1141 Ok(response_frame(7)),
1142 Ok(WsMessage::Pong(vec![].into())),
1143 Ok(response_frame(8)),
1144 ];
1145 let stream = stream::iter(frames);
1146 let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(8);
1147 let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
1148
1149 for expected in [7u64, 8] {
1150 let msg = rx.recv().await.expect("msg").expect("decode ok");
1151 if let Message::Response(resp) = *msg {
1152 assert_eq!(resp.id, CallId::new(expected as usize));
1153 }
1154 }
1155 assert!(rx.recv().await.is_none());
1156 task.await.expect("reader task join");
1157 }
1158
1159 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1160 async fn malformed_frames_do_not_block_subsequent_valid_frames() {
1161 let frames = vec![
1162 Ok(WsMessage::Text("{not valid json".to_string().into())),
1163 Ok(response_frame(42)),
1164 ];
1165 let stream = stream::iter(frames);
1166 let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(8);
1167 let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
1168
1169 let msg = rx.recv().await.expect("msg").expect("decode ok");
1170 if let Message::Response(resp) = *msg {
1171 assert_eq!(resp.id, CallId::new(42));
1172 }
1173 assert!(rx.recv().await.is_none());
1174 task.await.expect("reader task join");
1175 }
1176
1177 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1178 async fn close_frame_terminates_the_reader() {
1179 let frames = vec![
1180 Ok(response_frame(1)),
1181 Ok(WsMessage::Close(None)),
1182 Ok(response_frame(2)), // unreachable after Close
1183 ];
1184 let stream = stream::iter(frames);
1185 let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(8);
1186 let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
1187
1188 let msg = rx.recv().await.expect("msg").expect("decode ok");
1189 if let Message::Response(resp) = *msg {
1190 assert_eq!(resp.id, CallId::new(1));
1191 }
1192 assert!(
1193 rx.recv().await.is_none(),
1194 "reader must exit on Close; frames after Close must not appear"
1195 );
1196 task.await.expect("reader task join");
1197 }
1198
1199 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1200 async fn transport_error_is_forwarded_once_then_reader_exits() {
1201 let frames = vec![
1202 Ok(response_frame(1)),
1203 Err(tokio_tungstenite::tungstenite::Error::ConnectionClosed),
1204 Ok(response_frame(2)),
1205 ];
1206 let stream = stream::iter(frames);
1207 let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(8);
1208 let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
1209
1210 let msg = rx.recv().await.expect("msg").expect("ok");
1211 assert!(matches!(*msg, Message::Response(_)));
1212 match rx.recv().await {
1213 Some(Err(CdpError::Ws(_))) => {}
1214 other => panic!("expected forwarded Ws error, got {other:?}"),
1215 }
1216 assert!(rx.recv().await.is_none());
1217 task.await.expect("reader task join");
1218 }
1219
1220 /// Back-pressure property: with the smallest possible channel and
1221 /// many frames, the reader task awaits capacity after each send and
1222 /// never deadlocks. This is the core "no deadlock" proof for the
1223 /// new design — if the reader held anything across its `.await` that
1224 /// the consumer needed, the consumer's `recv().await` would block
1225 /// forever. Completion under a 5s watchdog proves it doesn't.
1226 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1227 async fn bounded_channel_does_not_deadlock_under_backpressure() {
1228 const N: u64 = 512;
1229 let frames: Vec<_> = (1..=N).map(|id| Ok(response_frame(id))).collect();
1230 let stream = stream::iter(frames);
1231
1232 let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(1);
1233 let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
1234
1235 let deadline = std::time::Duration::from_secs(5);
1236 let collected = tokio::time::timeout(deadline, async {
1237 let mut seen = 0u64;
1238 while let Some(frame) = rx.recv().await {
1239 let msg = frame.expect("decode ok");
1240 if let Message::Response(resp) = *msg {
1241 seen += 1;
1242 assert_eq!(
1243 resp.id,
1244 CallId::new(seen as usize),
1245 "back-pressure must preserve FIFO order"
1246 );
1247 }
1248 }
1249 seen
1250 })
1251 .await
1252 .expect("reader must make forward progress despite cap-1 back-pressure");
1253
1254 assert_eq!(collected, N, "all frames must arrive");
1255 task.await.expect("reader task join");
1256 }
1257
1258 /// Large message (>1 MB) is decoded correctly on the background
1259 /// task. This is the specific scenario the reader-task refactor
1260 /// was built for — we don't measure time here (benches cover that),
1261 /// we just prove the end-to-end path works without corruption or
1262 /// deadlock.
1263 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1264 async fn large_message_decodes_without_corruption() {
1265 let big = 2 * 1024 * 1024; // 2 MB payload
1266 let frames = vec![Ok(large_response_frame(100, big)), Ok(response_frame(101))];
1267 let stream = stream::iter(frames);
1268 let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(4);
1269 let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
1270
1271 let first = rx.recv().await.expect("msg").expect("ok");
1272 if let Message::Response(resp) = *first {
1273 assert_eq!(resp.id, CallId::new(100));
1274 }
1275 let second = rx.recv().await.expect("msg").expect("ok");
1276 if let Message::Response(resp) = *second {
1277 assert_eq!(resp.id, CallId::new(101));
1278 }
1279 assert!(rx.recv().await.is_none());
1280 task.await.expect("reader task join");
1281 }
1282
1283 /// FIFO ordering under the pipelined reader when large-frame
1284 /// decodes run in parallel via `spawn_blocking`.
1285 ///
1286 /// This test submits an interleaved sequence of large and small
1287 /// frames. Large frames take the `spawn_blocking` path (decode
1288 /// on the blocking pool, variable completion order); small
1289 /// frames take the inline path (decode immediately). The
1290 /// pipeline's `FuturesOrdered` queue must emit them to the
1291 /// Handler in strict arrival order regardless of which
1292 /// blocking-pool thread finishes first.
1293 ///
1294 /// If the ordering guarantee were ever broken — e.g. by
1295 /// accidentally swapping `FuturesOrdered` for `FuturesUnordered`
1296 /// — id sequence checks here would catch it immediately.
1297 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1298 async fn pipelined_large_and_small_frames_keep_fifo_order() {
1299 let big = 2 * 1024 * 1024; // 2 MB payload — forces spawn_blocking
1300 let frames = vec![
1301 Ok(large_response_frame(1, big)),
1302 Ok(response_frame(2)),
1303 Ok(response_frame(3)),
1304 Ok(large_response_frame(4, big)),
1305 Ok(response_frame(5)),
1306 Ok(large_response_frame(6, big)),
1307 Ok(response_frame(7)),
1308 Ok(response_frame(8)),
1309 ];
1310 let expected: Vec<usize> = (1..=8).collect();
1311
1312 let stream = stream::iter(frames);
1313 let (tx, mut rx) = mpsc::channel::<Result<Box<Message<CdpEventMessage>>>>(16);
1314 let task = tokio::spawn(ws_read_loop::<CdpEventMessage, _>(stream, tx));
1315
1316 let deadline = std::time::Duration::from_secs(10);
1317 let observed = tokio::time::timeout(deadline, async {
1318 let mut ids = Vec::with_capacity(expected.len());
1319 while let Some(frame) = rx.recv().await {
1320 let msg = frame.expect("decode ok");
1321 if let Message::Response(resp) = *msg {
1322 ids.push(CallId::new(ids.len() + 1));
1323 assert_eq!(
1324 resp.id,
1325 *ids.last().unwrap(),
1326 "pipelined reader must emit frames in strict arrival order \
1327 regardless of per-frame decode latency"
1328 );
1329 }
1330 }
1331 ids
1332 })
1333 .await
1334 .expect("pipelined reader should make forward progress within 10s");
1335
1336 assert_eq!(
1337 observed.len(),
1338 expected.len(),
1339 "all {} frames must reach the Handler",
1340 expected.len()
1341 );
1342 task.await.expect("reader task join");
1343 }
1344}
1345
1346#[cfg(test)]
1347mod connect_headers_tests {
1348 //! Unit tests for `ConnectHeaders`: the size accounting, replacement,
1349 //! the reserved-name guard, and the promise that an empty set leaves an
1350 //! upgrade request exactly as tungstenite built it.
1351
1352 use super::ConnectHeaderError;
1353 use super::ConnectHeaders;
1354 use tokio_tungstenite::tungstenite::client::IntoClientRequest;
1355 use tokio_tungstenite::tungstenite::http::HeaderMap;
1356
1357 #[test]
1358 fn empty_apply_is_a_no_op() {
1359 // Two `into_client_request` calls differ in `Sec-WebSocket-Key`, so
1360 // snapshot one request and compare it with itself after `apply`.
1361 let url = "ws://127.0.0.1:9222/devtools/browser/x";
1362 let mut req = url.into_client_request().unwrap();
1363 let before = req.headers().clone();
1364 ConnectHeaders::new().apply(&mut req);
1365 assert_eq!(&before, req.headers());
1366 }
1367
1368 #[test]
1369 fn apply_adds_and_replaces_by_name() {
1370 let url = "ws://127.0.0.1:9222/devtools/browser/x";
1371 let mut req = url.into_client_request().unwrap();
1372 let set = ConnectHeaders::new()
1373 .with("x-a", "1")
1374 .unwrap()
1375 .with("x-a", "2")
1376 .unwrap()
1377 .with("x-b", "3")
1378 .unwrap();
1379 assert_eq!(set.len(), 2);
1380 set.apply(&mut req);
1381 assert_eq!(req.headers().get_all("x-a").iter().count(), 1);
1382 assert_eq!(req.headers()["x-a"], "2");
1383 assert_eq!(req.headers()["x-b"], "3");
1384 }
1385
1386 #[test]
1387 fn total_bytes_tracks_replacement() {
1388 let mut set = ConnectHeaders::new();
1389 set.insert("x-a", "12345").unwrap();
1390 // "x-a: 12345\r\n"
1391 assert_eq!(set.total_bytes(), 3 + 2 + 5 + 2);
1392 set.insert("x-a", "1").unwrap();
1393 assert_eq!(set.total_bytes(), 3 + 2 + 1 + 2);
1394 set.insert("x-b", "").unwrap();
1395 assert_eq!(set.total_bytes(), 8 + 7);
1396 }
1397
1398 #[test]
1399 fn limit_counts_the_whole_set() {
1400 let half = "x".repeat(ConnectHeaders::MAX_TOTAL_BYTES / 2);
1401 let mut set = ConnectHeaders::new();
1402 set.insert("x-a", half.as_str()).unwrap();
1403 let err = set.insert("x-b", half.as_str()).unwrap_err();
1404 assert!(matches!(err, ConnectHeaderError::TooLarge { .. }), "{err}");
1405 // A refused insert leaves the set as it was.
1406 assert_eq!(set.len(), 1);
1407 // Replacing the big one with a small one frees the room back up.
1408 set.insert("x-a", "1").unwrap();
1409 set.insert("x-b", half.as_str()).unwrap();
1410 }
1411
1412 #[test]
1413 fn reserved_names_are_case_insensitive() {
1414 for name in ["Host", "CONNECTION", "Upgrade", "Sec-WebSocket-Protocol"] {
1415 let err = ConnectHeaders::new().with(name, "x").unwrap_err();
1416 assert!(
1417 matches!(err, ConnectHeaderError::Reserved(_)),
1418 "{name}: {err}"
1419 );
1420 }
1421 }
1422
1423 #[test]
1424 fn obs_text_bytes_are_refused() {
1425 // `HeaderValue::from_bytes` accepts 0x80..=0xFF; the handshake writer
1426 // does not, so the set must not either.
1427 let value = tokio_tungstenite::tungstenite::http::HeaderValue::from_bytes(&[0xE9]).unwrap();
1428 let err = ConnectHeaders::new().with("x-a", value).unwrap_err();
1429 assert!(matches!(err, ConnectHeaderError::NonAscii(_)), "{err}");
1430 }
1431
1432 #[test]
1433 fn header_map_conversion_validates_every_entry() {
1434 let mut map = HeaderMap::new();
1435 map.insert("x-a", "1".parse().unwrap());
1436 map.insert("host", "h".parse().unwrap());
1437 let err = ConnectHeaders::try_from(map).unwrap_err();
1438 assert!(matches!(err, ConnectHeaderError::Reserved(_)), "{err}");
1439 }
1440}