Polish up code

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,
-}

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");
+    let local_address = tcp_service
-    tracing::info!("Safe-chain proxy running on {local_address}");
+        .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<Response, Infallible> {
-    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)
         }
-    };
+    }
 }