use std::fmt;
use std::path::PathBuf;
use std::sync::Arc;
use parking_lot::Mutex;
use rustc_hash::FxHashMap;
use serde_json::Value;
use tokio::process::Child;
use tracing::{debug, info};
use uuid::Uuid;
use crate::driver::{Driver, FirefoxOptions, Profile};
use crate::error::{Error, Result};
use crate::identifiers::{FrameId, InterceptId, SessionId, TabId};
use crate::protocol::{
BrowsingContextCommand, Command, ProxyCommand, Request, Response, SessionCommand,
};
use crate::transport::ConnectionPool;
use super::Tab;
use super::proxy::ProxyConfig;
struct ProcessGuard {
child: Option<Child>,
pid: u32,
}
impl ProcessGuard {
fn new(child: Child) -> Self {
let pid = child.id().unwrap_or(0);
debug!(pid, "Process guard created");
Self {
child: Some(child),
pid,
}
}
#[inline]
fn pid(&self) -> u32 {
self.pid
}
}
impl Drop for ProcessGuard {
fn drop(&mut self) {
if let Some(mut child) = self.child.take()
&& let Err(e) = child.start_kill()
{
debug!(pid = self.pid, error = %e, "Failed to send kill signal in Drop");
}
}
}
pub(crate) struct WindowInner {
pub uuid: Uuid,
pub session_id: SessionId,
process: Mutex<ProcessGuard>,
pub pool: Arc<ConnectionPool>,
#[allow(dead_code)]
profile: Profile,
tabs: Mutex<FxHashMap<TabId, Tab>>,
pub initial_tab_id: TabId,
pub intercept_handlers: Mutex<FxHashMap<InterceptId, String>>,
}
#[derive(Clone)]
pub struct Window {
pub(crate) inner: Arc<WindowInner>,
}
impl fmt::Debug for Window {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Window")
.field("uuid", &self.inner.uuid)
.field("session_id", &self.inner.session_id)
.field("port", &self.inner.pool.port())
.finish_non_exhaustive()
}
}
impl Window {
pub(crate) fn new(
pool: Arc<ConnectionPool>,
process: Child,
profile: Profile,
session_id: SessionId,
initial_tab_id: TabId,
) -> Self {
let uuid = Uuid::new_v4();
let initial_tab = Tab::new(initial_tab_id, FrameId::main(), session_id, None);
let mut tabs = FxHashMap::default();
tabs.insert(initial_tab_id, initial_tab);
debug!(
uuid = %uuid,
session_id = %session_id,
tab_id = %initial_tab_id,
port = pool.port(),
"Window created"
);
Self {
inner: Arc::new(WindowInner {
uuid,
session_id,
process: Mutex::new(ProcessGuard::new(process)),
pool,
profile,
tabs: Mutex::new(tabs),
initial_tab_id,
intercept_handlers: Mutex::new(FxHashMap::default()),
}),
}
}
}
impl Window {
#[inline]
#[must_use]
pub fn session_id(&self) -> SessionId {
self.inner.session_id
}
#[inline]
#[must_use]
pub fn uuid(&self) -> &Uuid {
&self.inner.uuid
}
#[inline]
#[must_use]
pub fn port(&self) -> u16 {
self.inner.pool.port()
}
#[inline]
#[must_use]
pub fn pid(&self) -> u32 {
self.inner.process.lock().pid()
}
}
impl Window {
pub async fn close(&self) -> Result<()> {
debug!(uuid = %self.inner.uuid, "Closing window");
self.inner.pool.remove(self.inner.session_id);
let child = {
let mut guard = self.inner.process.lock();
guard.child.take()
};
if let Some(mut child) = child {
let pid = child.id().unwrap_or(0);
debug!(pid, "Killing Firefox process");
if let Err(e) = child.kill().await {
debug!(pid, error = %e, "Failed to kill process");
}
if let Err(e) = child.wait().await {
debug!(pid, error = %e, "Failed to wait for process");
}
info!(pid, "Process terminated");
}
info!(uuid = %self.inner.uuid, "Window closed");
Ok(())
}
}
impl Window {
#[must_use]
pub fn tab(&self) -> Tab {
Tab::new(
self.inner.initial_tab_id,
FrameId::main(),
self.inner.session_id,
Some(self.clone()),
)
}
pub async fn new_tab(&self) -> Result<Tab> {
let command = Command::BrowsingContext(BrowsingContextCommand::NewTab);
let response = self.send_command(command).await?;
let tab_id_u32 = response
.result
.as_ref()
.and_then(|v| v.get("tabId"))
.and_then(|v| v.as_u64())
.ok_or_else(|| Error::protocol("Expected tabId in NewTab response"))?;
let new_tab_id = TabId::new(tab_id_u32 as u32)
.ok_or_else(|| Error::protocol("Invalid tabId in NewTab response"))?;
let tab = Tab::new(
new_tab_id,
FrameId::main(),
self.inner.session_id,
Some(self.clone()),
);
self.inner.tabs.lock().insert(new_tab_id, tab.clone());
debug!(session_id = %self.inner.session_id, tab_id = %new_tab_id, "New tab created");
Ok(tab)
}
#[inline]
#[must_use]
pub fn tab_count(&self) -> usize {
self.inner.tabs.lock().len()
}
pub async fn steal_logs(&self) -> Result<Vec<Value>> {
let command = Command::Session(SessionCommand::StealLogs);
let response = self.send_command(command).await?;
let logs = response
.result
.as_ref()
.and_then(|v| v.get("logs"))
.and_then(|v| v.as_array())
.cloned()
.unwrap_or_default();
Ok(logs)
}
}
impl Window {
pub async fn set_proxy(&self, config: ProxyConfig) -> Result<()> {
debug!(
session_id = %self.inner.session_id,
proxy_type = %config.proxy_type.as_str(),
host = %config.host,
port = config.port,
"Setting window proxy"
);
let command = Command::Proxy(ProxyCommand::SetWindowProxy {
proxy_type: config.proxy_type.as_str().to_string(),
host: config.host,
port: config.port,
username: config.username,
password: config.password,
proxy_dns: config.proxy_dns,
});
self.send_command(command).await?;
Ok(())
}
pub async fn clear_proxy(&self) -> Result<()> {
debug!(session_id = %self.inner.session_id, "Clearing window proxy");
let command = Command::Proxy(ProxyCommand::ClearWindowProxy);
self.send_command(command).await?;
Ok(())
}
}
impl Window {
pub(crate) async fn send_command(&self, command: Command) -> Result<Response> {
let request = Request::new(self.inner.initial_tab_id, FrameId::main(), command);
self.inner.pool.send(self.inner.session_id, request).await
}
}
pub struct WindowBuilder<'a> {
driver: &'a Driver,
options: FirefoxOptions,
profile: Option<PathBuf>,
}
impl<'a> WindowBuilder<'a> {
pub(crate) fn new(driver: &'a Driver) -> Self {
Self {
driver,
options: FirefoxOptions::new(),
profile: None,
}
}
#[must_use]
pub fn headless(mut self) -> Self {
self.options = self.options.with_headless();
self
}
#[must_use]
pub fn window_size(mut self, width: u32, height: u32) -> Self {
self.options = self.options.with_window_size(width, height);
self
}
#[must_use]
pub fn profile(mut self, path: impl Into<PathBuf>) -> Self {
self.profile = Some(path.into());
self
}
pub async fn spawn(self) -> Result<Window> {
self.driver.spawn_window(self.options, self.profile).await
}
}
#[cfg(test)]
mod tests {
use super::Window;
#[test]
fn test_window_is_clone() {
fn assert_clone<T: Clone>() {}
assert_clone::<Window>();
}
#[test]
fn test_window_is_debug() {
fn assert_debug<T: std::fmt::Debug>() {}
assert_debug::<Window>();
}
}