diff --git a/Cargo.lock b/Cargo.lock index d78a5a2..521ba2e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -62,6 +62,12 @@ dependencies = [ "libc", ] +[[package]] +name = "anyhow" +version = "1.0.99" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0674a1ddeecb70197781e945de4b3b8ffb61fa939a5597bcf48503737663100" + [[package]] name = "async-compression" version = "0.4.29" @@ -1459,6 +1465,7 @@ dependencies = [ name = "proxy-scraper-checker" version = "0.1.0" dependencies = [ + "async-trait", "color-eyre", "crossterm", "dirs", @@ -1466,6 +1473,7 @@ dependencies = [ "foldhash 0.2.0", "futures", "hickory-resolver", + "http", "httpdate", "itertools 0.14.0", "maxminddb", @@ -1474,6 +1482,7 @@ dependencies = [ "rand", "ratatui", "reqwest", + "reqwest-middleware", "rlimit", "serde", "serde_json", @@ -1699,6 +1708,21 @@ dependencies = [ "webpki-roots", ] +[[package]] +name = "reqwest-middleware" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57f17d28a6e6acfe1733fe24bcd30774d13bffa4b8a22535b4c8c98423088d4e" +dependencies = [ + "anyhow", + "async-trait", + "http", + "reqwest", + "serde", + "thiserror 1.0.69", + "tower-service", +] + [[package]] name = "resolv-conf" version = "0.7.4" diff --git a/Cargo.toml b/Cargo.toml index 3b828e9..77d2304 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ license = "MIT" publish = false [dependencies] +async-trait = "=0.1.89" color-eyre = "=0.6.5" crossterm = { version = "=0.28.1", features = [ "event-stream", @@ -15,6 +16,7 @@ fancy-regex = "=0.16.1" foldhash = "=0.2.0" futures = { version = "=0.3.31", optional = true } hickory-resolver = "=0.25.2" +http = "=1.3.1" httpdate = "=1.0.3" itertools = "=0.14" maxminddb = { version = "=0.26.0", features = ["mmap"] } @@ -31,6 +33,7 @@ reqwest = { version = "=0.12.23", default-features = false, features = [ "socks", "system-proxy", ] } +reqwest-middleware = "=0.4.2" rlimit = "=0.10.2" serde = "=1.0.219" serde_json = "=1.0.143" diff --git a/src/checker.rs b/src/checker.rs index 6488f76..72110c3 100644 --- a/src/checker.rs +++ b/src/checker.rs @@ -57,7 +57,7 @@ pub async fn check_all( } Err(e) if tracing::event_enabled!(tracing::Level::DEBUG) => { tracing::debug!( - "{} | {}", + "{}: {}", proxy.to_string(true), pretty_error(&e) ); @@ -82,7 +82,10 @@ pub async fn check_all( match res { Ok(()) => {} Err(e) if e.is_panic() => { - tracing::error!("proxy checking task panicked: {}", e); + tracing::error!( + "Proxy checking task panicked: {}", + pretty_error(&e.into()) + ); } Err(e) => { return Err(e.into()); diff --git a/src/http.rs b/src/http.rs index 6916a45..92b755d 100644 --- a/src/http.rs +++ b/src/http.rs @@ -1,11 +1,11 @@ use std::{ - fmt::Display, + io, net::SocketAddr, sync::Arc, time::{Duration, SystemTime}, }; -use crate::{HashMap, config::Config}; +use crate::config::Config; const DEFAULT_MAX_RETRIES: u32 = 2; const INITIAL_RETRY_DELAY: Duration = Duration::from_millis(500); @@ -102,71 +102,55 @@ fn calculate_retry_timeout( Some(base.mul_f64(jitter)) } -pub async fn fetch_text( - http_client: reqwest::Client, - url: U, - basic_auth: Option<&BasicAuth>, - headers: Option<&HashMap>, -) -> crate::Result { - let mut attempt: u32 = 0; - loop { - let mut request = http_client.get(url.clone()); - if let Some(auth) = basic_auth { - request = - request.basic_auth(&auth.username, auth.password.as_ref()); - } - if let Some(headers) = headers { - for (k, v) in headers { - request = request.header(k, v); - } - } - match request.send().await { - Ok(resp) => { - let status = resp.status(); - if status.is_client_error() || status.is_server_error() { +pub struct RetryMiddleware; + +#[async_trait::async_trait] +impl reqwest_middleware::Middleware for RetryMiddleware { + async fn handle( + &self, + req: reqwest::Request, + extensions: &mut http::Extensions, + next: reqwest_middleware::Next<'_>, + ) -> reqwest_middleware::Result { + let mut attempt: u32 = 0; + loop { + let duplicate_request = req.try_clone().ok_or_else(|| { + reqwest_middleware::Error::middleware(io::Error::other( + "Request object is not cloneable", + )) + })?; + + match next.clone().run(duplicate_request, extensions).await { + Ok(resp) => { + let status = resp.status(); + if status.is_client_error() || status.is_server_error() { + if attempt < DEFAULT_MAX_RETRIES + && RETRY_STATUSES.contains(&status) + && let Some(delay) = calculate_retry_timeout( + Some(resp.headers()), + attempt, + ) + { + tokio::time::sleep(delay).await; + attempt = attempt.saturating_add(1); + continue; + } + resp.error_for_status_ref()?; + } + return Ok(resp); + } + Err(err) => { if attempt < DEFAULT_MAX_RETRIES - && RETRY_STATUSES.contains(&status) - && let Some(delay) = calculate_retry_timeout( - Some(resp.headers()), - attempt, - ) + && err.is_connect() + && let Some(delay) = + calculate_retry_timeout(None, attempt) { - tracing::info!( - "Request to {} returned status {}. Retrying \ - attempt {}/{} after {:?}", - url, - status, - attempt.saturating_add(1), - DEFAULT_MAX_RETRIES, - delay - ); tokio::time::sleep(delay).await; attempt = attempt.saturating_add(1); continue; } - resp.error_for_status_ref()?; + return Err(err); } - return Ok(resp.text().await?); - } - Err(err) => { - if attempt < DEFAULT_MAX_RETRIES - && err.is_connect() - && let Some(delay) = calculate_retry_timeout(None, attempt) - { - tracing::info!( - "Connection error while requesting {}: {}. Retrying \ - attempt {}/{} after {:?}", - url, - err, - attempt.saturating_add(1), - DEFAULT_MAX_RETRIES, - delay - ); - tokio::time::sleep(delay).await; - attempt = attempt.saturating_add(1); - continue; - } - return Err(err.into()); } } } @@ -175,7 +159,7 @@ pub async fn fetch_text( pub fn create_reqwest_client( config: &Config, dns_resolver: Arc, -) -> reqwest::Result { +) -> reqwest::Result { let mut builder = reqwest::ClientBuilder::new() .user_agent(&config.scraping.user_agent) .timeout(config.scraping.timeout) @@ -186,5 +170,10 @@ pub fn create_reqwest_client( builder = builder.proxy(reqwest::Proxy::all(proxy.clone())?); } - builder.build() + let client = builder.build()?; + let client_with_middleware = reqwest_middleware::ClientBuilder::new(client) + .with(RetryMiddleware) + .build(); + + Ok(client_with_middleware) } diff --git a/src/ipdb.rs b/src/ipdb.rs index d5e72d5..bea158d 100644 --- a/src/ipdb.rs +++ b/src/ipdb.rs @@ -121,7 +121,7 @@ impl DbType { pub async fn download( self, - http_client: reqwest::Client, + http_client: reqwest_middleware::ClientWithMiddleware, #[cfg(feature = "tui")] tx: tokio::sync::mpsc::UnboundedSender, ) -> crate::Result<()> { let db_path = self.db_path().await?; diff --git a/src/main.rs b/src/main.rs index 8ee651c..853ac0e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -105,7 +105,7 @@ fn create_logging_filter( async fn download_output_dependencies( config: &config::Config, - http_client: reqwest::Client, + http_client: reqwest_middleware::ClientWithMiddleware, token: tokio_util::sync::CancellationToken, #[cfg(feature = "tui")] tx: tokio::sync::mpsc::UnboundedSender< event::Event, @@ -236,7 +236,7 @@ fn watch_signals( tracing::warn!( "Failed to listen for {} signal: {}", signal_name, - e + utils::pretty_error(&e.into()) ); continue; } diff --git a/src/scraper.rs b/src/scraper.rs index 058c18c..17216a2 100644 --- a/src/scraper.rs +++ b/src/scraper.rs @@ -8,7 +8,6 @@ use crate::event::{AppEvent, Event}; use crate::{ HashSet, config::{Config, Source}, - http, parsers::PROXY_REGEX, proxy::{Proxy, ProxyType}, utils::pretty_error, @@ -16,7 +15,7 @@ use crate::{ async fn scrape_one( config: Arc, - http_client: reqwest::Client, + http_client: reqwest_middleware::ClientWithMiddleware, proto: ProxyType, proxies: Arc>>, source: Arc, @@ -25,13 +24,24 @@ async fn scrape_one( let text_result = if let Ok(u) = url::Url::parse(&source.url) { match u.scheme() { "http" | "https" => { - http::fetch_text( - http_client, - u, - source.basic_auth.as_ref(), - source.headers.as_ref(), - ) - .await + let mut request = http_client.get(u); + drop(http_client); + + if let Some(auth) = &source.basic_auth { + request = request + .basic_auth(&auth.username, auth.password.as_ref()); + } + + if let Some(headers) = &source.headers { + for (k, v) in headers { + request = request.header(k, v); + } + } + + match request.send().await { + Ok(resp) => resp.text().await.map_err(Into::into), + Err(e) => Err(e.into()), + } } _ => { drop(http_client); @@ -53,7 +63,7 @@ async fn scrape_one( let text = match text_result { Ok(text) => text, Err(e) => { - tracing::warn!("{} | {}", source.url, pretty_error(&e)); + tracing::warn!("{}: {}", source.url, pretty_error(&e)); return Ok(()); } }; @@ -64,7 +74,7 @@ async fn scrape_one( && i >= config.scraping.max_proxies_per_source { tracing::warn!( - "{} | Too many proxies (> {}) - skipped", + "{}: too many proxies (> {}) - skipped", source.url, config.scraping.max_proxies_per_source ); @@ -74,7 +84,7 @@ async fn scrape_one( } if matches.is_empty() { - tracing::warn!("{} | No proxies found", source.url); + tracing::warn!("{}: no proxies found", source.url); return Ok(()); } @@ -129,7 +139,7 @@ async fn scrape_one( pub async fn scrape_all( config: Arc, - http_client: reqwest::Client, + http_client: reqwest_middleware::ClientWithMiddleware, token: tokio_util::sync::CancellationToken, #[cfg(feature = "tui")] tx: tokio::sync::mpsc::UnboundedSender, ) -> crate::Result> {