From d20ad40145a9e355d8add45058da5ce1cbe3d874 Mon Sep 17 00:00:00 2001
From: link2xt <link2xt@testrun.org>
Date: Wed, 28 Aug 2024 02:58:21 +0000
Subject: [PATCH] add socks support

---
 src/net/http.rs | 29 +++++++++++++++++++++++------
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/src/net/http.rs b/src/net/http.rs
index e51003e830..109b65f89b 100644
--- a/src/net/http.rs
+++ b/src/net/http.rs
@@ -10,6 +10,7 @@ use serde::Serialize;
 use crate::context::Context;
 use crate::net::session::SessionStream;
 use crate::net::tls::wrap_tls;
+use crate::socks::Socks5Config;
 
 /// HTTP(S) GET response.
 #[derive(Debug)]
@@ -31,7 +32,6 @@ pub async fn read_url(context: &Context, url: &str) -> Result<String> {
     Ok(text.to_string())
 }
 
-// TODO add support for SOCKS5
 async fn get_http_sender<B>(
     context: &Context,
     parsed_url: hyper::Uri,
@@ -43,6 +43,7 @@ where
 {
     let scheme = parsed_url.scheme_str().context("URL has no scheme")?;
     let host = parsed_url.host().context("URL has no host")?;
+    let socks5_config_opt = Socks5Config::from_database(&context.sql).await?;
 
     let stream: Box<dyn SessionStream> = match scheme {
         "http" => {
@@ -53,16 +54,32 @@ where
             // better resolve from scratch each time to prevent
             // cache poisoning attacks from having lasting effects.
             let load_cache = false;
-            let tcp_stream = crate::net::connect_tcp(context, host, port, load_cache).await?;
-            Box::new(tcp_stream)
+            if let Some(socks5_config) = socks5_config_opt {
+                let socks5_stream = socks5_config
+                    .connect(context, host, port, load_cache)
+                    .await?;
+                Box::new(socks5_stream)
+            } else {
+                let tcp_stream = crate::net::connect_tcp(context, host, port, load_cache).await?;
+                Box::new(tcp_stream)
+            }
         }
         "https" => {
             let port = parsed_url.port_u16().unwrap_or(443);
             let load_cache = true;
-            let tcp_stream = crate::net::connect_tcp(context, host, port, load_cache).await?;
             let strict_tls = true;
-            let tls_stream = wrap_tls(strict_tls, host, &[], tcp_stream).await?;
-            Box::new(tls_stream)
+
+            if let Some(socks5_config) = socks5_config_opt {
+                let socks5_stream = socks5_config
+                    .connect(context, host, port, load_cache)
+                    .await?;
+                let tls_stream = wrap_tls(strict_tls, host, &[], socks5_stream).await?;
+                Box::new(tls_stream)
+            } else {
+                let tcp_stream = crate::net::connect_tcp(context, host, port, load_cache).await?;
+                let tls_stream = wrap_tls(strict_tls, host, &[], tcp_stream).await?;
+                Box::new(tls_stream)
+            }
         }
         _ => bail!("Unknown URL scheme"),
     };