Skip to content

Commit 8815c27

Browse files
authored
[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
1 parent 39ceed7 commit 8815c27

File tree

15 files changed

+1823
-44
lines changed

15 files changed

+1823
-44
lines changed

rust/Cargo.Bazel.lock

Lines changed: 1349 additions & 18 deletions
Large diffs are not rendered by default.

rust/Cargo.lock

Lines changed: 247 additions & 10 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

rust/Cargo.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "selenium-manager"
3-
version = "1.0.0-M2"
3+
version = "1.0.0-M3"
44
edition = "2021"
55
authors = ["Selenium <selenium-developers@googlegroups.com"]
66
license = "Apache-2.0"
@@ -31,6 +31,7 @@ exitcode = "1.1.2"
3131
[dev-dependencies]
3232
assert_cmd = "2.0.7"
3333
rstest = "0.16.0"
34+
wiremock = "0.5.17"
3435

3536
[profile.release]
3637
opt-level = 'z' # Optimize for size

rust/README.md

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ Selenium Manager can be executed using Cargo as follows:
1616

1717
```
1818
$ cargo run -- --help
19-
selenium-manager 1.0.0-M2
19+
selenium-manager 1.0.0-M3
2020
Selenium Manager is a CLI tool that automatically manages the browser/driver infrastructure required by Selenium.
2121
2222
Usage: selenium-manager [OPTIONS]
@@ -31,6 +31,10 @@ Options:
3131
Major browser version (e.g., 105, 106, etc. Also: beta, dev, canary -or nightly- is accepted)
3232
-P, --browser-path <BROWSER_PATH>
3333
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")
34+
-p, --proxy <PROXY>
35+
HTTP proxy for network connection (e.g., https://myproxy.net:8080)
36+
-t, --timeout <TIMEOUT>
37+
Timeout for network requests (in seconds) [default: 120]
3438
-D, --debug
3539
Display DEBUG messages
3640
-T, --trace

rust/src/chrome.rs

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ use std::path::PathBuf;
2424
use crate::config::ARCH::ARM64;
2525
use crate::config::OS::{LINUX, MACOS, WINDOWS};
2626
use crate::downloads::read_content_from_link;
27-
use crate::files::{compose_driver_path_in_cache, BrowserPath};
27+
use crate::files::{compose_driver_path_in_cache, BrowserPath, PARSE_ERROR};
2828
use crate::logger::Logger;
2929
use crate::metadata::{
3030
create_driver_metadata, get_driver_version_from_metadata, get_metadata, write_metadata,
@@ -69,6 +69,10 @@ impl SeleniumManager for ChromeManager {
6969
&self.http_client
7070
}
7171

72+
fn set_http_client(&mut self, http_client: Client) {
73+
self.http_client = http_client;
74+
}
75+
7276
fn get_browser_path_map(&self) -> HashMap<BrowserPath, &str> {
7377
HashMap::from([
7478
(
@@ -180,13 +184,15 @@ impl SeleniumManager for ChromeManager {
180184
"Reading {} version from {}",
181185
&self.driver_name, driver_url
182186
));
183-
let content = read_content_from_link(self.get_http_client(), driver_url);
184-
match content {
187+
match read_content_from_link(self.get_http_client(), driver_url) {
185188
Ok(version) => {
186189
driver_version = version;
187190
break;
188191
}
189-
_ => {
192+
Err(err) => {
193+
if !err.to_string().eq(PARSE_ERROR) {
194+
return Err(err);
195+
}
190196
self.log.warn(format!(
191197
"Error getting version of {} {}. Retrying with {} {} (attempt {}/{})",
192198
&self.driver_name,

rust/src/config.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
// under the License.
1717

1818
use crate::config::OS::{LINUX, MACOS, WINDOWS};
19+
use crate::REQUEST_TIMEOUT_SEC;
1920
use std::env::consts::{ARCH, OS};
2021

2122
pub struct ManagerConfig {
@@ -24,6 +25,8 @@ pub struct ManagerConfig {
2425
pub os: String,
2526
pub arch: String,
2627
pub browser_path: String,
28+
pub proxy: String,
29+
pub timeout: u64,
2730
}
2831

2932
impl ManagerConfig {
@@ -34,6 +37,8 @@ impl ManagerConfig {
3437
os: OS.to_string(),
3538
arch: ARCH.to_string(),
3639
browser_path: "".to_string(),
40+
proxy: "".to_string(),
41+
timeout: REQUEST_TIMEOUT_SEC,
3742
}
3843
}
3944

@@ -45,6 +50,8 @@ impl ManagerConfig {
4550
os: config.os.as_str().to_string(),
4651
arch: config.arch.as_str().to_string(),
4752
browser_path: config.browser_path.as_str().to_string(),
53+
proxy: config.proxy.as_str().to_string(),
54+
timeout: config.timeout,
4855
}
4956
}
5057
}

rust/src/edge.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,10 @@ impl SeleniumManager for EdgeManager {
6969
&self.http_client
7070
}
7171

72+
fn set_http_client(&mut self, http_client: Client) {
73+
self.http_client = http_client;
74+
}
75+
7276
fn get_browser_path_map(&self) -> HashMap<BrowserPath, &str> {
7377
HashMap::from([
7478
(

rust/src/files.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ use zip::ZipArchive;
3232
use crate::config::OS::WINDOWS;
3333
use crate::Logger;
3434

35+
pub const PARSE_ERROR: &str = "Wrong browser/driver version";
3536
const CACHE_FOLDER: &str = ".cache/selenium";
3637
const ZIP: &str = "zip";
3738
const GZ: &str = "gz";
@@ -180,7 +181,7 @@ pub fn get_binary_extension(os: &str) -> &str {
180181

181182
pub fn parse_version(version_text: String) -> Result<String, Box<dyn Error>> {
182183
if version_text.to_ascii_lowercase().contains("error") {
183-
return Err("Wrong browser/driver version".into());
184+
return Err(PARSE_ERROR.into());
184185
}
185186
let re = Regex::new(r"[^\d^.]").unwrap();
186187
Ok(re.replace_all(&version_text, "").to_string())

rust/src/firefox.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,10 @@ impl SeleniumManager for FirefoxManager {
6767
&self.http_client
6868
}
6969

70+
fn set_http_client(&mut self, http_client: Client) {
71+
self.http_client = http_client;
72+
}
73+
7074
fn get_browser_path_map(&self) -> HashMap<BrowserPath, &str> {
7175
HashMap::from([
7276
(

rust/src/iexplorer.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,10 @@ impl SeleniumManager for IExplorerManager {
6464
&self.http_client
6565
}
6666

67+
fn set_http_client(&mut self, http_client: Client) {
68+
self.http_client = http_client;
69+
}
70+
6771
fn get_browser_path_map(&self) -> HashMap<BrowserPath, &str> {
6872
HashMap::new()
6973
}

rust/src/lib.rs

Lines changed: 73 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,12 @@ use crate::iexplorer::IExplorerManager;
2323
use std::fs;
2424

2525
use crate::config::{str_to_os, ManagerConfig};
26-
use reqwest::Client;
26+
use reqwest::{Client, ClientBuilder, Proxy};
2727
use std::collections::HashMap;
2828
use std::error::Error;
2929
use std::path::PathBuf;
3030
use std::process::Command;
31+
use std::time::Duration;
3132

3233
use crate::downloads::download_driver_to_tmp_folder;
3334
use crate::files::{parse_version, uncompress, BrowserPath};
@@ -46,6 +47,7 @@ pub mod iexplorer;
4647
pub mod logger;
4748
pub mod metadata;
4849

50+
pub const REQUEST_TIMEOUT_SEC: u64 = 120; // The timeout is applied from when the request starts connecting until the response body has finished
4951
pub const STABLE: &str = "stable";
5052
pub const BETA: &str = "beta";
5153
pub const DEV: &str = "dev";
@@ -70,6 +72,8 @@ pub trait SeleniumManager {
7072

7173
fn get_http_client(&self) -> &Client;
7274

75+
fn set_http_client(&mut self, http_client: Client);
76+
7377
fn get_browser_path_map(&self) -> HashMap<BrowserPath, &str>;
7478

7579
fn discover_browser_version(&self) -> Option<String>;
@@ -191,9 +195,20 @@ pub trait SeleniumManager {
191195
}
192196
}
193197
}
194-
let driver_version = self
195-
.request_driver_version()
196-
.unwrap_or_else(|err| err.to_string());
198+
let driver_version = match self.request_driver_version() {
199+
Ok(version) => {
200+
if version.is_empty() {
201+
return Err(format!(
202+
"The {} version cannot be discovered",
203+
self.get_driver_name()
204+
));
205+
}
206+
version
207+
}
208+
Err(err) => {
209+
return Err(err.to_string());
210+
}
211+
};
197212
self.get_logger().debug(format!(
198213
"Required driver: {} {}",
199214
self.get_driver_name(),
@@ -312,6 +327,52 @@ pub trait SeleniumManager {
312327
config.browser_path = browser_path;
313328
self.set_config(config);
314329
}
330+
331+
fn get_proxy(&self) -> &str {
332+
self.get_config().proxy.as_str()
333+
}
334+
335+
fn set_proxy(&mut self, proxy: String) -> Result<(), Box<dyn Error>> {
336+
let mut config = ManagerConfig::clone(self.get_config());
337+
config.proxy = proxy.to_string();
338+
self.set_config(config);
339+
340+
if !proxy.is_empty() {
341+
self.get_logger().debug(format!("Using proxy: {}", &proxy));
342+
self.update_http_proxy()?;
343+
}
344+
Ok(())
345+
}
346+
347+
fn get_timeout(&self) -> u64 {
348+
self.get_config().timeout
349+
}
350+
351+
fn set_timeout(&mut self, timeout: u64) -> Result<(), Box<dyn Error>> {
352+
let mut config = ManagerConfig::clone(self.get_config());
353+
config.timeout = timeout;
354+
self.set_config(config);
355+
356+
if timeout != REQUEST_TIMEOUT_SEC {
357+
self.get_logger()
358+
.debug(format!("Using timeout of {} seconds", timeout));
359+
self.update_http_proxy()?;
360+
}
361+
Ok(())
362+
}
363+
364+
fn update_http_proxy(&mut self) -> Result<(), Box<dyn Error>> {
365+
let proxy = self.get_proxy();
366+
let timeout = self.get_timeout();
367+
368+
let mut builder = http_client_builder().timeout(Duration::from_secs(timeout));
369+
if !proxy.is_empty() {
370+
builder = builder.proxy(Proxy::all(proxy)?);
371+
}
372+
let http_client = builder.build()?;
373+
self.set_http_client(http_client);
374+
Ok(())
375+
}
315376
}
316377

317378
// ----------------------------------------------------------
@@ -370,9 +431,8 @@ pub fn clear_cache(log: &Logger) {
370431
}
371432

372433
pub fn create_default_http_client() -> Client {
373-
Client::builder()
374-
.danger_accept_invalid_certs(true)
375-
.use_rustls_tls()
434+
http_client_builder()
435+
.timeout(Duration::from_secs(REQUEST_TIMEOUT_SEC))
376436
.build()
377437
.unwrap_or_default()
378438
}
@@ -388,3 +448,9 @@ fn get_index_version(full_version: &str, index: usize) -> Result<String, Box<dyn
388448
.ok_or(format!("Wrong version: {}", full_version))?
389449
.to_string())
390450
}
451+
452+
pub fn http_client_builder() -> ClientBuilder {
453+
Client::builder()
454+
.danger_accept_invalid_certs(true)
455+
.use_rustls_tls()
456+
}

rust/src/main.rs

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,10 @@ use std::process::exit;
2121

2222
use clap::Parser;
2323

24-
use exitcode::DATAERR;
24+
use exitcode::{DATAERR, UNAVAILABLE};
2525

2626
use selenium_manager::logger::Logger;
27+
use selenium_manager::REQUEST_TIMEOUT_SEC;
2728
use selenium_manager::{
2829
clear_cache, get_manager_by_browser, get_manager_by_driver, SeleniumManager,
2930
};
@@ -62,6 +63,14 @@ struct Cli {
6263
#[clap(short = 'O', long, value_parser, default_value = "LOGGER")]
6364
output: String,
6465

66+
/// HTTP proxy for network connection (e.g., https://myproxy.net:8080)
67+
#[clap(short = 'p', long, value_parser)]
68+
proxy: Option<String>,
69+
70+
/// Timeout for network requests (in seconds)
71+
#[clap(short = 't', long, value_parser, default_value_t = REQUEST_TIMEOUT_SEC)]
72+
timeout: u64,
73+
6574
/// Display DEBUG messages
6675
#[clap(short = 'D', long)]
6776
debug: bool,
@@ -108,6 +117,20 @@ fn main() -> Result<(), Box<dyn Error>> {
108117
selenium_manager.set_browser_version(cli.browser_version.unwrap_or_default());
109118
selenium_manager.set_driver_version(cli.driver_version.unwrap_or_default());
110119
selenium_manager.set_browser_path(cli.browser_path.unwrap_or_default());
120+
match selenium_manager.set_timeout(cli.timeout) {
121+
Ok(_) => {}
122+
Err(err) => {
123+
selenium_manager.get_logger().error(err.to_string());
124+
flush_and_exit(UNAVAILABLE, selenium_manager.get_logger());
125+
}
126+
}
127+
match selenium_manager.set_proxy(cli.proxy.unwrap_or_default()) {
128+
Ok(_) => {}
129+
Err(err) => {
130+
selenium_manager.get_logger().error(err.to_string());
131+
flush_and_exit(UNAVAILABLE, selenium_manager.get_logger());
132+
}
133+
}
111134

112135
match selenium_manager.resolve_driver() {
113136
Ok(driver_path) => {

rust/tests/output_tests.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,4 +57,4 @@ fn shell_output_test() {
5757

5858
let driver = Path::new(output);
5959
assert!(driver.exists());
60-
}
60+
}

0 commit comments

Comments
 (0)