drission/cdp/
websocket.rs

1//! CDP 后端的 **WebSocket 帧监听** [`ChromiumWsListener`](对齐 camoufox `WsListener`)。
2//!
3//! 基于 CDP `Network.webSocketFrameSent`/`Received` 事件(需 `Network.enable`,与网络监听同域,
4//! **不**涉及 `Runtime.enable`,反检测友好)。文本帧(opcode 1)`data` 为原文,其余帧为 base64。
5
6use std::collections::{HashMap, VecDeque};
7use std::sync::Arc;
8use std::time::Duration;
9
10use serde_json::{Value, json};
11use tokio::sync::Mutex;
12use tokio::sync::broadcast::error::RecvError;
13use tokio::time::{Instant, sleep};
14
15use crate::Result;
16use crate::cdp::core::{CdpCore, EventBuf};
17use crate::protocol::Connection;
18
19/// 帧方向。
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum WsDirection {
22    Sent,
23    Received,
24}
25
26/// 一个 WebSocket 帧(对齐 camoufox `WsMessage`)。
27#[derive(Debug, Clone)]
28pub struct WsMessage {
29    pub direction: WsDirection,
30    pub url: String,
31    pub opcode: u8,
32    /// 文本帧为原文;其余帧为 base64。
33    pub data: String,
34}
35
36impl WsMessage {
37    pub fn is_text(&self) -> bool {
38        self.opcode == 1
39    }
40    pub fn is_binary(&self) -> bool {
41        self.opcode == 2
42    }
43    /// 文本(文本帧原样;其余帧尝试 base64 解码为 UTF-8)。
44    pub fn text(&self) -> String {
45        if self.is_text() {
46            self.data.clone()
47        } else {
48            crate::util::base64_decode(&self.data)
49                .and_then(|b| String::from_utf8(b).ok())
50                .unwrap_or_default()
51        }
52    }
53    /// 原始字节(文本帧=UTF-8 字节;其余=base64 解码)。
54    pub fn bytes(&self) -> Vec<u8> {
55        if self.is_text() {
56            self.data.clone().into_bytes()
57        } else {
58            crate::util::base64_decode(&self.data).unwrap_or_default()
59        }
60    }
61    /// 把负载当 JSON 解析。
62    pub fn json(&self) -> Value {
63        serde_json::from_str(&self.text()).unwrap_or(Value::Null)
64    }
65}
66
67/// WS 监听过滤(对齐 camoufox `WsFilter`)。
68#[derive(Debug, Clone, Default)]
69pub struct WsFilter {
70    pub url_contains: Option<String>,
71    pub direction: Option<WsDirection>,
72    /// 是否保留 ping/pong/close 等控制帧(默认 false:只留 text/binary 数据帧)。
73    pub with_control: bool,
74}
75
76impl WsFilter {
77    pub fn new() -> Self {
78        Self::default()
79    }
80    pub fn url_contains(mut self, s: impl Into<String>) -> Self {
81        self.url_contains = Some(s.into());
82        self
83    }
84    pub fn direction(mut self, d: WsDirection) -> Self {
85        self.direction = Some(d);
86        self
87    }
88    pub fn with_control(mut self) -> Self {
89        self.with_control = true;
90        self
91    }
92    fn matches(&self, m: &WsMessage) -> bool {
93        if !(self.with_control || m.opcode == 1 || m.opcode == 2) {
94            return false;
95        }
96        if let Some(s) = &self.url_contains {
97            if !m.url.contains(s) {
98                return false;
99            }
100        }
101        if let Some(d) = self.direction {
102            if m.direction != d {
103                return false;
104            }
105        }
106        true
107    }
108}
109
110const BUFFER_CAP: usize = 500;
111
112/// WebSocket 帧监听句柄(`tab.websocket()` 返回)。
113pub struct ChromiumWsListener {
114    core: Arc<CdpCore>,
115}
116
117impl ChromiumWsListener {
118    pub(crate) fn new(core: Arc<CdpCore>) -> Self {
119        Self { core }
120    }
121
122    /// 开始监听(默认只留 text/binary 数据帧)。
123    pub async fn start(&self) -> Result<()> {
124        self.start_with(WsFilter::default()).await
125    }
126
127    /// 带过滤开始监听。
128    pub async fn start_with(&self, filter: WsFilter) -> Result<()> {
129        self.stop().await?;
130        self.core.send("Network.enable", json!({})).await?;
131        let buf = self.core.ws.lock().await.buf.clone();
132        let task = tokio::spawn(ws_pump(
133            self.core.conn.clone(),
134            self.core.session_id.clone(),
135            filter,
136            buf,
137        ));
138        let mut g = self.core.ws.lock().await;
139        g.running = true;
140        g.abort = Some(task.abort_handle());
141        Ok(())
142    }
143
144    /// 是否正在监听。
145    pub async fn listening(&self) -> bool {
146        self.core.ws.lock().await.running
147    }
148
149    /// 等待一个帧(超时返回 `None`)。
150    pub async fn wait(&self, timeout: Option<Duration>) -> Result<Option<WsMessage>> {
151        let buf = self.core.ws.lock().await.buf.clone();
152        let deadline = Instant::now() + timeout.unwrap_or_else(|| self.core.timeout());
153        loop {
154            if let Some(m) = buf.lock().await.pop_front() {
155                return Ok(Some(m));
156            }
157            if Instant::now() >= deadline {
158                return Ok(None);
159            }
160            sleep(Duration::from_millis(50)).await;
161        }
162    }
163
164    /// 在总超时内尽量收集 `n` 个帧。
165    pub async fn wait_count(&self, n: usize, timeout: Option<Duration>) -> Result<Vec<WsMessage>> {
166        let buf = self.core.ws.lock().await.buf.clone();
167        let deadline = Instant::now() + timeout.unwrap_or_else(|| self.core.timeout());
168        let mut out = Vec::with_capacity(n);
169        while out.len() < n {
170            if let Some(m) = buf.lock().await.pop_front() {
171                out.push(m);
172                continue;
173            }
174            if Instant::now() >= deadline {
175                break;
176            }
177            sleep(Duration::from_millis(50)).await;
178        }
179        Ok(out)
180    }
181
182    /// 取出当前缓冲全部帧并清空。
183    pub async fn messages(&self) -> Vec<WsMessage> {
184        let buf = self.core.ws.lock().await.buf.clone();
185        let mut g = buf.lock().await;
186        g.drain(..).collect()
187    }
188
189    /// 停止监听。
190    pub async fn stop(&self) -> Result<()> {
191        let (abort, buf) = {
192            let mut g = self.core.ws.lock().await;
193            g.running = false;
194            (g.abort.take(), g.buf.clone())
195        };
196        buf.lock().await.clear();
197        if let Some(a) = abort {
198            a.abort();
199        }
200        Ok(())
201    }
202}
203
204async fn ws_pump(
205    conn: Connection,
206    session_id: String,
207    filter: WsFilter,
208    buf: Arc<Mutex<VecDeque<WsMessage>>>,
209) {
210    let mut events = conn.subscribe();
211    let mut urls: HashMap<String, String> = HashMap::new();
212    loop {
213        let ev = match events.recv().await {
214            Ok(ev) => ev,
215            Err(RecvError::Lagged(_)) => continue,
216            Err(RecvError::Closed) => break,
217        };
218        if ev.session_id.as_deref() != Some(session_id.as_str()) {
219            continue;
220        }
221        match ev.method.as_str() {
222            "Network.webSocketCreated" => {
223                if let (Some(id), Some(url)) =
224                    (ev.params["requestId"].as_str(), ev.params["url"].as_str())
225                {
226                    urls.insert(id.to_string(), url.to_string());
227                }
228            }
229            "Network.webSocketFrameSent" | "Network.webSocketFrameReceived" => {
230                let dir = if ev.method.ends_with("Sent") {
231                    WsDirection::Sent
232                } else {
233                    WsDirection::Received
234                };
235                let id = ev.params["requestId"].as_str().unwrap_or_default();
236                let url = urls.get(id).cloned().unwrap_or_default();
237                let resp = &ev.params["response"];
238                let m = WsMessage {
239                    direction: dir,
240                    url,
241                    opcode: resp["opcode"].as_u64().unwrap_or(0) as u8,
242                    data: resp["payloadData"].as_str().unwrap_or_default().to_string(),
243                };
244                if filter.matches(&m) {
245                    let mut g = buf.lock().await;
246                    if g.len() >= BUFFER_CAP {
247                        g.pop_front();
248                    }
249                    g.push_back(m);
250                }
251            }
252            _ => {}
253        }
254    }
255}
256
257/// WS 监听共享状态(放 [`CdpCore`])。
258pub(crate) type WsShared = EventBuf<WsMessage>;
drission/cdp/websocket.rs

drission/cdp/
websocket.rs
