Add proxy option

james58899 · Oct 27, 2023 · 1a814c2 · 1a814c2
1 parent 935c242
commit 1a814c2
Show file tree

Hide file tree

Showing 8 changed files with 64 additions and 20 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -26,7 +26,7 @@ parking_lot = { version = "0.12", features = ["hardware-lock-elision", "deadlock
 pin-project-lite = "0.2"
 rand = { version = "0.8", default-features = false, features = ["small_rng"] }
 regex = "1.10"
-reqwest = { version = "0.11", default-features = false, features = ["rustls-tls", "stream"] }
+reqwest = { version = "0.11", default-features = false, features = ["rustls-tls", "stream", "socks"] }
 tempfile = "3.8"
 tokio = { version = "1", features = ["full", "parking_lot"] }
 tokio-stream = { version = "0.1", default-features = false, features = ["fs"] }

diff --git a/src/gallery_downloader.rs b/src/gallery_downloader.rs
@@ -35,7 +35,7 @@ impl GalleryDownloader {
     pub fn new<P: AsRef<Path>>(client: Arc<RPCClient>, download_dir: P) -> GalleryDownloader {
         GalleryDownloader {
             client,
-            reqwest: util::create_http_client(Duration::from_secs(300)),
+            reqwest: util::create_http_client(Duration::from_secs(300), None),
             download_dir: download_dir.as_ref().to_path_buf(),
         }
     }

diff --git a/src/main.rs b/src/main.rs
@@ -33,6 +33,7 @@ use openssl::{
 };
 use parking_lot::{Mutex, RwLock};
 use regex::Regex;
+use reqwest::Proxy;
 use tempfile::TempPath;
 #[cfg(not(target_env = "msvc"))]
 use tikv_jemallocator::Jemalloc;
@@ -130,6 +131,10 @@ struct Args {
     /// Disable server command ip check
     #[arg(long, default_value_t = false)]
     disable_ip_origin_check: bool,
+
+    /// Configure proxy for fetch cache
+    #[arg(long)]
+    proxy: Option<String>,
 }
 
 type DownloadState = RwLock<HashMap<[u8; 20], (Arc<TempPath>, watch::Receiver<u64>)>>;
@@ -141,6 +146,7 @@ struct AppState {
     download_state: DownloadState,
     cache_manager: Arc<CacheManager>,
     command_channel: Sender<Command>,
+    has_proxy: bool,
 }
 
 pub enum Command {
@@ -192,18 +198,32 @@ async fn main() -> Result<(), Box<dyn Error>> {
     )
     .await?;
 
-    // command channel
+    // Proxy
+    let proxy = match args.proxy.as_ref().map(Proxy::all) {
+        Some(Ok(proxy)) => {
+            info!("Using proxy for fetch cache: {}", args.proxy.unwrap());
+            Some(proxy)
+        }
+        Some(Err(err)) => {
+            error!("Parser proxy setting error: {}", err);
+            None
+        }
+        None => None,
+    };
+    let has_proxy = proxy.is_some();
+    // Command channel
     let (tx, mut rx) = mpsc::channel::<Command>(1);
     let (server, cert_changer) = create_server(
         args.port.unwrap_or_else(|| init_settings.client_port()),
         client.get_cert().await.unwrap(),
         AppState {
             runtime: Handle::current(),
-            reqwest: create_http_client(Duration::from_secs(30)),
+            reqwest: create_http_client(Duration::from_secs(30), proxy),
             rpc: client.clone(),
             download_state: Default::default(),
             cache_manager: cache_manager.clone(),
             command_channel: tx.clone(),
+            has_proxy,
         },
     );
     let server_handle = server.handle();

diff --git a/src/route/cache.rs b/src/route/cache.rs
@@ -1,4 +1,4 @@
-use std::{io::SeekFrom, ops::RangeInclusive, sync::Arc};
+use std::{io::SeekFrom, ops::RangeInclusive, sync::Arc, time::Duration};
 
 use actix_files::NamedFile;
 use actix_web::{
@@ -25,7 +25,7 @@ use tokio::{
 use crate::{
     cache_manager::CacheFileInfo,
     route::{forbidden, parse_additional},
-    util::string_to_hash,
+    util::{create_http_client, string_to_hash},
     AppState,
 };
 
@@ -120,8 +120,17 @@ async fn hath(
                 };
 
                 let mut download = 0;
-                let request = data.reqwest.get(sources.next().unwrap());
-                if let Ok(mut stream) = request.send().await.and_then(|r| r.error_for_status()).map(|r| r.bytes_stream()) {
+                let source = sources.next().unwrap();
+                let mut request = data.reqwest.get(source).send().await;
+                if let Err(ref err) = request {
+                    error!("Cache download error: {}", err);
+
+                    // Retry without proxy
+                    if data.has_proxy && err.is_connect() {
+                        request = create_http_client(Duration::from_secs(30), None).get(source).send().await;
+                    }
+                };
+                if let Ok(mut stream) = request.and_then(|r| r.error_for_status()).map(|r| r.bytes_stream()) {
                     while let Some(bytes) = stream.next().await {
                         let bytes = match &bytes {
                             Ok(it) => it,

diff --git a/src/route/server_command.rs b/src/route/server_command.rs
@@ -15,7 +15,7 @@ use reqwest::{
 
 use crate::{
     route::{forbidden, parse_additional, speed_test::random_response},
-    util::string_to_hash,
+    util::{create_http_client, string_to_hash},
     AppState, Command, MAX_KEY_TIME_DRIFT,
 };
 
@@ -88,16 +88,13 @@ async fn servercmd(
                 )
                 .unwrap();
                 debug!("Speedtest thread start: {}", url);
-                let reqwest = data.reqwest.clone();
+                let reqwest = create_http_client(Duration::from_secs(60), None); // No proxy http client
                 requests.push(tokio::spawn(async move {
                     for retry in 0..3 {
                         if retry > 0 {
                             debug!("Retrying.. ({} tries left)", 3 - retry);
                         }
-                        let request = reqwest
-                            .get(url.clone())
-                            .header(CONNECTION, HeaderValue::from_static("Close"))
-                            .timeout(Duration::from_secs(60));
+                        let request = reqwest.get(url.clone()).header(CONNECTION, HeaderValue::from_static("Close"));
                         match request.send().await.and_then(|r| r.error_for_status()) {
                             Ok(res) => {
                                 let start = Instant::now();

diff --git a/src/rpc.rs b/src/rpc.rs
@@ -126,7 +126,7 @@ impl RPCClient {
             clock_offset: AtomicI64::new(0),
             id,
             key: key.to_string(),
-            reqwest: create_http_client(Duration::from_secs(600)),
+            reqwest: create_http_client(Duration::from_secs(600), None),
             rpc_servers: RwLock::new(vec![]),
             running: AtomicBool::new(false),
             settings: Arc::new(Settings {

diff --git a/src/util.rs b/src/util.rs
@@ -2,6 +2,7 @@ use std::time::Duration;
 
 use futures::future::try_join_all;
 use openssl::sha::Sha1;
+use reqwest::Proxy;
 use tokio::fs::create_dir_all;
 
 use crate::CLIENT_VERSION;
@@ -12,18 +13,22 @@ pub fn string_to_hash(str: String) -> String {
     hex::encode(hasher.finish())
 }
 
-pub fn create_http_client(timeout: Duration) -> reqwest::Client {
-    reqwest::ClientBuilder::new()
+pub fn create_http_client(timeout: Duration, proxy: Option<Proxy>) -> reqwest::Client {
+    let mut builder = reqwest::ClientBuilder::new()
         .user_agent(format!("Hentai@Home {CLIENT_VERSION}"))
         .tcp_keepalive(Duration::from_secs(75)) // Linux default keepalive inverval
         .connect_timeout(Duration::from_secs(5))
         .timeout(timeout)
         .pool_idle_timeout(Duration::from_secs(3600))
         .pool_max_idle_per_host(8)
         .http1_title_case_headers()
-        .http1_only()
-        .build()
-        .unwrap()
+        .http1_only();
+
+    if let Some(proxy) = proxy {
+        builder = builder.proxy(proxy);
+    }
+
+    builder.build().unwrap()
 }
 
 pub async fn create_dirs(dirs: Vec<&str>) -> Result<Vec<()>, std::io::Error> {