Feature/structure and testing changes (#22)

* refactor: refactored code to facilitate testing in pipeline

* refactor: added some tests

* refactor: added some tests

* refactor: added some tests

* version bump

* made changes so arguments will accept Option<&T> instead of &Option<T>

* added auth related code from main branch

* allow support for passing DNS in the socks_address

* accepted clippy suggestions

* bump the dependencies version

.github/workflows/rust.yml

@@ -2,21 +2,18 @@ name: Rust
 
 on:
   push:
-    branches: [ main ]
+    branches: [main]
   pull_request:
-    branches: [ main ]
+    branches: [main]
 
 env:
   CARGO_TERM_COLOR: always
 
 jobs:
   build:
-
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v2
+      - uses: actions/checkout@v2
-    - name: Build
+      - name: Run tests
-      run: cargo build --verbose
+        run: cargo test --verbose -- --nocapture
-    - name: Run tests
-      run: cargo test --verbose

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "sthp"
-version = "0.4.0"
+version = "0.5.0-alpha1"
 license = "MIT"
 authors = ["Karan Gauswami <[email protected]>"]
 edition = "2021"
@@ -12,14 +12,17 @@ description = "Convert Socks5 proxy into Http proxy"
 
 [dependencies]
 color-eyre = { version = "0.6", default-features = false }
-http = "0.2.9"
+hyper = { version = "1.3", features = ["client","server","http1"] }
-hyper = { version = "1.0.0-rc.4", features = ["client","server","http1"] }
 clap = { version = "4", features = ["derive"] }
 tokio-socks = "0.5"
-tokio = { version = "1.28", features = ["macros", "rt-multi-thread"] }
+tokio = { version = "1.38", features = ["macros", "rt-multi-thread"] }
 bytes = "1.4.0"
-http-body-util = "0.1.0-rc.2"
+http-body-util = "0.1.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" }
+hyper-util = { version="0.1.5",features = ["tokio"] }
 base64 = "0.22.1"
+
+[dev-dependencies]
+socksprox = { version = "0.1" }
+reqwest = { version = "0.12" }

README.md

@@ -27,6 +27,15 @@ sthp -p 8080 -s 127.0.0.1:1080
 
 This will create proxy server on 8080 and use localhost:1080 as a Socks5 Proxy
 
