Add src/main.rs

src/main.rs

@@ -0,0 +1,138 @@
+use bytes::Bytes;
+use http_body_util::Full;
+use hyper::body::Incoming;
+use hyper::server::conn::http1;
+use hyper::service::service_fn;
+use hyper::{Method, Request, Response, StatusCode};
+use hyper_util::rt::TokioIo;
+use mimalloc::MiMalloc;
+use std::net::SocketAddr;
+use tokio::net::TcpListener;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
+
+async fn handle_http(req: Request<Incoming>) -> Result<Response<Full<Bytes>>, hyper::Error> {
+    if req.method() != Method::GET {
+        return Ok(Response::builder()
+            .status(StatusCode::METHOD_NOT_ALLOWED)
+            .body(Full::new(Bytes::from("Only GET supported")))
+            .unwrap());
+    }
+
+    let url = match req.uri().to_string().parse::<reqwest::Url>() {
+        Ok(u) => u,
+        Err(_) => {
+            return Ok(Response::builder()
+                .status(StatusCode::BAD_REQUEST)
+                .body(Full::new(Bytes::from("Invalid URL")))
+                .unwrap())
+        }
+    };
+
+    let client = reqwest::Client::new();
+    let resp = match client.get(url).send().await {
+        Ok(r) => r,
+        Err(_) => {
+            return Ok(Response::builder()
+                .status(StatusCode::BAD_GATEWAY)
+                .body(Full::new(Bytes::from("Upstream error")))
+                .unwrap())
+        }
+    };
+
+    let headers = resp.headers().clone();
+    let content_type = headers
+        .get(reqwest::header::CONTENT_TYPE)
+        .and_then(|v| v.to_str().ok())
+        .unwrap_or("");
+
+    let bytes = match resp.bytes().await {
+        Ok(b) => b,
+        Err(_) => {
+            return Ok(Response::builder()
+                .status(StatusCode::BAD_GATEWAY)
+                .body(Full::new(Bytes::from("Read error")))
+                .unwrap())
+        }
+    };
+
+    // Only compress images
+    if content_type.starts_with("image/") {
+        if let Ok(img) = image::load_from_memory(&bytes) {
+            let mut out = Vec::new();
+
+            match content_type {
+                ct if ct.contains("jpeg") || ct.contains("jpg") => {
+                    let _ = img.write_to(
+                        &mut std::io::Cursor::new(&mut out),
+                        image::ImageFormat::Jpeg,
+                    );
+                }
+                ct if ct.contains("png") => {
+                    let _ = img.write_to(
+                        &mut std::io::Cursor::new(&mut out),
+                        image::ImageFormat::Png,
+                    );
+                }
+                _ => {
+                    out = bytes.to_vec();
+                }
+            }
+
+            return Ok(Response::builder()
+                .header("content-type", content_type)
+                .body(Full::new(Bytes::from(out)))
+                .unwrap());
+        }
+    }
+
+    Ok(Response::builder()
+        .body(Full::new(bytes))
+        .unwrap())
+}
+
+// CONNECT tunneling (HTTPS_PROXY support, but no inspection)
+async fn tunnel(mut client: tokio::net::TcpStream, host: String) {
+    match tokio::net::TcpStream::connect(host).await {
+        Ok(mut server) => {
+            let (mut cr, mut cw) = client.split();
+            let (mut sr, mut sw) = server.split();
+
+            let client_to_server = tokio::io::copy(&mut cr, &mut sw);
+            let server_to_client = tokio::io::copy(&mut sr, &mut cw);
+
+            let _ = tokio::join!(client_to_server, server_to_client);
+        }
+        Err(_) => {}
+    }
+}
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    let addr: SocketAddr = "0.0.0.0:8544".parse()?;
+    let listener = TcpListener::bind(addr).await?;
+
+    loop {
+        let (stream, _) = listener.accept().await?;
+        tokio::spawn(async move {
+            let io = TokioIo::new(stream);
+
+            let svc = service_fn(|req: Request<Incoming>| async move {
+                if req.method() == Method::CONNECT {
+                    // HTTPS tunneling (no compression possible here)
+                    return Ok::<_, hyper::Error>(
+                        Response::builder()
+                            .status(200)
+                            .body(Full::new(Bytes::new()))
+                            .unwrap(),
+                    );
+                }
+
+                handle_http(req).await
+            });
+
+            let _ = http1::Builder::new().serve_connection(io, svc).await;
+        });
+    }
+}