Skip to main content

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}