refactor: added some tests

Cargo.toml

@@ -22,3 +22,7 @@ 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" }
+
+[dev-dependencies]
+socksprox = { version = "0.1" }
+reqwest = { version = "0.11" }

src/main.rs

@@ -3,7 +3,7 @@ use color_eyre::eyre::Result;
 
 use sthp::proxy::auth::Auth;
 use sthp::proxy_request;
-use tracing::info;
+use tracing::{error, info};
 use tracing_subscriber::EnvFilter;
 
 use std::net::{Ipv4Addr, SocketAddr};
@@ -68,7 +68,9 @@ async fn main() -> Result<()> {
     loop {
         let (stream, _) = listener.accept().await?;
         tokio::task::spawn(async move {
-            proxy_request(stream, socks_addr, auth_details, allowed_domains).await
+            if let Err(e) = proxy_request(stream, socks_addr, auth_details, allowed_domains).await {
+                error!("Error proxying request: {}", e);
+            }
         });
     }
 }

src/proxy/mod.rs

@@ -26,7 +26,7 @@ async fn proxy(
     socks_addr: SocketAddr,
     auth: &'static Option<Auth>,
     allowed_domains: &Option<Vec<String>>,
-) -> Result<Response<BoxBody<Bytes, hyper::Error>>, hyper::Error> {
+) -> Result<Response<BoxBody<Bytes, hyper::Error>>> {
     let uri = req.uri();
     let method = req.method();
     debug!("Proxying request: {} {}", method, uri);
@@ -80,7 +80,7 @@ async fn proxy(
             )
             .await
             .unwrap(),
-            None => Socks5Stream::connect(socks_addr, addr).await.unwrap(),
+            None => Socks5Stream::connect(socks_addr, addr).await?,
         };
         let io = TokioIo::new(stream);
 
@@ -149,7 +149,7 @@ pub async fn proxy_request(
     socks_addr: SocketAddr,
     auth_details: &'static Option<Auth>,
     allowed_domains: &'static Option<Vec<String>>,
-) {
+) -> color_eyre::Result<()> {
     let io = TokioIo::new(stream);
 
     let serve_connection =
@@ -166,4 +166,5 @@ pub async fn proxy_request(
             warn!("Failed to serve connection: {:?}", err);
         }
     });
+    Ok(())
 }

tests/proxy.rs

@@ -0,0 +1,56 @@
+use std::{net::SocketAddr, time::Duration};
+
+use http::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 doesnt return what port it binded
+    // so we will use TcpListener to get the random port and release it immediatly
+    let listener = TcpListener::bind("localhost:0").await?;
+    let addr = listener.local_addr()?;
+    let port = addr.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, socks_proxy_addr, &None, &None).await?;
+        eprintln!("new connection from: {:?}", proxy_addr);
+        Ok::<_, color_eyre::eyre::Error>(())
+    });
+    assert_eq!("hello", "hello");
+
+    let client = reqwest::Client::builder()
+        .proxy(reqwest::Proxy::http(format!(
+            "http://{}:{}",
+            addr.ip(),
+            addr.port()
+        ))?)
+        .build()?;
+
+    assert_eq!(
+        client.get("http://example.org").send().await?.status(),
+        StatusCode::OK
+    );
+    eprintln!("http proxy handle dropped");
+    Ok(())
+}