From b8a24068684719855dfe38d2b865d6204d4b6215 Mon Sep 17 00:00:00 2001
From: GoogleGeek <ffmpeg@gmail.com>
Date: Mon, 10 Jun 2024 16:50:50 +0800
Subject: [PATCH] support HTTP Basic Auth

---
 Cargo.lock  |  7 +++++++
 Cargo.toml  |  1 +
 src/main.rs | 45 ++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 7d3f7f0..b8c3254 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -89,6 +89,12 @@ dependencies = [
  "rustc-demangle",
 ]
 
+[[package]]
+name = "base64"
+version = "0.22.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
+
 [[package]]
 name = "bytes"
 version = "1.5.0"
@@ -567,6 +573,7 @@ dependencies = [
 name = "sthp"
 version = "0.4.0"
 dependencies = [
+ "base64",
  "bytes",
  "clap",
  "color-eyre",
diff --git a/Cargo.toml b/Cargo.toml
index 276e99d..8396cae 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -22,3 +22,4 @@ http-body-util = "0.1.0-rc.2"
 tracing = "0.1.37"
 tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
 hyper-util = { git = "https://github.com/hyperium/hyper-util.git", rev = "229757e565e0935a7a3b1d0f9e9ab88d9310e779" }
+base64 = "0.22.1"
diff --git a/src/main.rs b/src/main.rs
index 576111c..fb4167e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -17,6 +17,9 @@ use hyper::server::conn::http1;
 use hyper::service::service_fn;
 use hyper::upgrade::Upgraded;
 use hyper::{Method, Request, Response};
+use hyper::header::{HeaderValue, PROXY_AUTHENTICATE};
+use base64::engine::general_purpose;
+use base64::Engine;
 
 use hyper_util::rt::TokioIo;
 use tokio::net::TcpListener;
@@ -53,6 +56,10 @@ struct Cli {
     /// Comma-separated list of allowed domains
     #[arg(long, value_delimiter = ',')]
     allowed_domains: Option<Vec<String>>,
+
+    /// HTTP Basic Auth credentials in the format "user:passwd"
+    #[arg(long, default_value = "")]
+    http_basic: Option<String>,
 }
 
 #[tokio::main]
@@ -72,6 +79,8 @@ async fn main() -> Result<()> {
     let addr = SocketAddr::from((args.listen_ip, port));
     let allowed_domains = args.allowed_domains;
     let allowed_domains = &*Box::leak(Box::new(allowed_domains));
+    let http_basic = args.http_basic.map(|hb| format!("Basic {}", general_purpose::STANDARD.encode(hb)));
+    let http_basic = &*Box::leak(Box::new(http_basic));
 
     let listener = TcpListener::bind(addr).await?;
     info!("Listening on http://{}", addr);
@@ -80,7 +89,7 @@ async fn main() -> Result<()> {
         let (stream, _) = listener.accept().await?;
         let io = TokioIo::new(stream);
 
-        let serve_connection = service_fn(move |req| proxy(req, socks_addr, auth, allowed_domains));
+        let serve_connection = service_fn(move |req| proxy(req, socks_addr, auth, &http_basic, allowed_domains));
 
         tokio::task::spawn(async move {
             if let Err(err) = http1::Builder::new()
@@ -100,9 +109,43 @@ async fn proxy(
     req: Request<hyper::body::Incoming>,
     socks_addr: SocketAddr,
     auth: &'static Option<Auth>,
+    http_basic: &Option<String>,
     allowed_domains: &Option<Vec<String>>,
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, hyper::Error> {
     let uri = req.uri();
+    let mut http_authed = false;
+    let hm = req.headers();
+
+    if hm.contains_key("proxy-authorization") {
+        let config_auth = match http_basic {
+            Some(value) => value.clone(),
+            None => String::new(),
+        };
+        let http_auth = hm.get("proxy-authorization").unwrap();
+        if http_auth == &HeaderValue::from_str(&config_auth).unwrap() {
+            http_authed = true;
+        }
+    } else {
+        // When the request does not contain a Proxy-Authorization header,
+        // send a 407 response code and a Proxy-Authenticate header
+        let mut response = Response::new(full("Proxy authentication required"));
+        *response.status_mut() = http::StatusCode::PROXY_AUTHENTICATION_REQUIRED;
+        response.headers_mut().insert(
+            PROXY_AUTHENTICATE,
+            HeaderValue::from_static("Basic realm=\"proxy\""),
+        );
+        return Ok(response);
+    }
+
+    if !http_authed {
+        warn!("Failed to authenticate: {:?}", hm);
+        let mut resp = Response::new(full(
+            "Authorization failed, you are not allowed through the proxy.",
+        ));
+        *resp.status_mut() = http::StatusCode::FORBIDDEN;
+        return Ok(resp);
+    }
+
     let method = req.method();
     debug!("Proxying request: {} {}", method, uri);
     if let (Some(allowed_domains), Some(request_domain)) = (allowed_domains, req.uri().host()) {