Skip to main content

browser_commander/browser/webdriver/
mod.rs

1//! Managed ChromeDriver/geckodriver through command-stream, with typed Fantoccini
2//! sessions, portable managed downloads and optional native WebDriver BiDi.
3
4mod adapter;
5pub mod bidi;
6mod capabilities;
7mod snapshot;
8mod storage;
9pub use snapshot::{launch_webdriver_snapshot, WebDriverSnapshotResult};
10
11pub use capabilities::build_capabilities;
12pub use fantoccini::{Client as WebDriverClient, Locator};
13
14use crate::{
15    browser::profile_directory::{
16        create_temporary_user_data_dir_with_first_run, remove_user_data_dir,
17    },
18    downloads::{normalize_download_options, DownloadManager, DownloadSetting},
19    utilities::{start_process, ManagedProcess, StartProcessOptions},
20};
21use anyhow::{anyhow, Context, Result};
22use bidi::BidiClient;
23use serde_json::{json, Value};
24use std::{
25    collections::HashMap,
26    path::{Path, PathBuf},
27    sync::Arc,
28    time::Duration,
29};
30use tokio::{
31    net::TcpStream,
32    sync::OnceCell,
33    time::{sleep, timeout},
34};
35
36#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
37pub enum WebDriverBrowser {
38    #[default]
39    Chrome,
40    Firefox,
41    Safari,
42    SafariTechnologyPreview,
43}
44
45impl WebDriverBrowser {
46    pub fn is_safari(self) -> bool {
47        matches!(self, Self::Safari | Self::SafariTechnologyPreview)
48    }
49}
50
51#[derive(Clone, Debug)]
52pub struct WebDriverOptions {
53    pub browser: WebDriverBrowser,
54    /// Matching chromedriver/geckodriver; otherwise resolved through PATH.
55    pub driver_executable: Option<PathBuf>,
56    pub browser_executable: Option<PathBuf>,
57    /// An existing dedicated profile is retained on close. Never use a profile
58    /// that is already open in another browser.
59    pub user_data_dir: Option<PathBuf>,
60    pub headless: bool,
61    pub sandbox: bool,
62    pub args: Vec<String>,
63    pub preferences: Value,
64    pub local_state: Value,
65    pub default_browser_check: Option<bool>,
66    pub first_run: bool,
67    pub automation_parity: bool,
68    pub capabilities: serde_json::Map<String, Value>,
69    /// Added to the driver and inherited by its browser; parent stays unchanged.
70    pub env: Option<HashMap<String, String>>,
71    pub downloads: DownloadSetting,
72    pub launch_timeout: Duration,
73    /// Request webSocketUrl and connect when the driver supplies it.
74    pub bidi: bool,
75}
76
77impl Default for WebDriverOptions {
78    fn default() -> Self {
79        Self {
80            browser: WebDriverBrowser::Chrome,
81            driver_executable: None,
82            browser_executable: None,
83            user_data_dir: None,
84            headless: false,
85            sandbox: true,
86            args: Vec::new(),
87            preferences: json!({}),
88            local_state: json!({}),
89            default_browser_check: None,
90            first_run: false,
91            automation_parity: true,
92            capabilities: Default::default(),
93            env: None,
94            downloads: DownloadSetting::Off,
95            launch_timeout: Duration::from_secs(30),
96            bidi: true,
97        }
98    }
99}
100
101struct OwnedDriver {
102    child: Option<Arc<ManagedProcess>>,
103    profile: PathBuf,
104    temporary: bool,
105}
106
107impl OwnedDriver {
108    async fn stop(&mut self) -> Result<()> {
109        if let Some(child) = self.child.take() {
110            child.kill();
111            child.wait_timeout(Duration::from_secs(6)).await;
112        }
113        if self.temporary {
114            remove_user_data_dir(&self.profile).await?;
115        }
116        Ok(())
117    }
118}
119
120impl Drop for OwnedDriver {
121    fn drop(&mut self) {
122        let profile = self.profile.clone();
123        let temporary = self.temporary;
124        if let Some(child) = self.child.take() {
125            child.kill();
126            if let Ok(runtime) = tokio::runtime::Handle::try_current() {
127                runtime.spawn(async move {
128                    child.wait_timeout(Duration::from_secs(6)).await;
129                    if temporary {
130                        let _ = remove_user_data_dir(&profile).await;
131                    }
132                });
133                return;
134            }
135        }
136        if temporary {
137            let _ = std::fs::remove_dir_all(profile);
138        }
139    }
140}
141
142/// Owns driver, browser session, BiDi and the optional download watcher.
143/// All typed Fantoccini APIs are available through [`Self::client`]. Close
144/// waits for downloads before quitting; Drop stops the owned driver as well.
145pub struct ManagedWebDriver {
146    browser: WebDriverBrowser,
147    client: WebDriverClient,
148    bidi: Option<Arc<BidiClient>>,
149    downloads: Option<Arc<DownloadManager>>,
150    profile: PathBuf,
151    driver: tokio::sync::Mutex<OwnedDriver>,
152    closed: OnceCell<()>,
153}
154
155impl ManagedWebDriver {
156    pub fn client(&self) -> &WebDriverClient {
157        &self.client
158    }
159    pub fn bidi(&self) -> Option<&Arc<BidiClient>> {
160        self.bidi.as_ref()
161    }
162    pub fn downloads(&self) -> Option<&Arc<DownloadManager>> {
163        self.downloads.as_ref()
164    }
165    pub fn user_data_dir(&self) -> &Path {
166        &self.profile
167    }
168    pub fn driver_pid(&self) -> Option<u32> {
169        self.driver
170            .try_lock()
171            .ok()
172            .and_then(|driver| driver.child.as_ref().and_then(|child| child.pid()))
173    }
174    pub(crate) fn process_handle(&self) -> crate::browser::browser_process::BrowserProcess {
175        let child = self
176            .driver
177            .try_lock()
178            .expect("new driver is unlocked")
179            .child
180            .as_ref()
181            .unwrap()
182            .clone();
183        crate::browser::browser_process::BrowserProcess::from_control(child)
184    }
185    /// Reject unsupported Safari features before starting an operation.
186    pub fn require_feature(&self, feature: &str) -> Result<(), crate::core::engine::EngineError> {
187        if self.browser.is_safari() {
188            Err(crate::browser::safari::unsupported(feature))
189        } else {
190            Ok(())
191        }
192    }
193    pub async fn close(&self) -> Result<()> {
194        self.closed
195            .get_or_try_init(|| async {
196                if let Some(manager) = &self.downloads {
197                    manager.dispose().await;
198                }
199                if let Some(bidi) = &self.bidi {
200                    bidi.close().await;
201                }
202                // Even a dead session must release its driver and copied profile.
203                let result = timeout(Duration::from_secs(6), self.client.clone().close()).await;
204                self.driver.lock().await.stop().await?;
205                if let Ok(Err(error)) = result {
206                    tracing::debug!(%error, "WebDriver session was already closed");
207                }
208                Ok::<(), anyhow::Error>(())
209            })
210            .await
211            .map(|_| ())
212    }
213}
214
215impl Drop for ManagedWebDriver {
216    fn drop(&mut self) {
217        if let Ok(runtime) = tokio::runtime::Handle::try_current() {
218            let manager = self.downloads.take();
219            let bidi = self.bidi.take();
220            runtime.spawn(async move {
221                if let Some(manager) = manager {
222                    manager.dispose().await;
223                }
224                if let Some(bidi) = bidi {
225                    bidi.close().await;
226                }
227            });
228        }
229    }
230}
231
232/// Start a locally installed driver. No npm package or Node process is used.
233pub async fn launch_webdriver(options: WebDriverOptions) -> Result<ManagedWebDriver> {
234    launch_owned(options, false).await
235}
236
237async fn launch_owned(
238    mut options: WebDriverOptions,
239    own_profile: bool,
240) -> Result<ManagedWebDriver> {
241    if options.launch_timeout.is_zero() {
242        return Err(anyhow!("launch_timeout must be positive"));
243    }
244    // Validate before creating profiles or starting a process.
245    build_capabilities(&options, Path::new("validation-profile"), None)?;
246    let safari = options.browser.is_safari();
247    if safari && !cfg!(target_os = "macos") && options.driver_executable.is_none() {
248        return Err(crate::browser::safari::unsupported("Safari launch outside macOS").into());
249    }
250    let temporary = !safari && (own_profile || options.user_data_dir.is_none());
251    let profile = if safari {
252        PathBuf::new()
253    } else {
254        match &options.user_data_dir {
255            Some(path) => {
256                std::fs::create_dir_all(path)?;
257                std::fs::canonicalize(path)?
258            }
259            None => create_temporary_user_data_dir_with_first_run(None, options.first_run)?,
260        }
261    };
262    let mut owner = OwnedDriver {
263        child: None,
264        profile: profile.clone(),
265        temporary,
266    };
267    if options.browser == WebDriverBrowser::Chrome {
268        use crate::browser::profile_directory::{
269            configure_user_data_dir_for_profile, prepare_user_data_dir_with_first_run,
270        };
271        prepare_user_data_dir_with_first_run(&profile, options.first_run)?;
272        let selected = options
273            .args
274            .iter()
275            .find_map(|arg| arg.strip_prefix("--profile-directory="))
276            .unwrap_or("Default");
277        configure_user_data_dir_for_profile(
278            &profile,
279            selected,
280            options.default_browser_check,
281            &options.preferences,
282            &options.local_state,
283        )?;
284        if options.automation_parity
285            && !options
286                .args
287                .iter()
288                .any(|arg| arg.starts_with("--remote-debugging-port"))
289        {
290            options.args.push(format!(
291                "--remote-debugging-port={}",
292                crate::browser::debugging_port::reserve_loopback_port()?
293            ));
294        }
295    }
296    let manager = normalize_download_options(options.downloads.clone())
297        .map(DownloadManager::create)
298        .transpose()?;
299    let staging = manager
300        .as_ref()
301        .map(|manager| crate::downloads::sources::prepare_staging_directory(&manager.directory))
302        .transpose()?;
303    let capabilities = build_capabilities(&options, &profile, staging.as_deref())?;
304    let port = crate::browser::debugging_port::reserve_loopback_port()?;
305    let executable = options.driver_executable.clone().unwrap_or_else(|| {
306        PathBuf::from(match options.browser {
307            WebDriverBrowser::Chrome => "chromedriver",
308            WebDriverBrowser::Firefox => "geckodriver",
309            WebDriverBrowser::Safari => "/usr/bin/safaridriver",
310            WebDriverBrowser::SafariTechnologyPreview => {
311                "/Applications/Safari Technology Preview.app/Contents/MacOS/safaridriver"
312            }
313        })
314    });
315    let args = match options.browser {
316        WebDriverBrowser::Chrome => {
317            vec![format!("--port={port}"), "--allowed-ips=127.0.0.1".into()]
318        }
319        WebDriverBrowser::Firefox => vec![
320            "--host".into(),
321            "127.0.0.1".into(),
322            "--port".into(),
323            port.to_string(),
324            // geckodriver otherwise shares port 9222 across sessions, even
325            // when its WebDriver HTTP port is allocated dynamically.
326            "--websocket-port".into(),
327            "0".into(),
328        ],
329        WebDriverBrowser::Safari | WebDriverBrowser::SafariTechnologyPreview => {
330            vec!["--port".into(), port.to_string()]
331        }
332    };
333    let output = Arc::new(std::sync::Mutex::new(Vec::new()));
334    let tail = output.clone();
335    let listener: crate::utilities::OutputListener = Arc::new(move |chunk| {
336        tracing::debug!(output = %String::from_utf8_lossy(chunk), "WebDriver output");
337        let mut tail = tail.lock().unwrap();
338        tail.extend_from_slice(chunk);
339        if tail.len() > 8192 {
340            let excess = tail.len() - 8192;
341            tail.drain(..excess);
342        }
343    });
344    owner.child = Some(Arc::new(
345        start_process(
346            executable.to_string_lossy().as_ref(),
347            &args,
348            StartProcessOptions {
349                env: {
350                    let mut env = options.env.clone().unwrap_or_default();
351                    if options.browser == WebDriverBrowser::Firefox && !options.sandbox {
352                        env.insert("MOZ_DISABLE_CONTENT_SANDBOX".into(), "1".into());
353                    }
354                    (!env.is_empty()).then_some(env)
355                },
356                on_stdout: vec![listener.clone()],
357                on_stderr: vec![listener],
358                ..Default::default()
359            },
360        )
361        .await
362        .with_context(|| format!("start WebDriver {}", executable.display()))?,
363    ));
364    let endpoint = format!("http://127.0.0.1:{port}");
365    let launched = timeout(options.launch_timeout, async {
366        loop {
367            if let Some(code) = owner.child.as_ref().and_then(|child| child.exit_code()) {
368                return Err(anyhow!(
369                    "WebDriver exited with {code}: {}",
370                    String::from_utf8_lossy(&output.lock().unwrap())
371                ));
372            }
373            if TcpStream::connect(("127.0.0.1", port)).await.is_ok() {
374                break;
375            }
376            sleep(Duration::from_millis(25)).await;
377        }
378        let client = fantoccini::ClientBuilder::rustls()?
379            .capabilities(capabilities)
380            .connect(&endpoint)
381            .await?;
382        // Safari reports an empty URL until its first navigation.
383        if safari {
384            client.goto("about:blank").await?;
385        }
386        Ok::<_, anyhow::Error>(client)
387    })
388    .await;
389    let client = match launched {
390        Ok(Ok(client)) => client,
391        result => {
392            owner.stop().await?;
393            return match result {
394                Ok(Err(error)) => Err(if safari {
395                    crate::browser::safari::launch_error(error, options.browser)
396                } else {
397                    error
398                }),
399                _ => Err(anyhow!(
400                    "WebDriver launch timed out: {}",
401                    String::from_utf8_lossy(&output.lock().unwrap())
402                )),
403            };
404        }
405    };
406    let mut browser = ManagedWebDriver {
407        browser: options.browser,
408        client,
409        bidi: None,
410        downloads: manager,
411        profile,
412        driver: tokio::sync::Mutex::new(owner),
413        closed: OnceCell::new(),
414    };
415    let setup = async {
416        if options.bidi {
417            if let Some(url) = browser
418                .client
419                .capabilities()
420                .and_then(|caps| caps.get("webSocketUrl"))
421                .and_then(Value::as_str)
422            {
423                browser.bidi = Some(Arc::new(BidiClient::connect(url).await?));
424            }
425        }
426        if let Some(manager) = &browser.downloads {
427            manager.attach_filesystem().await?;
428        }
429        Ok::<(), anyhow::Error>(())
430    }
431    .await;
432    if let Err(error) = setup {
433        browser.close().await?;
434        return Err(error);
435    }
436    Ok(browser)
437}