diff --git a/proxy/src/main.rs b/proxy/src/main.rs index a4dca9b..00283cd 100644 --- a/proxy/src/main.rs +++ b/proxy/src/main.rs @@ -7,6 +7,19 @@ use rama::telemetry::tracing::{ mod server; use server::proxy::run_server; +/// CLI arguments for configuring proxy behavior. +#[derive(Parser)] +#[command( + about = "A security-focused HTTP/HTTPS proxy for Safe-chain", + version, + author +)] +struct Args { + /// TCP port binding. Use 0 for OS-assigned port (recommended for avoiding conflicts). + #[arg(short, long, default_value_t = 0)] + port: u16, +} + #[tokio::main] async fn main() { let args = Args::parse(); @@ -14,6 +27,10 @@ async fn main() { run_server(args.port).await; } +/// Configures structured logging with runtime control via `RUST_LOG` environment variable. +/// +/// Defaults to INFO level to balance visibility with performance. +/// Use `RUST_LOG=debug` or `RUST_LOG=trace` for troubleshooting. fn setup_tracing() { tracing::subscriber::registry() .with(fmt::layer()) @@ -25,9 +42,3 @@ fn setup_tracing() { .init(); tracing::info!("Tracing is set up"); } - -#[derive(Parser, Debug)] -struct Args { - #[arg(short, long, default_value_t = 0)] - port: u16, -} diff --git a/proxy/src/server/proxy.rs b/proxy/src/server/proxy.rs index 896adab..1bd0356 100644 --- a/proxy/src/server/proxy.rs +++ b/proxy/src/server/proxy.rs @@ -1,3 +1,8 @@ +//! HTTP/HTTPS proxy implementation using the Rama framework. +//! +//! Supports both CONNECT tunneling (for HTTPS) and plain HTTP proxying. +//! Includes graceful shutdown, body size limits, and structured logging. + use rama::{ extensions::ExtensionsMut, http::{ @@ -22,6 +27,14 @@ use rama::{ }; use std::{convert::Infallible, time::Duration}; +/// Maximum allowed body size for proxied requests and responses. +/// Protects against memory exhaustion from excessively large payloads. +const MAX_BODY_SIZE: usize = 500 * 1024 * 1024; // 500 MB + +/// Starts the proxy server with graceful shutdown support. +/// +/// Spawns the server task and waits for a shutdown signal (e.g., Ctrl+C). +/// Active connections are given up to 30 seconds to complete before forced termination. pub async fn run_server(port: u16) { let graceful = rama::graceful::Shutdown::default(); @@ -33,19 +46,24 @@ pub async fn run_server(port: u16) { .expect("graceful shutdown"); } +/// Core server task that binds to a port and serves HTTP/HTTPS traffic. +/// +/// Configures the HTTP server with: +/// - CONNECT method upgrade for HTTPS tunneling +/// - Hop-by-hop header removal (Connection, Keep-Alive, etc.) +/// - Body size limits to prevent resource exhaustion +/// - Request/response tracing for observability async fn server_task(guard: rama::graceful::ShutdownGuard, port: u16) { - let tcp_address = format!("127.0.0.1:{}", port); - let tcp_service = TcpListener::build() - .bind(tcp_address) + .bind(format!("127.0.0.1:{}", port)) .await - .expect("bind tcp proxy"); + .unwrap_or_else(|e| panic!("Failed to bind tcp proxy to 127.0.0.1:{}: {}", port, e)); - let local_address = tcp_service.local_addr().expect("tcp proxy assigned a port"); - tracing::info!("Safe-chain proxy running on {local_address}"); + let local_address = tcp_service + .local_addr() + .expect("Could not get bound local address for TCP server"); let exec = Executor::graceful(guard.clone()); - let http_service = HttpServer::auto(exec).service( ( TraceLayer::new_for_http(), @@ -61,18 +79,25 @@ async fn server_task(guard: rama::graceful::ShutdownGuard, port: u16) { .into_layer(service_fn(http_plain_proxy)), ); + tracing::info!(proxy.address = %local_address, "safe-chain proxy running"); + tcp_service .serve_graceful( guard, ( // protect the http proxy from too large bodies, both from request and response end - BodyLimitLayer::symmetric(500 * 1024 * 1024), + BodyLimitLayer::symmetric(MAX_BODY_SIZE), ) .into_layer(http_service), ) .await; } +/// Handles HTTPS CONNECT requests by establishing a TCP tunnel. +/// +/// Extracts the target host:port from the request and stores it in request extensions +/// for use by the TCP forwarder. Returns 200 OK to signal successful tunnel establishment, +/// or 400 BAD REQUEST if the target cannot be determined. async fn http_connect_accept(mut req: Request) -> Result<(Response, Request), Response> { match RequestContext::try_from(&req).map(|ctx| ctx.host_with_port()) { Ok(authority) => { @@ -84,23 +109,31 @@ async fn http_connect_accept(mut req: Request) -> Result<(Response, Request), Re req.extensions_mut().insert(ProxyTarget(authority)); } Err(err) => { - tracing::error!("error extracting authority: {err:?}"); + tracing::error!(uri = %req.uri(), "error extracting authority: {err:?}"); return Err(StatusCode::BAD_REQUEST.into_response()); } } - return Ok((StatusCode::OK.into_response(), req)); + Ok((StatusCode::OK.into_response(), req)) } +/// Forwards plain HTTP requests to their destination. +/// +/// Uses an HTTP client to relay requests transparently. Returns 502 BAD GATEWAY +/// if the upstream server is unreachable or returns an error. The `Infallible` return +/// type indicates this handler always produces a response (never panics the service). async fn http_plain_proxy(req: Request) -> Result { - let client = EasyHttpWebClient::default(); + let uri = req.uri().clone(); - return match client.serve(req).await { + let client = EasyHttpWebClient::default(); + tracing::info!(uri = %uri, "serving http over proxy"); + + match client.serve(req).await { Ok(resp) => Ok(resp), Err(err) => { - tracing::error!("Error forwarding request: {err:?}"); + tracing::error!(uri = %uri, "error forwarding request: {err:?}"); let resp = StatusCode::BAD_GATEWAY.into_response(); Ok(resp) } - }; + } }