use crate::error::{CdpError, CdpResult};
use crate::event_filter::EventFilter;
use crate::protocol::{WsCommand, WsResponse};
use crate::rest_client::get_websocket_url;
use futures_util::{SinkExt, StreamExt};
use serde::Serialize;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::broadcast;
use tokio::sync::mpsc;
use tokio::time::timeout;
use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite::Message;
use tracing::{debug, error, trace};
#[derive(Clone)]
pub struct CdpClient {
command_tx: mpsc::UnboundedSender<String>,
pending_requests: Arc<Mutex<HashMap<u64, oneshot::Sender<WsResponse>>>>,
event_tx: broadcast::Sender<WsResponse>,
next_id: Arc<AtomicU64>,
default_timeout: Duration,
is_alive: Arc<AtomicBool>,
}
impl CdpClient {
pub async fn new(host: &str, default_timeout: Duration) -> CdpResult<Self> {
let ws_url = get_websocket_url(host).await?;
let (ws_stream, _) = connect_async(ws_url).await?;
let (mut ws_sink, mut ws_stream) = ws_stream.split();
let pending_requests = Arc::new(Mutex::new(
HashMap::<u64, oneshot::Sender<WsResponse>>::new(),
));
let pending_requests_clone = Arc::clone(&pending_requests);
let (command_tx, mut command_rx) = mpsc::unbounded_channel::<String>();
let (event_tx, _) = broadcast::channel(128);
let event_tx_clone = event_tx.clone();
let is_alive = Arc::new(AtomicBool::new(true));
let is_alive_reader = is_alive.clone();
tokio::spawn(async move {
let result: CdpResult<()> = async {
while let Some(json_str) = command_rx.recv().await {
debug!("About to send command {}", json_str);
ws_sink.send(Message::Text(json_str.into())).await?;
}
Ok(())
}
.await;
if let Err(e) = result {
error!("Fatal error in CDP Writer task: {}", e);
}
});
tokio::spawn(async move {
let reader_result: CdpResult<()> = async {
while let Some(result_stream) = ws_stream.next().await {
let msg = result_stream.map_err(CdpError::from)?;
if let Message::Text(received_text) = msg {
let response: WsResponse = serde_json::from_str(&received_text)?;
if let Some(id) = response.id {
let mut map = pending_requests_clone.lock().map_err(|_| {
CdpError::InternalError(
"Could not acquire mutex lock on pending requests cloned map"
.to_string(),
)
})?;
if let Some(responder) = map.remove(&id) {
trace!("Received expected id {} from CDP", id);
let _ = responder.send(response);
} else {
trace!("Discarded message with id {} (no listener found)", id);
}
} else if let Some(ref method) = response.method {
trace!("CDP Event: {}", method);
let _ = event_tx_clone.send(response.clone());
} else {
debug!(
"Received unexpected message from CDP (no id, no method): {}",
received_text
);
}
}
}
Ok(())
}
.await;
is_alive_reader.store(false, Ordering::SeqCst);
match reader_result {
Err(e) => debug!("CDP Connection lost due to error: {}", e),
Ok(_) => debug!("CDP Connection closed by server/gracefully"),
}
if let Ok(mut map) = pending_requests_clone.lock() {
debug!(
"Closing {} pending requests due to disconnection",
map.len()
);
map.clear(); }
});
let next_id = Arc::new(AtomicU64::new(1));
Ok(Self {
command_tx,
pending_requests,
event_tx,
next_id,
default_timeout,
is_alive,
})
}
pub fn subscribe(&self) -> broadcast::Receiver<WsResponse> {
self.event_tx.subscribe()
}
pub fn on_domain(&self, domain: &'static str) -> EventFilter {
EventFilter::new(self.subscribe(), domain)
}
pub async fn send_raw_command<P: Serialize>(
&self,
method: &str,
params: P,
) -> CdpResult<WsResponse> {
if !self.is_alive.load(Ordering::SeqCst) {
return Err(CdpError::Disconnected);
}
let id = self.next_id.fetch_add(1, Ordering::SeqCst);
let (tx, rx) = oneshot::channel();
{
let mut map = self.pending_requests.lock().map_err(|_| {
CdpError::InternalError(
"Mutex poisoned: another thread panicked while holding the lock".into(),
)
})?;
map.insert(id, tx);
}
let cmd = WsCommand {
id,
method: method.to_string(),
params: Some(params),
};
let json_payload = serde_json::to_string(&cmd)?;
self.command_tx.send(json_payload).map_err(|e| {
self.is_alive.store(false, Ordering::SeqCst);
CdpError::InternalError(format!("Failed to send command to writer task: {}", e))
})?;
match timeout(self.default_timeout, rx).await {
Ok(Ok(response)) => {
if let Some(error_obj) = response.error {
return Err(CdpError::ProtocolError {
code: error_obj["code"].as_i64().unwrap_or(-1),
message: error_obj["message"]
.as_str()
.unwrap_or("Unknown CDP error")
.to_string(),
});
}
Ok(response)
}
Ok(Err(_)) => Err(CdpError::Disconnected),
Err(_) => {
if let Ok(mut map) = self.pending_requests.lock() {
map.remove(&id);
}
Err(CdpError::Timeout {
method: method.to_string(),
timeout: self.default_timeout,
})
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
#[tokio::test]
async fn test_command_timeout_and_cleanup() {
let (command_tx, _command_rx) = mpsc::unbounded_channel();
let client = CdpClient {
command_tx,
pending_requests: Arc::new(Mutex::new(HashMap::new())),
event_tx: broadcast::channel(1).0,
next_id: Arc::new(AtomicU64::new(1)),
default_timeout: Duration::from_millis(100),
is_alive: Arc::new(AtomicBool::new(true)),
};
let result = client
.send_raw_command("Page.navigate", serde_json::json!({"url": "about:blank"}))
.await;
match result {
Err(CdpError::Timeout { method, .. }) => {
assert_eq!(method, "Page.navigate");
println!("✅ Timeout detected successfully");
}
other => panic!("❌ Expected Timeout, but got: {:?}", other),
}
let pending = client.pending_requests.lock().unwrap();
assert_eq!(pending.len(), 0, "❌ ID was not removed from HashMap");
println!("✅ HashMap is clean, no memory leaks");
}
#[tokio::test]
async fn test_protocol_error_mapping() {
let (command_tx, _command_rx) = mpsc::unbounded_channel();
let client = CdpClient {
command_tx,
pending_requests: Arc::new(Mutex::new(HashMap::new())),
event_tx: broadcast::channel(1).0,
next_id: Arc::new(AtomicU64::new(1)),
default_timeout: Duration::from_secs(1),
is_alive: Arc::new(AtomicBool::new(true)),
};
let client_clone = client.clone();
let handle = tokio::spawn(async move {
client_clone
.send_raw_command("Page.navigate", serde_json::json!({"url": "invalid-url"}))
.await
});
tokio::time::sleep(Duration::from_millis(10)).await;
let responder = {
let mut map = client.pending_requests.lock().unwrap();
map.remove(&1)
.expect("The client should have registered request with ID 1")
};
let error_response = WsResponse {
id: Some(1),
result: None,
method: None,
params: None,
error: Some(serde_json::json!({
"code": -32000,
"message": "Cannot navigate to invalid URL"
})),
};
responder.send(error_response).unwrap();
let result = handle.await.unwrap();
match result {
Err(CdpError::ProtocolError { code, message }) => {
assert_eq!(code, -32000);
assert_eq!(message, "Cannot navigate to invalid URL");
println!("✅ Protocol Error mapped correctly");
}
other => panic!("❌ Expected ProtocolError, but got: {:?}", other),
}
}
#[tokio::test]
async fn test_event_broadcasting_and_filtering() {
let (command_tx, _command_rx) = mpsc::unbounded_channel();
let (event_tx, _) = broadcast::channel(32);
let client = CdpClient {
command_tx,
pending_requests: Arc::new(Mutex::new(HashMap::new())),
event_tx: event_tx.clone(),
next_id: Arc::new(AtomicU64::new(1)),
default_timeout: Duration::from_secs(1),
is_alive: Arc::new(AtomicBool::new(true)),
};
let mut page_events = client.on_domain("Page");
let mock_event = WsResponse {
id: None, result: None,
method: Some("Page.loadEventFired".to_string()),
params: Some(serde_json::json!({
"timestamp": 12345.678
})),
error: None,
};
let event_to_send = mock_event.clone();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(10)).await;
let _ = event_tx.send(event_to_send);
});
let received_event = timeout(Duration::from_secs(1), page_events.next())
.await
.expect("Timeout waiting for Page event")
.expect("Stream closed unexpectedly")
.expect("Failed to receive event from broadcast channel");
assert_eq!(
received_event.method.as_deref(),
Some("Page.loadEventFired")
);
let timestamp = received_event
.params
.as_ref()
.and_then(|p| p.get("timestamp"))
.and_then(|t| t.as_f64());
assert_eq!(timestamp, Some(12345.678));
println!("✅ Event correctly filtered and received by domain subscriber");
}
#[tokio::test]
async fn test_broadcast_lagged_error() {
let (command_tx, _command_rx) = mpsc::unbounded_channel();
let (event_tx, _) = broadcast::channel(4);
let client = CdpClient {
command_tx,
pending_requests: Arc::new(Mutex::new(HashMap::new())),
event_tx: event_tx.clone(),
next_id: Arc::new(AtomicU64::new(1)),
default_timeout: Duration::from_secs(1),
is_alive: Arc::new(AtomicBool::new(true)),
};
let mut page_events = client.on_domain("Page");
for i in 0..10 {
let mock_event = WsResponse {
id: None,
result: None,
method: Some("Page.event".to_string()),
params: Some(serde_json::json!({ "index": i })),
error: None,
};
let _ = event_tx.send(mock_event);
}
let result = page_events.next().await;
match result {
Some(Err(CdpError::InternalError(msg))) => {
assert!(msg.contains("channel overflow") || msg.contains("lagged"));
println!("✅ Correctly detected lagged subscriber (buffer overflow)");
}
Some(Ok(event)) => {
panic!(
"❌ Expected a Lagged error, but received a successful event: {:?}",
event
);
}
None => panic!("❌ Stream closed unexpectedly"),
_ => {}
}
println!("✅ Buffer overflow test passed: Oldest messages were dropped as expected");
}
#[tokio::test]
async fn test_connection_drop_during_request() {
let (command_tx, _command_rx) = mpsc::unbounded_channel();
let pending_requests = Arc::new(Mutex::new(HashMap::new()));
let client = CdpClient {
command_tx,
pending_requests: pending_requests.clone(),
event_tx: broadcast::channel(1).0,
next_id: Arc::new(AtomicU64::new(1)),
default_timeout: Duration::from_secs(5),
is_alive: Arc::new(AtomicBool::new(true)),
};
let client_clone = client.clone();
let request_handle = tokio::spawn(async move {
client_clone
.send_raw_command("Debugger.enable", serde_json::json!({}))
.await
});
tokio::time::sleep(Duration::from_millis(10)).await;
{
let mut map = pending_requests.lock().unwrap();
map.clear();
}
client.is_alive.store(false, Ordering::SeqCst);
let result = request_handle.await.unwrap();
match result {
Err(CdpError::Disconnected) => {
println!("✅ Correctly detected disconnection during pending request");
}
other => panic!("❌ Expected CdpError::Disconnected, but got: {:?}", other),
}
}
}