[rust] Support for web proxy in Selenium Manager (#11575)

* [rust] Support for web proxy in Selenium Manager * [rust] Include CLI tests for web proxy * [rust] Include flag for network requests timeout (30 seconds by default) * [rust] Check parse error in fallback for chromedriver * [rust] Increase default request timeout to 60 seconds * [rust] Increase request timeout to 120 seconds * [rust] Improve proxy and timeout tests * [rust] Include --clear-cache and --debug in timeout test * [rust] Include test using mock proxy
SeleniumHQ · Jan 30, 2023 · 8815c27 · 8815c27
1 parent 39ceed7
commit 8815c27
Show file tree

Hide file tree

Showing 15 changed files with 1,823 additions and 44 deletions.
diff --git a/rust/Cargo.Bazel.lock b/rust/Cargo.Bazel.lock
diff --git a/rust/Cargo.lock b/rust/Cargo.lock
diff --git a/rust/Cargo.toml b/rust/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "selenium-manager"
-version = "1.0.0-M2"
+version = "1.0.0-M3"
 edition = "2021"
 authors = ["Selenium <selenium-developers@googlegroups.com"]
 license = "Apache-2.0"
@@ -31,6 +31,7 @@ exitcode = "1.1.2"
 [dev-dependencies]
 assert_cmd = "2.0.7"
 rstest = "0.16.0"
+wiremock = "0.5.17"
 
 [profile.release]
 opt-level = 'z'     # Optimize for size

diff --git a/rust/README.md b/rust/README.md
@@ -16,7 +16,7 @@ Selenium Manager can be executed using Cargo as follows:
 
 ```
 $ cargo run -- --help
-selenium-manager 1.0.0-M2
+selenium-manager 1.0.0-M3
 Selenium Manager is a CLI tool that automatically manages the browser/driver infrastructure required by Selenium.
 
 Usage: selenium-manager [OPTIONS]
@@ -31,6 +31,10 @@ Options:
           Major browser version (e.g., 105, 106, etc. Also: beta, dev, canary -or nightly- is accepted)
   -P, --browser-path <BROWSER_PATH>
           Browser path (absolute) for browser version detection (e.g., /usr/bin/google-chrome, "/Applications/Google\ Chrome.app/Contents/MacOS/Google\ Chrome", "C:\\Program Files\\Google\\Chrome\\Application\\chrome.exe")
+  -p, --proxy <PROXY>
+          HTTP proxy for network connection (e.g., https://myproxy.net:8080)
+  -t, --timeout <TIMEOUT>
+          Timeout for network requests (in seconds) [default: 120]
   -D, --debug
           Display DEBUG messages
   -T, --trace

diff --git a/rust/src/chrome.rs b/rust/src/chrome.rs
@@ -24,7 +24,7 @@ use std::path::PathBuf;
 use crate::config::ARCH::ARM64;
 use crate::config::OS::{LINUX, MACOS, WINDOWS};
 use crate::downloads::read_content_from_link;
-use crate::files::{compose_driver_path_in_cache, BrowserPath};
+use crate::files::{compose_driver_path_in_cache, BrowserPath, PARSE_ERROR};
 use crate::logger::Logger;
 use crate::metadata::{
     create_driver_metadata, get_driver_version_from_metadata, get_metadata, write_metadata,
@@ -69,6 +69,10 @@ impl SeleniumManager for ChromeManager {
         &self.http_client
     }
 
+    fn set_http_client(&mut self, http_client: Client) {
+        self.http_client = http_client;
+    }
+
     fn get_browser_path_map(&self) -> HashMap<BrowserPath, &str> {
         HashMap::from([
             (
@@ -180,13 +184,15 @@ impl SeleniumManager for ChromeManager {
                         "Reading {} version from {}",
                         &self.driver_name, driver_url
                     ));
-                    let content = read_content_from_link(self.get_http_client(), driver_url);
-                    match content {
+                    match read_content_from_link(self.get_http_client(), driver_url) {
                         Ok(version) => {
                             driver_version = version;
                             break;
                         }
-                        _ => {
+                        Err(err) => {
+                            if !err.to_string().eq(PARSE_ERROR) {
+                                return Err(err);
+                            }
                             self.log.warn(format!(
                                 "Error getting version of {} {}. Retrying with {} {} (attempt {}/{})",
                                 &self.driver_name,

diff --git a/rust/src/config.rs b/rust/src/config.rs
@@ -16,6 +16,7 @@
 // under the License.
 
 use crate::config::OS::{LINUX, MACOS, WINDOWS};
+use crate::REQUEST_TIMEOUT_SEC;
 use std::env::consts::{ARCH, OS};
 
 pub struct ManagerConfig {
@@ -24,6 +25,8 @@ pub struct ManagerConfig {
     pub os: String,
     pub arch: String,
     pub browser_path: String,
+    pub proxy: String,
+    pub timeout: u64,
 }
 
 impl ManagerConfig {
@@ -34,6 +37,8 @@ impl ManagerConfig {
             os: OS.to_string(),
             arch: ARCH.to_string(),
             browser_path: "".to_string(),
+            proxy: "".to_string(),
+            timeout: REQUEST_TIMEOUT_SEC,
         }
     }
 
@@ -45,6 +50,8 @@ impl ManagerConfig {
             os: config.os.as_str().to_string(),
             arch: config.arch.as_str().to_string(),
             browser_path: config.browser_path.as_str().to_string(),
+            proxy: config.proxy.as_str().to_string(),
+            timeout: config.timeout,
         }
     }
 }

diff --git a/rust/src/edge.rs b/rust/src/edge.rs
@@ -69,6 +69,10 @@ impl SeleniumManager for EdgeManager {
         &self.http_client
     }
 
+    fn set_http_client(&mut self, http_client: Client) {
+        self.http_client = http_client;
+    }
+
     fn get_browser_path_map(&self) -> HashMap<BrowserPath, &str> {
         HashMap::from([
             (

diff --git a/rust/src/files.rs b/rust/src/files.rs
@@ -32,6 +32,7 @@ use zip::ZipArchive;
 use crate::config::OS::WINDOWS;
 use crate::Logger;
 
+pub const PARSE_ERROR: &str = "Wrong browser/driver version";
 const CACHE_FOLDER: &str = ".cache/selenium";
 const ZIP: &str = "zip";
 const GZ: &str = "gz";
@@ -180,7 +181,7 @@ pub fn get_binary_extension(os: &str) -> &str {
 
 pub fn parse_version(version_text: String) -> Result<String, Box<dyn Error>> {
     if version_text.to_ascii_lowercase().contains("error") {
-        return Err("Wrong browser/driver version".into());
+        return Err(PARSE_ERROR.into());
     }
     let re = Regex::new(r"[^\d^.]").unwrap();
     Ok(re.replace_all(&version_text, "").to_string())

diff --git a/rust/src/firefox.rs b/rust/src/firefox.rs
@@ -67,6 +67,10 @@ impl SeleniumManager for FirefoxManager {
         &self.http_client
     }
 
+    fn set_http_client(&mut self, http_client: Client) {
+        self.http_client = http_client;
+    }
+
     fn get_browser_path_map(&self) -> HashMap<BrowserPath, &str> {
         HashMap::from([
             (

diff --git a/rust/src/iexplorer.rs b/rust/src/iexplorer.rs
@@ -64,6 +64,10 @@ impl SeleniumManager for IExplorerManager {
         &self.http_client
     }
 
+    fn set_http_client(&mut self, http_client: Client) {
+        self.http_client = http_client;
+    }
+
     fn get_browser_path_map(&self) -> HashMap<BrowserPath, &str> {
         HashMap::new()
     }

diff --git a/rust/src/lib.rs b/rust/src/lib.rs
@@ -23,11 +23,12 @@ use crate::iexplorer::IExplorerManager;
 use std::fs;
 
 use crate::config::{str_to_os, ManagerConfig};
-use reqwest::Client;
+use reqwest::{Client, ClientBuilder, Proxy};
 use std::collections::HashMap;
 use std::error::Error;
 use std::path::PathBuf;
 use std::process::Command;
+use std::time::Duration;
 
 use crate::downloads::download_driver_to_tmp_folder;
 use crate::files::{parse_version, uncompress, BrowserPath};
@@ -46,6 +47,7 @@ pub mod iexplorer;
 pub mod logger;
 pub mod metadata;
 
+pub const REQUEST_TIMEOUT_SEC: u64 = 120; // The timeout is applied from when the request starts connecting until the response body has finished
 pub const STABLE: &str = "stable";
 pub const BETA: &str = "beta";
 pub const DEV: &str = "dev";
@@ -70,6 +72,8 @@ pub trait SeleniumManager {
 
     fn get_http_client(&self) -> &Client;
 
+    fn set_http_client(&mut self, http_client: Client);
+
     fn get_browser_path_map(&self) -> HashMap<BrowserPath, &str>;
 
     fn discover_browser_version(&self) -> Option<String>;
@@ -191,9 +195,20 @@ pub trait SeleniumManager {
                 }
             }
         }
-        let driver_version = self
-            .request_driver_version()
-            .unwrap_or_else(|err| err.to_string());
+        let driver_version = match self.request_driver_version() {
+            Ok(version) => {
+                if version.is_empty() {
+                    return Err(format!(
+                        "The {} version cannot be discovered",
+                        self.get_driver_name()
+                    ));
+                }
+                version
+            }
+            Err(err) => {
+                return Err(err.to_string());
+            }
+        };
         self.get_logger().debug(format!(
             "Required driver: {} {}",
             self.get_driver_name(),
@@ -312,6 +327,52 @@ pub trait SeleniumManager {
         config.browser_path = browser_path;
         self.set_config(config);
     }
+
+    fn get_proxy(&self) -> &str {
+        self.get_config().proxy.as_str()
+    }
+
+    fn set_proxy(&mut self, proxy: String) -> Result<(), Box<dyn Error>> {
+        let mut config = ManagerConfig::clone(self.get_config());
+        config.proxy = proxy.to_string();
+        self.set_config(config);
+
+        if !proxy.is_empty() {
+            self.get_logger().debug(format!("Using proxy: {}", &proxy));
+            self.update_http_proxy()?;
+        }
+        Ok(())
+    }
+
+    fn get_timeout(&self) -> u64 {
+        self.get_config().timeout
+    }
+
+    fn set_timeout(&mut self, timeout: u64) -> Result<(), Box<dyn Error>> {
+        let mut config = ManagerConfig::clone(self.get_config());
+        config.timeout = timeout;
+        self.set_config(config);
+
+        if timeout != REQUEST_TIMEOUT_SEC {
+            self.get_logger()
+                .debug(format!("Using timeout of {} seconds", timeout));
+            self.update_http_proxy()?;
+        }
+        Ok(())
+    }
+
+    fn update_http_proxy(&mut self) -> Result<(), Box<dyn Error>> {
+        let proxy = self.get_proxy();
+        let timeout = self.get_timeout();
+
+        let mut builder = http_client_builder().timeout(Duration::from_secs(timeout));
+        if !proxy.is_empty() {
+            builder = builder.proxy(Proxy::all(proxy)?);
+        }
+        let http_client = builder.build()?;
+        self.set_http_client(http_client);
+        Ok(())
+    }
 }
 
 // ----------------------------------------------------------
@@ -370,9 +431,8 @@ pub fn clear_cache(log: &Logger) {
 }
 
 pub fn create_default_http_client() -> Client {
-    Client::builder()
-        .danger_accept_invalid_certs(true)
-        .use_rustls_tls()
+    http_client_builder()
+        .timeout(Duration::from_secs(REQUEST_TIMEOUT_SEC))
         .build()
         .unwrap_or_default()
 }
@@ -388,3 +448,9 @@ fn get_index_version(full_version: &str, index: usize) -> Result<String, Box<dyn
         .ok_or(format!("Wrong version: {}", full_version))?
         .to_string())
 }
+
+pub fn http_client_builder() -> ClientBuilder {
+    Client::builder()
+        .danger_accept_invalid_certs(true)
+        .use_rustls_tls()
+}
diff --git a/rust/src/main.rs b/rust/src/main.rs
@@ -21,9 +21,10 @@ use std::process::exit;
 
 use clap::Parser;
 
-use exitcode::DATAERR;
+use exitcode::{DATAERR, UNAVAILABLE};
 
 use selenium_manager::logger::Logger;
+use selenium_manager::REQUEST_TIMEOUT_SEC;
 use selenium_manager::{
     clear_cache, get_manager_by_browser, get_manager_by_driver, SeleniumManager,
 };
@@ -62,6 +63,14 @@ struct Cli {
     #[clap(short = 'O', long, value_parser, default_value = "LOGGER")]
     output: String,
 
+    /// HTTP proxy for network connection (e.g., https://myproxy.net:8080)
+    #[clap(short = 'p', long, value_parser)]
+    proxy: Option<String>,
+
+    /// Timeout for network requests (in seconds)
+    #[clap(short = 't', long, value_parser, default_value_t = REQUEST_TIMEOUT_SEC)]
+    timeout: u64,
+
     /// Display DEBUG messages
     #[clap(short = 'D', long)]
     debug: bool,
@@ -108,6 +117,20 @@ fn main() -> Result<(), Box<dyn Error>> {
     selenium_manager.set_browser_version(cli.browser_version.unwrap_or_default());
     selenium_manager.set_driver_version(cli.driver_version.unwrap_or_default());
     selenium_manager.set_browser_path(cli.browser_path.unwrap_or_default());
+    match selenium_manager.set_timeout(cli.timeout) {
+        Ok(_) => {}
+        Err(err) => {
+            selenium_manager.get_logger().error(err.to_string());
+            flush_and_exit(UNAVAILABLE, selenium_manager.get_logger());
+        }
+    }
+    match selenium_manager.set_proxy(cli.proxy.unwrap_or_default()) {
+        Ok(_) => {}
+        Err(err) => {
+            selenium_manager.get_logger().error(err.to_string());
+            flush_and_exit(UNAVAILABLE, selenium_manager.get_logger());
+        }
+    }
 
     match selenium_manager.resolve_driver() {
         Ok(driver_path) => {

diff --git a/rust/tests/output_tests.rs b/rust/tests/output_tests.rs
@@ -57,4 +57,4 @@ fn shell_output_test() {
 
     let driver = Path::new(output);
     assert!(driver.exists());
-}
+}