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

browser_commander/browser/webdriver/
mod.rs
