Skip to main content

playwright_cdp/cdp/
transport.rs

1//! WebSocket transport for CDP.
2//!
3//! One JSON object per WebSocket text frame — exactly CDP's wire format. This
4//! is a trimmed, self-contained version of `playwright-rust`'s transport (no
5//! driver dependency, no trait gymnastics): `connect` returns a writer half
6//! plus a bounded receiver of parsed JSON values, and spawns the read pump.
7
8use crate::{Error, Result};
9use futures_util::{SinkExt, StreamExt};
10use serde_json::Value;
11use std::collections::HashMap;
12use std::str::FromStr;
13use tokio::net::TcpStream;
14use tokio::sync::mpsc;
15use tokio_tungstenite::tungstenite::client::IntoClientRequest;
16use tokio_tungstenite::tungstenite::http::header::{HeaderName, HeaderValue};
17use tokio_tungstenite::tungstenite::Message as WsMessage;
18use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
19
20/// Capacity of the transport-to-dispatch channel. Bounded so a slow dispatch
21/// loop exerts backpressure on the socket instead of buffering unboundedly.
22pub(crate) const MESSAGE_CHANNEL_CAPACITY: usize = 256;
23
24type WsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
25
26/// Connect to a CDP WebSocket endpoint.
27///
28/// Returns a [`WebSocketWriter`] (the sending half) and a bounded receiver of
29/// incoming parsed JSON messages. A background task pumps the read side.
30pub async fn connect(
31    url: &str,
32    headers: Option<HashMap<String, String>>,
33) -> Result<(WebSocketWriter, mpsc::Receiver<Value>)> {
34    let mut request = url
35        .into_client_request()
36        .map_err(|e| Error::TransportError(format!("invalid CDP endpoint {url}: {e}")))?;
37
38    if let Some(extra) = headers {
39        let h = request.headers_mut();
40        for (k, v) in extra {
41            let name = HeaderName::from_str(&k)
42                .map_err(|e| Error::TransportError(format!("bad header name {k}: {e}")))?;
43            let value = HeaderValue::from_str(&v)
44                .map_err(|e| Error::TransportError(format!("bad header value: {e}")))?;
45            h.insert(name, value);
46        }
47    }
48
49    let (ws_stream, _resp) = tokio_tungstenite::connect_async(request)
50        .await
51        .map_err(|e| Error::ConnectionFailed(format!("WebSocket connect to {url} failed: {e}")))?;
52
53    let (sink, stream) = ws_stream.split();
54    let (tx, rx) = mpsc::channel::<Value>(MESSAGE_CHANNEL_CAPACITY);
55
56    tokio::spawn(pump_reads(stream, tx));
57
58    Ok((WebSocketWriter { sink }, rx))
59}
60
61async fn pump_reads(
62    mut stream: futures_util::stream::SplitStream<WsStream>,
63    tx: mpsc::Sender<Value>,
64) {
65    while let Some(msg_result) = stream.next().await {
66        match msg_result {
67            Ok(WsMessage::Text(text)) => match serde_json::from_str::<Value>(&text) {
68                Ok(v) => {
69                    if tx.send(v).await.is_err() {
70                        break; // dispatch side gone
71                    }
72                }
73                Err(e) => {
74                    tracing::trace!(error = %e, "ignoring non-JSON CDP frame");
75                }
76            },
77            Ok(WsMessage::Binary(_)) => {} // CDP uses Text frames only
78            Ok(WsMessage::Close(_)) | Err(_) => break,
79            // Ping/Pong handled by tungstenite automatically.
80            _ => {}
81        }
82    }
83}
84
85/// Sending half of the CDP WebSocket connection.
86pub struct WebSocketWriter {
87    sink: futures_util::stream::SplitSink<WsStream, WsMessage>,
88}
89
90impl WebSocketWriter {
91    /// Send a JSON value as a single text frame.
92    pub async fn send(&mut self, value: &Value) -> Result<()> {
93        let s = serde_json::to_string(value)?;
94        self.sink
95            .send(WsMessage::Text(s.into()))
96            .await
97            .map_err(|e| Error::TransportError(format!("WebSocket send failed: {e}")))?;
98        Ok(())
99    }
100
101    /// Close the socket cleanly.
102    pub async fn close(&mut self) -> Result<()> {
103        let _ = self
104            .sink
105            .close()
106            .await
107            .map_err(|e| Error::TransportError(format!("WebSocket close failed: {e}")));
108        Ok(())
109    }
110}