remove no_httpauth flag from cli and some cleanup (#23)

Signed-off-by: KaranGauswami <[email protected]>

README.md

@@ -40,7 +40,6 @@ Options:
   -u, --username <USERNAME>                Socks5 username
   -P, --password <PASSWORD>                Socks5 password
       --http-basic <USER:PASSWD>           HTTP Basic Auth
-      --no-httpauth <1/0>                  Ignore HTTP Basic Auth, [default: 1]
   -s, --socks-address <SOCKS_ADDRESS>      Socks5 proxy address [default: 127.0.0.1:1080]
       --allowed-domains <ALLOWED_DOMAINS>  Comma-separated list of allowed domains
   -h, --help                               Print help information

src/main.rs

@@ -1,7 +1,7 @@
 mod auth;
 
 use crate::auth::Auth;
-use clap::{Args, Parser, value_parser};
+use clap::{Args, Parser};
 use color_eyre::eyre::Result;
 
 use tokio_socks::tcp::Socks5Stream;
@@ -10,16 +10,16 @@ use tracing_subscriber::EnvFilter;
 
 use std::net::{Ipv4Addr, SocketAddr};
 
+use base64::engine::general_purpose;
+use base64::Engine;
 use bytes::Bytes;
 use http_body_util::{combinators::BoxBody, BodyExt, Empty, Full};
 use hyper::client::conn::http1::Builder;
+use hyper::header::{HeaderValue, PROXY_AUTHENTICATE};
 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;
@@ -60,10 +60,6 @@ struct Cli {
     /// HTTP Basic Auth credentials in the format "user:passwd"
     #[arg(long)]
     http_basic: Option<String>,
-
-    /// Disable HTTP authentication [default: 1]
-    #[arg(long, value_parser = value_parser ! (u8).range(0..=1), default_value_t = 1)]
-    no_httpauth: u8,
 }
 
 #[tokio::main]
@@ -83,18 +79,30 @@ 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 = args
+        .http_basic
+        .map(|hb| format!("Basic {}", general_purpose::STANDARD.encode(hb)))
+        .map(|hb| HeaderValue::from_str(&hb))
+        .transpose()?;
     let http_basic = &*Box::leak(Box::new(http_basic));
-    let no_httpauth = args.no_httpauth == 1;
 
     let listener = TcpListener::bind(addr).await?;
     info!("Listening on http://{}", addr);
 
     loop {
-        let (stream, _) = listener.accept().await?;
+        let (stream, client_addr) = listener.accept().await?;
         let io = TokioIo::new(stream);
 
-        let serve_connection = service_fn(move |req| proxy(req, socks_addr, auth, &http_basic, allowed_domains, no_httpauth));
+        let serve_connection = service_fn(move |req| {
+            proxy(
+                req,
+                client_addr,
+                socks_addr,
+                auth,
+                http_basic,
+                allowed_domains,
+            )
+        });
 
         tokio::task::spawn(async move {
             if let Err(err) = http1::Builder::new()
@@ -112,44 +120,39 @@ async fn main() -> Result<()> {
 
 async fn proxy(
     req: Request<hyper::body::Incoming>,
+    client_addr: SocketAddr,
     socks_addr: SocketAddr,
     auth: &'static Option<Auth>,
-    http_basic: &Option<String>,
+    basic_http_header: &Option<HeaderValue>,
     allowed_domains: &Option<Vec<String>>,
-    no_httpauth: bool,
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, hyper::Error> {
-    let mut http_authed = false;
+    let mut authenticated = false;
     let hm = req.headers();
 
-    if no_httpauth {
+    if let Some(basic_http_header) = basic_http_header {
-        http_authed = true;
+        let Some(http_auth) = hm.get("proxy-authorization") else {
-    } else if hm.contains_key("proxy-authorization") {
+            // When the request does not contain a Proxy-Authorization header,
-        let config_auth = match http_basic {
+            // send a 407 response code and a Proxy-Authenticate header
-            Some(value) => value.clone(),
+            let mut response = Response::new(full("Proxy Authentication Required"));
-            None => String::new(),
+            *response.status_mut() = http::StatusCode::PROXY_AUTHENTICATION_REQUIRED;
+            response.headers_mut().insert(
+                PROXY_AUTHENTICATE,
+                HeaderValue::from_static("Basic realm=\"proxy\""),
+            );
+            return Ok(response);
         };
-        let http_auth = hm.get("proxy-authorization").unwrap();
+        if http_auth == basic_http_header {
-        if http_auth == &HeaderValue::from_str(&config_auth).unwrap() {
+            authenticated = true;
-            http_authed = true;
         }
     } else {
-        // When the request does not contain a Proxy-Authorization header,
+        authenticated = true;
-        // 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 {
+    if !authenticated {
-        warn!("Failed to authenticate: {:?}", hm);
+        warn!("Failed auth attempt from: {}", client_addr);
-        let mut resp = Response::new(full(
+        // http response code reference taken from tinyproxy
-            "Authorization failed, you are not allowed through the proxy.",
+        let mut resp = Response::new(full("Unauthorized"));
-        ));
+        *resp.status_mut() = http::StatusCode::UNAUTHORIZED;
-        *resp.status_mut() = http::StatusCode::FORBIDDEN;
         return Ok(resp);
     }
 
@@ -194,7 +197,7 @@ async fn proxy(
     } else {
         let host = req.uri().host().expect("uri has no host");
         let port = req.uri().port_u16().unwrap_or(80);
-        let addr = format!("{}:{}", host, port);
+        let addr = (host, port);
 
         let stream = match auth {
             Some(auth) => Socks5Stream::connect_with_password(