+```bash
+sthp -p 8080 -s example.com:8080
+```
+
+This will create proxy server on 8080 and use example:1080 as a Socks5 Proxy
+
+> [!NOTE]  
+> The --socks-address (-s) flag does not support adding a schema at the start (e.g., socks:// or socks5h://). Currently, it only supports socks5h, which means DNS resolution will be done on the SOCKS server.
+
 ### Options
 
 There are a few options for using `sthp`.

src/lib.rs

@@ -0,0 +1,2 @@
+pub mod proxy;
+pub use proxy::proxy_request;

src/main.rs

@@ -1,32 +1,21 @@
-mod auth;
-
-use crate::auth::Auth;
 use clap::{Args, Parser};
-use color_eyre::eyre::Result;
+use color_eyre::eyre::{OptionExt, Result};
 
-use tokio_socks::tcp::Socks5Stream;
+use sthp::proxy::auth::Auth;
-use tracing::{debug, info, warn};
+use sthp::proxy_request;
+use tracing::{error, info};
 use tracing_subscriber::EnvFilter;
 
-use std::net::{Ipv4Addr, SocketAddr};
+use std::net::{Ipv4Addr, SocketAddr, ToSocketAddrs};
 
 use base64::engine::general_purpose;
 use base64::Engine;
-use bytes::Bytes;
+use hyper::header::HeaderValue;
-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_util::rt::TokioIo;
 use tokio::net::TcpListener;
 
 #[derive(Debug, Args)]
 #[group()]
-struct Auths {
+struct AuthParams {
     /// Socks5 username
     #[arg(short = 'u', long, required = false)]
     username: String,
@@ -36,6 +25,12 @@ struct Auths {
     password: String,
 }
 
+fn socket_addr(s: &str) -> Result<SocketAddr> {
+    let mut address = s.to_socket_addrs()?;
+    let address = address.next();
+    address.ok_or_eyre("no IP address found for the hostname".to_string())
+}
+
 #[derive(Parser, Debug)]
 #[command(author, version, about,long_about=None)]
 struct Cli {
@@ -47,10 +42,10 @@ struct Cli {
     listen_ip: Ipv4Addr,
 
     #[command(flatten)]
-    auth: Option<Auths>,
+    auth: Option<AuthParams>,
 
     /// Socks5 proxy address
-    #[arg(short, long, default_value = "127.0.0.1:1080")]
+    #[arg(short, long, default_value = "127.0.0.1:1080", value_parser=socket_addr)]
     socks_address: SocketAddr,
 
     /// Comma-separated list of allowed domains
@@ -72,10 +67,10 @@ async fn main() -> Result<()> {
 
     let socks_addr = args.socks_address;
     let port = args.port;
-    let auth = args
+    let auth_details = args
         .auth
         .map(|auth| Auth::new(auth.username, auth.password));
-    let auth = &*Box::leak(Box::new(auth));
+    let auth_details = &*Box::leak(Box::new(auth_details));
     let addr = SocketAddr::from((args.listen_ip, port));
     let allowed_domains = args.allowed_domains;
     let allowed_domains = &*Box::leak(Box::new(allowed_domains));
@@ -91,183 +86,19 @@ async fn main() -> Result<()> {
 
     loop {
         let (stream, client_addr) = listener.accept().await?;
-        let io = TokioIo::new(stream);
-
-        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()
+            if let Err(e) = proxy_request(
-                .preserve_header_case(true)
+                stream,
-                .title_case_headers(true)
+                client_addr,
-                .serve_connection(io, serve_connection)
-                .with_upgrades()
-                .await
-            {
-                warn!("Failed to serve connection: {:?}", err);
-            }
-        });
-    }
-}
-
-async fn proxy(
-    req: Request<hyper::body::Incoming>,
-    client_addr: SocketAddr,
-    socks_addr: SocketAddr,
-    auth: &'static Option<Auth>,
-    basic_http_header: &Option<HeaderValue>,
-    allowed_domains: &Option<Vec<String>>,
-) -> Result<Response<BoxBody<Bytes, hyper::Error>>, hyper::Error> {
-    let mut authenticated = false;
-    let hm = req.headers();
-
-    if let Some(basic_http_header) = basic_http_header {
-        let Some(http_auth) = hm.get("proxy-authorization") 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_auth == basic_http_header {
-            authenticated = true;
-        }
-    } else {
-        authenticated = true;
-    }
-
-    if !authenticated {
-        warn!("Failed auth attempt from: {}", client_addr);
-        // http response code reference taken from tinyproxy
-        let mut resp = Response::new(full("Unauthorized"));
-        *resp.status_mut() = http::StatusCode::UNAUTHORIZED;
-        return Ok(resp);
-    }
-
-    let method = req.method();
-    debug!("Proxying request: {} {}", method, req.uri());
-    if let (Some(allowed_domains), Some(request_domain)) = (allowed_domains, req.uri().host()) {
-        let domain = request_domain.to_owned();
-        if !allowed_domains.contains(&domain) {
-            warn!(
-                "Access to domain {} is not allowed through the proxy.",
-                domain
-            );
-            let mut resp = Response::new(full(
-                "Access to this domain is not allowed through the proxy.",
-            ));
-            *resp.status_mut() = http::StatusCode::FORBIDDEN;
-            return Ok(resp);
-        }
-    }
-
-    if Method::CONNECT == req.method() {
-        if let Some(addr) = host_addr(req.uri()) {
-            tokio::task::spawn(async move {
-                match hyper::upgrade::on(req).await {
-                    Ok(upgraded) => {
-                        if let Err(e) = tunnel(upgraded, addr, socks_addr, auth).await {
-                            warn!("server io error: {}", e);
-                        };
-                    }
-                    Err(e) => warn!("upgrade error: {}", e),
-                }
-            });
-
-            Ok(Response::new(empty()))
-        } else {
-            warn!("CONNECT host is not socket addr: {:?}", req.uri());
-            let mut resp = Response::new(full("CONNECT must be to a socket address"));
-            *resp.status_mut() = http::StatusCode::BAD_REQUEST;
-
-            Ok(resp)
-        }
-    } else {
-        let host = req.uri().host().expect("uri has no host");
-        let port = req.uri().port_u16().unwrap_or(80);
-        let addr = (host, port);
-
-        let stream = match auth {
-            Some(auth) => Socks5Stream::connect_with_password(
                 socks_addr,
-                addr,
+                auth_details.as_ref(),
-                &auth.username,
+                allowed_domains.as_ref(),
-                &auth.password,
+                http_basic.as_ref(),
             )
             .await
-            .unwrap(),
+            {
-            None => Socks5Stream::connect(socks_addr, addr).await.unwrap(),
+                error!("Error proxying request: {}", e);
-        };
-        let io = TokioIo::new(stream);
-
-        let (mut sender, conn) = Builder::new()
-            .preserve_header_case(true)
-            .title_case_headers(true)
-            .handshake(io)
-            .await?;
-        tokio::task::spawn(async move {
-            if let Err(err) = conn.await {
-                warn!("Connection failed: {:?}", err);
             }
         });
-
-        let resp = sender.send_request(req).await?;
-        Ok(resp.map(|b| b.boxed()))
     }
 }
-
-fn host_addr(uri: &http::Uri) -> Option<String> {
-    uri.authority().map(|auth| auth.to_string())
-}
-
-fn empty() -> BoxBody<Bytes, hyper::Error> {
-    Empty::<Bytes>::new()
-        .map_err(|never| match never {})
-        .boxed()
-}
-
-fn full<T: Into<Bytes>>(chunk: T) -> BoxBody<Bytes, hyper::Error> {
-    Full::new(chunk.into())
-        .map_err(|never| match never {})
-        .boxed()
-}
-
-async fn tunnel(
-    upgraded: Upgraded,
-    addr: String,
-    socks_addr: SocketAddr,
-    auth: &Option<Auth>,
-) -> Result<()> {
-    let mut stream = match auth {
-        Some(auth) => {
-            Socks5Stream::connect_with_password(socks_addr, addr, &auth.username, &auth.password)
-                .await?
-        }
-        None => Socks5Stream::connect(socks_addr, addr).await?,
-    };
-
-    let mut upgraded = TokioIo::new(upgraded);
-
-    // Proxying data
-    let (from_client, from_server) =
-        tokio::io::copy_bidirectional(&mut upgraded, &mut stream).await?;
-
-    // Print message when done
-    debug!(
-        "client wrote {} bytes and received {} bytes",
-        from_client, from_server
-    );
-    Ok(())
-}

src/proxy/mod.rs

@@ -0,0 +1,213 @@
+use auth::Auth;
+use color_eyre::eyre::Result;
+
+pub mod auth;
+
+use hyper::header::{HeaderValue, PROXY_AUTHENTICATE};
+use hyper::service::service_fn;
+use tokio::net::TcpStream;
+use tokio_socks::tcp::Socks5Stream;
+use tracing::{debug, warn};
+
+use std::net::SocketAddr;
+
+use bytes::Bytes;
+use http_body_util::{combinators::BoxBody, BodyExt, Empty, Full};
+
+use hyper::upgrade::Upgraded;
+use hyper::{Method, Request, Response};
+
+use hyper_util::rt::TokioIo;
+
+use hyper::client::conn::http1::Builder;
+use hyper::server::conn::http1;
+
+async fn proxy(
+    req: Request<hyper::body::Incoming>,
+    client_addr: SocketAddr,
+    socks_addr: SocketAddr,
+    auth: Option<&'static Auth>,
+    allowed_domains: Option<&'static Vec<String>>,
+    basic_http_header: Option<&HeaderValue>,
+) -> Result<Response<BoxBody<Bytes, hyper::Error>>> {
+    let mut authenticated = false;
+    let hm = req.headers();
+
+    if let Some(basic_http_header) = basic_http_header {
+        let Some(http_auth) = hm.get("proxy-authorization") 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() = hyper::StatusCode::PROXY_AUTHENTICATION_REQUIRED;
+            response.headers_mut().insert(
+                PROXY_AUTHENTICATE,
+                HeaderValue::from_static("Basic realm=\"proxy\""),
+            );
+            return Ok(response);
+        };
+        if http_auth == basic_http_header {
+            authenticated = true;
+        }
+    } else {
+        authenticated = true;
+    }
+
+    if !authenticated {
+        warn!("Failed auth attempt from: {}", client_addr);
+        // http response code reference taken from tinyproxy
+        let mut resp = Response::new(full("Unauthorized"));
+        *resp.status_mut() = hyper::StatusCode::UNAUTHORIZED;
+        return Ok(resp);
+    }
+
+    let uri = req.uri();
+    let method = req.method();
+    debug!("Proxying request: {} {}", method, uri);
+    if let (Some(allowed_domains), Some(request_domain)) = (allowed_domains, req.uri().host()) {
+        let domain = request_domain.to_owned();
+        if !allowed_domains.contains(&domain) {
+            warn!(
+                "Access to domain {} is not allowed through the proxy.",
+                domain
+            );
+            let mut resp = Response::new(full(
+                "Access to this domain is not allowed through the proxy.",
+            ));
+            *resp.status_mut() = hyper::StatusCode::FORBIDDEN;
+            return Ok(resp);
+        }
+    }
+
+    if Method::CONNECT == req.method() {
+        if let Some(addr) = host_addr(req.uri()) {
+            tokio::task::spawn(async move {
+                match hyper::upgrade::on(req).await {
+                    Ok(upgraded) => {
+                        if let Err(e) = tunnel(upgraded, addr, socks_addr, auth).await {
+                            warn!("server io error: {}", e);
+                        };
+                    }
+                    Err(e) => warn!("upgrade error: {}", e),
+                }
+            });
+
+            Ok(Response::new(empty()))
+        } else {
+            warn!("CONNECT host is not socket addr: {:?}", req.uri());
+            let mut resp = Response::new(full("CONNECT must be to a socket address"));
+            *resp.status_mut() = hyper::StatusCode::BAD_REQUEST;
+
+            Ok(resp)
+        }
+    } 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 stream = match auth {
+            Some(auth) => Socks5Stream::connect_with_password(
+                socks_addr,
+                addr,
+                &auth.username,
+                &auth.password,
+            )
+            .await
+            .unwrap(),
+            None => Socks5Stream::connect(socks_addr, addr).await?,
+        };
+        let io = TokioIo::new(stream);
+
+        let (mut sender, conn) = Builder::new()
+            .preserve_header_case(true)
+            .title_case_headers(true)
+            .handshake(io)
+            .await?;
+        tokio::task::spawn(async move {
+            if let Err(err) = conn.await {
+                warn!("Connection failed: {:?}", err);
+            }
+        });
+
+        let resp = sender.send_request(req).await?;
+        Ok(resp.map(|b| b.boxed()))
+    }
+}
+
+fn host_addr(uri: &hyper::Uri) -> Option<String> {
+    uri.authority().map(|auth| auth.to_string())
+}
+
+fn empty() -> BoxBody<Bytes, hyper::Error> {
+    Empty::<Bytes>::new()
+        .map_err(|never| match never {})
+        .boxed()
+}
+
+fn full<T: Into<Bytes>>(chunk: T) -> BoxBody<Bytes, hyper::Error> {
+    Full::new(chunk.into())
+        .map_err(|never| match never {})
+        .boxed()
+}
+
+async fn tunnel(
+    upgraded: Upgraded,
+    addr: String,
+    socks_addr: SocketAddr,
+    auth: Option<&Auth>,
+) -> Result<()> {
+    let mut stream = match auth {
+        Some(auth) => {
+            Socks5Stream::connect_with_password(socks_addr, addr, &auth.username, &auth.password)
+                .await?
+        }
+        None => Socks5Stream::connect(socks_addr, addr).await?,
+    };
+
+    let mut upgraded = TokioIo::new(upgraded);
+
+    // Proxying data
+    let (from_client, from_server) =
+        tokio::io::copy_bidirectional(&mut upgraded, &mut stream).await?;
+
+    // Print message when done
+    debug!(
+        "client wrote {} bytes and received {} bytes",
+        from_client, from_server
+    );
+    Ok(())
+}
+
+pub async fn proxy_request(
+    stream: TcpStream,
+    client_addr: SocketAddr,
+    socks_addr: SocketAddr,
+    auth_details: Option<&'static Auth>,
+    allowed_domains: Option<&'static Vec<String>>,
+    basic_http_header: Option<&'static HeaderValue>,
+) -> color_eyre::Result<()> {
+    let io = TokioIo::new(stream);
+
+    let serve_connection = service_fn(move |req| {
+        proxy(
+            req,
+            client_addr,
+            socks_addr,
+            auth_details,
+            allowed_domains,
+            basic_http_header,
+        )
+    });
+
+    tokio::task::spawn(async move {
+        if let Err(err) = http1::Builder::new()
+            .preserve_header_case(true)
+            .title_case_headers(true)
+            .serve_connection(io, serve_connection)
+            .with_upgrades()
+            .await
+        {
+            warn!("Failed to serve connection: {:?}", err);
+        }
+    });
+    Ok(())
+}

tests/proxy.rs

@@ -0,0 +1,56 @@
+use std::net::SocketAddr;
+
+use hyper::StatusCode;
+use sthp::proxy_request;
+use tokio::net::TcpListener;
+
+use color_eyre::Result;
+use socksprox::Socks5Server;
+use tokio::task::JoinHandle;
+
+async fn start_socks_server() -> Result<(JoinHandle<()>, SocketAddr)> {
+    // TODO: Currently, Socks5Server does not return the port it is bound to.
+    // To work around this, we will use TcpListener to obtain a random available port.
+    // After retrieving the port, we will immediately release it.
+    let listener = TcpListener::bind("localhost:0").await?;
+    let addr = listener.local_addr()?;
+    let port = addr.port();
+    eprintln!("socks proxy will listen on port {}", port);
+    // release port
+    drop(listener);
+
+    let mut server = Socks5Server::new("localhost", port, None, None)
+        .await
+        .unwrap();
+    let join_handle = tokio::task::spawn(async move {
+        server.serve().await;
+    });
+    Ok((join_handle, addr))
+}
+
+#[tokio::test]
+async fn simple_test() -> Result<()> {
+    let (_, socks_proxy_addr) = start_socks_server().await?;
+    let listener = TcpListener::bind("localhost:0").await?;
+    let addr = listener.local_addr()?;
+    let _ = tokio::task::spawn(async move {
+        let (stream, proxy_addr) = listener.accept().await?;
+        proxy_request(stream, proxy_addr, socks_proxy_addr, None, None, None).await?;
+        eprintln!("new connection from: {:?}", proxy_addr);
+        Ok::<_, color_eyre::eyre::Error>(())
+    });
+    eprintln!("http proxy will listen on {}", addr);
+
+    let client = reqwest::Client::builder()
+        .proxy(reqwest::Proxy::http(format!(
+            "http://localhost:{}",
+            addr.port()
+        ))?)
+        .build()?;
+
+    assert_eq!(
+        client.get("http://example.org").send().await?.status(),
+        StatusCode::OK
+    );
+    Ok(())
+}