use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{mpsc, Mutex};
use crate::cdp::transport::CdpMessage;
use crate::cdp::types::{
NetworkLoadingFailedEvent, NetworkLoadingFinishedEvent, NetworkRequestWillBeSentEvent,
NetworkResponseReceivedEvent,
};
use crate::page::CapturedRequest;
#[derive(Debug, Clone)]
pub enum NetworkEvent {
RequestStarted(CapturedRequest),
ResponseReceived {
request_id: String,
status: i32,
status_text: String,
headers: HashMap<String, String>,
mime_type: Option<String>,
},
RequestCompleted {
request_id: String,
encoded_data_length: i64,
},
RequestFailed {
request_id: String,
error_text: String,
canceled: bool,
},
}
const MAX_INFLIGHT_REQUESTS: usize = 10_000;
pub struct NetworkWatcher {
requests: Arc<Mutex<HashMap<String, CapturedRequest>>>,
event_tx: mpsc::Sender<NetworkEvent>,
event_rx: Mutex<mpsc::Receiver<NetworkEvent>>,
}
impl NetworkWatcher {
pub fn new() -> Self {
let (event_tx, event_rx) = mpsc::channel(256);
Self {
requests: Arc::new(Mutex::new(HashMap::new())),
event_tx,
event_rx: Mutex::new(event_rx),
}
}
pub async fn process_event(&self, event: &CdpMessage) -> bool {
if let CdpMessage::Event { method, params, .. } = event {
match method.as_str() {
"Network.requestWillBeSent" => {
if let Ok(e) =
serde_json::from_value::<NetworkRequestWillBeSentEvent>(params.clone())
{
self.on_request_will_be_sent(e).await;
return true;
}
}
"Network.responseReceived" => {
if let Ok(e) =
serde_json::from_value::<NetworkResponseReceivedEvent>(params.clone())
{
self.on_response_received(e).await;
return true;
}
}
"Network.loadingFinished" => {
if let Ok(e) =
serde_json::from_value::<NetworkLoadingFinishedEvent>(params.clone())
{
self.on_loading_finished(e).await;
return true;
}
}
"Network.loadingFailed" => {
if let Ok(e) =
serde_json::from_value::<NetworkLoadingFailedEvent>(params.clone())
{
self.on_loading_failed(e).await;
return true;
}
}
_ => {}
}
}
false
}
async fn on_request_will_be_sent(&self, event: NetworkRequestWillBeSentEvent) {
let request = {
let mut requests = self.requests.lock().await;
if requests.contains_key(&event.request_id) {
tracing::trace!(
"Ignoring redirect requestWillBeSent for {} (keeping original)",
event.request_id
);
return;
}
if requests.len() >= MAX_INFLIGHT_REQUESTS {
if let Some(id) = requests.keys().next().cloned() {
requests.remove(&id);
tracing::debug!("Evicted a tracked request (cap={})", MAX_INFLIGHT_REQUESTS);
}
}
let request = CapturedRequest {
request_id: event.request_id.clone(),
url: event.request.url.clone(),
method: event.request.method.clone(),
headers: event.request.headers.clone(),
post_data: event.request.post_data.clone(),
resource_type: event.r#type.clone(),
status: None,
status_text: None,
response_headers: None,
mime_type: None,
timestamp: event.timestamp,
complete: false,
};
requests.insert(event.request_id.clone(), request.clone());
request
};
self.emit(NetworkEvent::RequestStarted(request));
}
async fn on_response_received(&self, event: NetworkResponseReceivedEvent) {
{
let mut requests = self.requests.lock().await;
if let Some(req) = requests.get_mut(&event.request_id) {
req.status = Some(event.response.status);
req.status_text = Some(event.response.status_text.clone());
req.response_headers = Some(event.response.headers.clone());
req.mime_type = event.response.mime_type.clone();
}
}
self.emit(NetworkEvent::ResponseReceived {
request_id: event.request_id,
status: event.response.status,
status_text: event.response.status_text,
headers: event.response.headers,
mime_type: event.response.mime_type,
});
}
async fn on_loading_finished(&self, event: NetworkLoadingFinishedEvent) {
{
let mut requests = self.requests.lock().await;
if let Some(req) = requests.get_mut(&event.request_id) {
req.complete = true;
}
}
self.emit(NetworkEvent::RequestCompleted {
request_id: event.request_id,
encoded_data_length: event.encoded_data_length,
});
}
async fn on_loading_failed(&self, event: NetworkLoadingFailedEvent) {
{
let mut requests = self.requests.lock().await;
requests.remove(&event.request_id);
}
self.emit(NetworkEvent::RequestFailed {
request_id: event.request_id,
error_text: event.error_text,
canceled: event.canceled.unwrap_or(false),
});
}
fn emit(&self, event: NetworkEvent) {
if let Err(e) = self.event_tx.try_send(event) {
tracing::warn!("Dropping network event (consumer not draining): {}", e);
}
}
pub async fn recv(&self) -> Option<NetworkEvent> {
let mut rx = self.event_rx.lock().await;
rx.recv().await
}
pub async fn try_recv(&self) -> Option<NetworkEvent> {
let mut rx = match self.event_rx.try_lock() {
Ok(rx) => rx,
Err(_) => return None,
};
rx.try_recv().ok()
}
pub async fn get_request(&self, request_id: &str) -> Option<CapturedRequest> {
let requests = self.requests.lock().await;
requests.get(request_id).cloned()
}
pub async fn get_all_requests(&self) -> Vec<CapturedRequest> {
let requests = self.requests.lock().await;
requests.values().cloned().collect()
}
pub async fn clear(&self) {
let mut requests = self.requests.lock().await;
requests.clear();
}
}
impl Default for NetworkWatcher {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_network_watcher_creation() {
let watcher = NetworkWatcher::new();
assert!(watcher.get_all_requests().await.is_empty());
}
fn event(method: &str, params: serde_json::Value) -> CdpMessage {
CdpMessage::Event {
method: method.to_string(),
params,
session_id: None,
}
}
fn request_will_be_sent(id: &str, url: &str) -> CdpMessage {
event(
"Network.requestWillBeSent",
serde_json::json!({
"requestId": id,
"request": {
"url": url,
"method": "GET",
"headers": {},
},
"timestamp": 1.0,
"type": "Document",
}),
)
}
#[tokio::test]
async fn test_completed_request_is_retained() {
let watcher = NetworkWatcher::new();
let id = "req-1";
assert!(
watcher
.process_event(&request_will_be_sent(id, "https://example.com/"))
.await
);
assert!(
watcher
.process_event(&event(
"Network.responseReceived",
serde_json::json!({
"requestId": id,
"response": {
"url": "https://example.com/",
"status": 200,
"statusText": "OK",
"headers": {},
"mimeType": "text/html",
},
}),
))
.await
);
assert!(
watcher
.process_event(&event(
"Network.loadingFinished",
serde_json::json!({
"requestId": id,
"timestamp": 2.0,
"encodedDataLength": 1234,
}),
))
.await
);
let req = watcher.get_request(id).await.expect("request retained");
assert!(req.complete);
assert_eq!(req.status, Some(200));
assert_eq!(req.status_text.as_deref(), Some("OK"));
}
#[tokio::test]
async fn test_redirect_keeps_original_entry() {
let watcher = NetworkWatcher::new();
let id = "req-redirect";
assert!(
watcher
.process_event(&request_will_be_sent(id, "https://first.example/"))
.await
);
assert!(
watcher
.process_event(&request_will_be_sent(id, "https://second.example/"))
.await
);
let req = watcher.get_request(id).await.expect("request present");
assert_eq!(req.url, "https://first.example/");
assert_eq!(watcher.get_all_requests().await.len(), 1);
}
#[tokio::test]
async fn test_loading_failed_removes_entry() {
let watcher = NetworkWatcher::new();
let id = "req-fail";
assert!(
watcher
.process_event(&request_will_be_sent(id, "https://example.com/"))
.await
);
assert!(watcher.get_request(id).await.is_some());
assert!(
watcher
.process_event(&event(
"Network.loadingFailed",
serde_json::json!({
"requestId": id,
"errorText": "net::ERR_FAILED",
"canceled": false,
}),
))
.await
);
assert!(watcher.get_request(id).await.is_none());
}
#[tokio::test]
async fn test_pump_does_not_block_when_events_not_drained() {
let watcher = NetworkWatcher::new();
const N: usize = 300;
for i in 0..N {
let id = format!("req-{i}");
let url = format!("https://example.com/{i}");
assert!(
watcher
.process_event(&request_will_be_sent(&id, &url))
.await
);
}
assert_eq!(watcher.get_all_requests().await.len(), N);
assert!(watcher.get_request("req-0").await.is_some());
assert!(watcher
.get_request(&format!("req-{}", N - 1))
.await
.is_some());
}
}