diff --git a/src/downloader/downloader.rs b/src/downloader/downloader.rs index 927948b..048199e 100644 --- a/src/downloader/downloader.rs +++ b/src/downloader/downloader.rs @@ -580,7 +580,7 @@ impl Downloader { overwrite: Option, ) -> Result, DownloaderError> { Self::new2( - Arc::new(WebClient::default()), + Arc::new(WebClient::with_no_timeout()), url, headers, path, @@ -833,7 +833,7 @@ async fn test_failed_downloader() { } let url = "https://a.com/ssdassaodasdas"; let pb = p.join("addd"); - let client = Arc::new(WebClient::default()); + let client = Arc::new(WebClient::with_no_timeout()); let mut retry_interval = NonTailList::::default(); retry_interval += Duration::new(0, 0); client @@ -966,7 +966,7 @@ async fn test_failed_multi_downloader() { } let url = "https://a.com/ssdassaodasdas"; let pb = p.join("addds"); - let client = Arc::new(WebClient::default()); + let client = Arc::new(WebClient::with_no_timeout()); let mut retry_interval = NonTailList::::default(); retry_interval += Duration::new(0, 0); client diff --git a/src/fanbox_api.rs b/src/fanbox_api.rs index 96fbbb3..2e109e3 100644 --- a/src/fanbox_api.rs +++ b/src/fanbox_api.rs @@ -1,3 +1,4 @@ +use crate::error::PixivDownloaderError; use crate::ext::atomic::AtomicQuick; use crate::ext::replace::ReplaceWith2; use crate::ext::rw_lock::GetRwLock; @@ -9,14 +10,39 @@ use crate::fanbox::post::FanboxPost; use crate::gettext; use crate::opthelper::get_helper; use crate::parser::metadata::MetaDataParser; -use crate::webclient::WebClient; +use crate::webclient::{ReqMiddleware, WebClient}; use json::JsonValue; use proc_macros::fanbox_api_quick_test; -use reqwest::IntoUrl; +use reqwest::{Client, IntoUrl, Request, RequestBuilder}; use std::ops::Deref; use std::sync::atomic::AtomicBool; use std::sync::Arc; use std::sync::RwLock; +use std::time::Duration; + +pub struct FanboxDownloadDetectMiddleware { + _unused: [u8; 0], +} + +impl Default for FanboxDownloadDetectMiddleware { + fn default() -> Self { + Self { _unused: [] } + } +} + +impl ReqMiddleware for FanboxDownloadDetectMiddleware { + fn handle(&self, r: Request, c: Client) -> Result { + let is_downloads_url = r.url().host_str().unwrap_or("") == "downloads.fanbox.cc"; + Ok(if is_downloads_url { + log::debug!(target: "fanbox_api", "Disable request timeout for {}", r.url()); + RequestBuilder::from_parts(c, r) + .timeout(Duration::MAX) + .build()? + } else { + r + }) + } +} /// Fanbox API client pub struct FanboxClientInternal { @@ -90,6 +116,8 @@ impl FanboxClientInternal { for (k, v) in headers.iter() { self.client.set_header(k, v); } + self.client + .add_req_middleware(Box::new(FanboxDownloadDetectMiddleware::default())); self.inited.qstore(true); true } diff --git a/src/opthelper.rs b/src/opthelper.rs index cd3742c..3a963c3 100644 --- a/src/opthelper.rs +++ b/src/opthelper.rs @@ -691,6 +691,30 @@ impl OptHelper { None => false, } } + + /// Set a timeout for only the connect phase of a client. + pub fn connect_timeout(&self) -> Duration { + let t = self + .opt + .get_ref() + .connect_timeout + .or_else(|| self.settings.get_ref().get_u64("connect-timeout")) + .unwrap_or(10_000); + Duration::from_millis(t) + } + + /// Set request timeout in milliseconds. + /// The timeout is applied from when the request starts connecting until the response body + /// has finished. Not used for downloader. + pub fn client_timeout(&self) -> Duration { + Duration::from_millis( + self.opt + .get_ref() + .client_timeout + .or_else(|| self.settings.get_ref().get_u64("client-timeout")) + .unwrap_or(30_000), + ) + } } impl Default for OptHelper { diff --git a/src/opts.rs b/src/opts.rs index b8cd7cd..8ac7646 100644 --- a/src/opts.rs +++ b/src/opts.rs @@ -138,6 +138,12 @@ pub struct CommandOpts { #[cfg(feature = "ugoira")] /// Whether to use ugoira cli. pub ugoira_cli: Option, + /// Set a timeout in milliseconds for only the connect phase of a client. + pub connect_timeout: Option, + /// Set request timeout in milliseconds. + /// The timeout is applied from when the request starts connecting until the response body + /// has finished. Not used for downloader. + pub client_timeout: Option, } impl CommandOpts { @@ -190,6 +196,8 @@ impl CommandOpts { ugoira: None, #[cfg(feature = "ugoira")] ugoira_cli: None, + connect_timeout: None, + client_timeout: None, } } @@ -356,6 +364,19 @@ pub fn parse_u64>(s: Option) -> Result, ParseIntErr } } +/// Prase Non Zero [u64] from string +pub fn parse_non_zero_u64>(s: Option) -> Result, ParseIntError> { + match s { + Some(s) => { + let s = s.as_ref(); + let s = s.trim(); + let c = s.parse::()?; + Ok(Some(c.get())) + } + None => Ok(None), + } +} + pub fn parse_nonempty_usize>(s: Option) -> Result, ParseIntError> { match s { Some(s) => { @@ -709,6 +730,13 @@ pub fn parse_cmd() -> Option { HasArg::Maybe, getopts::Occur::Optional, ); + opts.optopt( + "", + "connect-timeout", + gettext("Set a timeout in milliseconds for only the connect phase of a client."), + "TIME", + ); + opts.optopt("", "client-timeout", gettext("Set request timeout in milliseconds. The timeout is applied from when the request starts connecting until the response body has finished. Not used for downloader."), "TIME"); let result = match opts.parse(&argv[1..]) { Ok(m) => m, Err(err) => { @@ -1153,6 +1181,32 @@ pub fn parse_cmd() -> Option { return None; } } + match parse_non_zero_u64(result.opt_str("connect-timeout")) { + Ok(r) => re.as_mut().unwrap().connect_timeout = r, + Err(e) => { + log::error!( + "{} {}", + gettext("Failed to parse :") + .replace("", "connect-timeout") + .as_str(), + e + ); + return None; + } + } + match parse_non_zero_u64(result.opt_str("client-timeout")) { + Ok(r) => re.as_mut().unwrap().client_timeout = r, + Err(e) => { + log::error!( + "{} {}", + gettext("Failed to parse :") + .replace("", "client-timeout") + .as_str(), + e + ); + return None; + } + } re } diff --git a/src/settings.rs b/src/settings.rs index cf9cc80..aca8bc3 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -312,6 +312,13 @@ impl SettingStore { } } + pub fn get_u64(&self, key: &str) -> Option { + match self.data.get(key) { + Some(obj) => obj.as_u64(), + None => None, + } + } + pub fn have(&self, key: &str) -> bool { self.data.have(key) } diff --git a/src/settings_list.rs b/src/settings_list.rs index 7b22c83..b43fa02 100644 --- a/src/settings_list.rs +++ b/src/settings_list.rs @@ -74,6 +74,8 @@ pub fn get_settings_list() -> Vec { SettingDes::new("ugoira", gettext("The path to ugoira cli executable."), JsonValueType::Str, None).unwrap(), #[cfg(feature = "ugoira")] SettingDes::new("ugoira-cli", gettext("Whether to use ugoira cli."), JsonValueType::Boolean, None).unwrap(), + SettingDes::new("connect-timeout", gettext("Set a timeout in milliseconds for only the connect phase of a client."), JsonValueType::Number, Some(check_nonzero_u64)).unwrap(), + SettingDes::new("client-timeout", gettext("Set request timeout in milliseconds. The timeout is applied from when the request starts connecting until the response body has finished. Not used for downloader."), JsonValueType::Number, Some(check_nonzero_u64)).unwrap(), ] } @@ -114,6 +116,14 @@ fn check_u64(obj: &JsonValue) -> bool { r.is_some() } +fn check_nonzero_u64(obj: &JsonValue) -> bool { + let r = obj.as_u64(); + match r { + Some(u) => u > 0, + None => false, + } +} + #[inline] fn check_parse_size_u32(obj: &JsonValue) -> bool { parse_u32_size(obj).is_some() diff --git a/src/webclient.rs b/src/webclient.rs index e429ef4..abcce85 100644 --- a/src/webclient.rs +++ b/src/webclient.rs @@ -3,6 +3,7 @@ use crate::cookies::ManagedCookieJar; use crate::error::PixivDownloaderError; use crate::ext::atomic::AtomicQuick; use crate::ext::json::ToJson; +use crate::ext::replace::ReplaceWith2; use crate::ext::rw_lock::GetRwLock; use crate::formdata::FormData; use crate::gettext; @@ -107,6 +108,9 @@ pub struct WebClient { retry_interval: RwLock>>, /// Request middlewares req_middlewares: RwLock>>, + /// Set request timeout. The timeout is applied from when the request starts connecting until + /// the response body has finished. + timeout: RwLock>, } impl WebClient { @@ -121,6 +125,7 @@ impl WebClient { retry: Arc::new(AtomicI64::new(3)), retry_interval: RwLock::new(None), req_middlewares: RwLock::new(Vec::new()), + timeout: RwLock::new(None), } } @@ -224,6 +229,19 @@ impl WebClient { self.retry.qstore(retry) } + /// Set request timeout. The timeout is applied from when the request starts connecting until + /// the response body has finished. + pub fn set_timeout(&self, timeout: Option) { + self.timeout.replace_with2(timeout); + } + + /// Create a client with no timeout set. + pub fn with_no_timeout() -> Self { + let c = Self::default(); + c.set_timeout(None); + c + } + /// Send GET requests with parameters /// * `param` - GET parameters. Should be a JSON object/array. If value in map is not a string, will dump it /// # Examples @@ -567,6 +585,7 @@ impl Default for WebClient { if !chain.is_empty() { c = c.proxy(reqwest::Proxy::custom(move |url| chain.r#match(url))); } + c = c.connect_timeout(opt.connect_timeout()); let c = c.build().unwrap(); let c = Self::new(c); match opt.retry() { @@ -574,6 +593,7 @@ impl Default for WebClient { None => {} } c.get_retry_interval_as_mut().replace(opt.retry_interval()); + c.set_timeout(Some(opt.client_timeout())); c } }