From 42dd982d425044b80bf8327f1d4233263cec6649 Mon Sep 17 00:00:00 2001 From: gigirassy Date: Thu, 28 May 2026 02:13:14 +0200 Subject: [PATCH] Add src/main.rs --- src/main.rs | 138 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 138 insertions(+) create mode 100644 src/main.rs diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..15216b1 --- /dev/null +++ b/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) -> Result>, 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::() { + 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| 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; + }); + } +} \ No newline at end of file