Add timeout for web client

src/downloader/downloader.rs

@@ -580,7 +580,7 @@ impl Downloader<LocalFile> {
         overwrite: Option<bool>,
     ) -> Result<DownloaderResult<Self>, 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::<Duration>::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::<Duration>::default();
     retry_interval += Duration::new(0, 0);
     client

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<Request, PixivDownloaderError> {
+        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
     }

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 {

src/opts.rs

@@ -138,6 +138,12 @@ pub struct CommandOpts {
     #[cfg(feature = "ugoira")]
     /// Whether to use ugoira cli.
     pub ugoira_cli: Option<bool>,
+    /// Set a timeout in milliseconds for only the connect phase of a client.
+    pub connect_timeout: Option<u64>,
+    /// 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<u64>,
 }
 
 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<T: AsRef<str>>(s: Option<T>) -> Result<Option<u64>, ParseIntErr
     }
 }
 
+/// Prase Non Zero [u64] from string
+pub fn parse_non_zero_u64<T: AsRef<str>>(s: Option<T>) -> Result<Option<u64>, ParseIntError> {
+    match s {
+        Some(s) => {
+            let s = s.as_ref();
+            let s = s.trim();
+            let c = s.parse::<std::num::NonZeroU64>()?;
+            Ok(Some(c.get()))
+        }
+        None => Ok(None),
+    }
+}
+
 pub fn parse_nonempty_usize<T: AsRef<str>>(s: Option<T>) -> Result<Option<usize>, ParseIntError> {
     match s {
         Some(s) => {
@@ -709,6 +730,13 @@ pub fn parse_cmd() -> Option<CommandOpts> {
         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<CommandOpts> {
             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 <opt>:")
+                    .replace("<opt>", "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 <opt>:")
+                    .replace("<opt>", "client-timeout")
+                    .as_str(),
+                e
+            );
+            return None;
+        }
+    }
     re
 }
 

src/settings.rs

@@ -312,6 +312,13 @@ impl SettingStore {
         }
     }
 
+    pub fn get_u64(&self, key: &str) -> Option<u64> {
+        match self.data.get(key) {
+            Some(obj) => obj.as_u64(),
+            None => None,
+        }
+    }
+
     pub fn have(&self, key: &str) -> bool {
         self.data.have(key)
     }

src/settings_list.rs

@@ -74,6 +74,8 @@ pub fn get_settings_list() -> Vec<SettingDes> {
         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()

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<Option<NonTailList<Duration>>>,
     /// Request middlewares
     req_middlewares: RwLock<Vec<Box<dyn ReqMiddleware + Send + Sync>>>,
+    /// Set request timeout. The timeout is applied from when the request starts connecting until
+    /// the response body has finished.
+    timeout: RwLock<Option<Duration>>,
 }
 
 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<Duration>) {
+        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
     }
 }