use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::Duration;
use anyhow::{anyhow, Context, Result};
use tokio::time::sleep;
use webdriver_downloader::prelude::*;
#[cfg(all(feature = "geckodriver", feature = "chromedriver"))]
compile_error!("Only one of 'geckodriver' or 'chromedriver' features can be enabled at a time.");
#[cfg(not(any(feature = "geckodriver", feature = "chromedriver")))]
compile_error!("At least one of 'geckodriver' or 'chromedriver' features must be enabled.");
#[cfg(target_os = "windows")]
const DRIVER_EXT: &str = ".exe";
#[cfg(not(target_os = "windows"))]
const DRIVER_EXT: &str = "";
const BROWSER_BIN_PATH_ENV: &str = "BROWSER_PATH";
const WEBDRIVER_BIN_PATH_ENV: &str = "WEBDRIVER_PATH";
const INSTALL_PATH_ENV: &str = "WEBDRIVER_INSTALL_PATH";
const GECKODRIVER_NAME: &str = "geckodriver";
const CHROMEDRIVER_NAME: &str = "chromedriver";
const MAX_DOWNLOAD_RETRIES: u32 = 3;
const INITIAL_RETRY_DELAY: u64 = 2;
struct WebdriverDownloadConfig {
driver_name: &'static str,
get_browser_path: fn() -> Result<PathBuf>,
}
fn user_bin_dir() -> PathBuf {
if let Ok(bin) = env::var(INSTALL_PATH_ENV) {
return PathBuf::from(bin);
}
if let Some(home) = dirs::home_dir() {
#[cfg(target_os = "windows")]
{
return home.join(".local").join("bin");
}
#[cfg(not(target_os = "windows"))]
{
return home.join(".local").join("bin");
}
}
PathBuf::from(".")
}
fn is_webdriver_available(bin_name: &str) -> bool {
if let Ok(path) = env::var(WEBDRIVER_BIN_PATH_ENV) {
let bin_path = if cfg!(target_os = "windows") && !path.to_lowercase().ends_with(".exe") {
format!("{path}{DRIVER_EXT}")
} else {
path
};
let exe = Path::new(&bin_path);
if exe.exists() && exe.is_file() {
println!("{bin_name} found at path specified in {WEBDRIVER_BIN_PATH_ENV}: {bin_path}");
return true;
}
}
let bin_dir = user_bin_dir();
let bin_path = bin_dir.join(format!("{bin_name}{DRIVER_EXT}"));
if bin_path.exists() && bin_path.is_file() {
println!(
"{} found in user's bin directory: {}",
bin_name,
bin_path.display()
);
println!(
"cargo:rustc-env=WEBDRIVER_DOWNLOAD_PATH={}",
bin_path.to_string_lossy()
);
return true;
}
false
}
async fn download_with_retry(
driver_info: &impl WebdriverDownloadInfo,
reinstall: bool,
skip_verification: bool,
num_tries: usize,
) -> Result<()> {
let mut attempts = 0;
let mut last_error = None;
while attempts < MAX_DOWNLOAD_RETRIES {
match download(driver_info, reinstall, skip_verification, num_tries).await {
Ok(_) => {
return Ok(());
}
Err(e) => {
last_error = Some(e);
attempts += 1;
if attempts < MAX_DOWNLOAD_RETRIES {
let delay = Duration::from_secs(INITIAL_RETRY_DELAY * 2u64.pow(attempts - 1));
println!(
"cargo:warning=Download attempt {attempts} failed, retrying in {delay:?}..."
);
sleep(delay).await;
}
}
}
}
Err(anyhow!(
"Failed to download driver after {MAX_DOWNLOAD_RETRIES} attempts: {last_error:?}",
))
}
fn setup_driver(config: &WebdriverDownloadConfig) -> Result<()> {
if is_webdriver_available(config.driver_name) {
return Ok(());
}
println!(
"cargo::warning=You can specify {WEBDRIVER_BIN_PATH_ENV} to an existing {CHROMEDRIVER_NAME}/{GECKODRIVER_NAME} installation to avoid downloads."
);
println!(
"cargo::warning=You can override browser detection using {BROWSER_BIN_PATH_ENV} environment variable."
);
println!(
"cargo::warning={} selected but not installed, will be downloaded ... ",
config.driver_name
);
let browser_path = (config.get_browser_path)()
.with_context(|| format!("Failed to detect browser path for {}", config.driver_name))?;
let browser_version = get_browser_version(&browser_path).with_context(|| {
format!(
"Failed to get version for browser at {}",
browser_path.display()
)
})?;
println!("cargo::warning=Browser version detected: {browser_version}");
let webdriver_bin_dir = user_bin_dir();
println!("cargo::warning=Driver will be installed in: {webdriver_bin_dir:?}");
fs::create_dir_all(&webdriver_bin_dir).with_context(|| {
format!(
"Failed to create directory: {}",
webdriver_bin_dir.display()
)
})?;
let webdriver_bin = webdriver_bin_dir.join(config.driver_name);
println!(
"cargo::rerun-if-changed={}",
webdriver_bin.to_string_lossy()
);
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.context("Failed to create Tokio runtime")?;
match config.driver_name {
CHROMEDRIVER_NAME => {
let driver_info = ChromedriverInfo::new(webdriver_bin.clone(), browser_path);
runtime
.block_on(async { download_with_retry(&driver_info, false, true, 1).await })
.with_context(|| {
format!("Failed to download and install {}", config.driver_name)
})?;
}
GECKODRIVER_NAME => {
let driver_info = GeckodriverInfo::new(webdriver_bin.clone(), browser_path);
runtime
.block_on(async { download_with_retry(&driver_info, false, true, 1).await })
.with_context(|| {
format!("Failed to download and install {}", config.driver_name)
})?;
}
_ => return Err(anyhow!("Unsupported driver type: {}", config.driver_name)),
}
println!(
"cargo:rustc-env=WEBDRIVER_DOWNLOAD_PATH={}",
webdriver_bin.to_string_lossy()
);
Ok(())
}
#[cfg(feature = "chromedriver")]
fn get_chrome_path() -> Result<PathBuf> {
if let Ok(chrome_path) = env::var(BROWSER_BIN_PATH_ENV) {
let path = PathBuf::from(&chrome_path);
if path.exists() {
Ok(path)
} else {
Err(anyhow!("Chrome not found on path: {chrome_path}")).with_context(|| {
format!("Please set {BROWSER_BIN_PATH_ENV} to a valid Chrome installation")
})
}
} else {
let new_browser_path = os_specific::chromedriver_for_testing::default_browser_path()?;
let old_browser_path = os_specific::chromedriver_old::default_browser_path()?;
if new_browser_path.exists() {
Ok(new_browser_path)
} else if old_browser_path.exists() {
Ok(old_browser_path)
} else {
Err(anyhow!("Chrome browser not detected")).with_context(|| {
format!("Use {BROWSER_BIN_PATH_ENV} to point to a valid Chrome installation")
})
}
}
}
#[cfg(feature = "geckodriver")]
fn get_firefox_path() -> Result<PathBuf> {
if let Ok(firefox_path) = env::var(BROWSER_BIN_PATH_ENV) {
let path = PathBuf::from(firefox_path.clone());
if path.exists() {
Ok(path)
} else {
Err(anyhow!("Firefox not found on path: {firefox_path}")).with_context(|| {
format!("Please set {BROWSER_BIN_PATH_ENV} to a valid Firefox installation",)
})
}
} else {
let browser_path = os_specific::geckodriver::default_browser_path()?;
if browser_path.exists() {
Ok(browser_path)
} else {
Err(anyhow!("Firefox browser not detected")).with_context(|| {
format!("Use {BROWSER_BIN_PATH_ENV} to point to a valid Firefox installation",)
})
}
}
}
fn get_browser_version(path: &PathBuf) -> Result<String> {
let output = Command::new(path)
.arg("--version")
.output()
.with_context(|| format!("Failed to execute browser at {}", path.display()))?;
let out_str = String::from_utf8_lossy(&output.stdout);
out_str
.split_whitespace()
.find(|s| s.chars().next().unwrap_or(' ').is_ascii_digit())
.map(|v| v.to_string())
.ok_or(anyhow!(
"Failed to get browser version for browser: {}",
path.display()
))
.with_context(|| {
format!(
"Browser at {} did not return a valid version string",
path.display()
)
})
}
async fn download(
driver_info: &impl WebdriverDownloadInfo,
reinstall: bool,
skip_verification: bool,
num_tries: usize,
) -> Result<(), WebdriverDownloadError> {
if !reinstall && driver_info.is_installed().await {
println!("cargo::warning=Driver already installed ...");
Ok(())
} else {
if skip_verification {
driver_info.download_install().await?;
} else {
driver_info.download_verify_install(num_tries).await?;
}
println!("cargo::warning=Driver installed successfully ...");
Ok(())
}
}
fn main() -> Result<()> {
if cfg!(feature = "webdriver_download") {
println!("cargo:rerun-if-changed=src/lib.rs");
let webdriver_bin_dir = user_bin_dir();
println!(
"cargo::rerun-if-changed={}",
webdriver_bin_dir.to_string_lossy()
);
#[cfg(feature = "chromedriver")]
{
let config = WebdriverDownloadConfig {
driver_name: CHROMEDRIVER_NAME,
get_browser_path: get_chrome_path,
};
setup_driver(&config)?;
}
#[cfg(feature = "geckodriver")]
{
let config = WebdriverDownloadConfig {
driver_name: GECKODRIVER_NAME,
get_browser_path: get_firefox_path,
};
setup_driver(&config)?;
}
#[cfg(not(any(feature = "chromedriver", feature = "geckodriver")))]
{
println!("cargo::warning=No specific driver feature enabled, skipping driver setup");
}
} else {
#[cfg(feature = "chromedriver")]
{
let msg = format!("'webdriver_download' feature disabled. Please install a '{CHROMEDRIVER_NAME}' version manually and make the environment variable 'WEBDRIVER_PATH' point to it.");
println!("cargo::warning={msg}");
}
#[cfg(feature = "geckodriver")]
{
let msg = format!("'webdriver_download' feature disabled. Please install a '{GECKODRIVER_NAME}' version manually and make the environment variable 'WEBDRIVER_PATH' point to it.");
println!("cargo::warning={msg}");
}
#[cfg(not(any(feature = "chromedriver", feature = "geckodriver")))]
{
println!("cargo::warning='webdriver_download' feature disabled and no driver feature enabled");
}
}
Ok(())
}