From 000d1362479add1e00796103cf86fd11990f004c Mon Sep 17 00:00:00 2001 From: William Howard Date: Sun, 31 May 2026 12:26:56 +0100 Subject: [PATCH] code quality improvements --- Cargo.lock | 1 - Cargo.toml | 1 - src/client.rs | 8 ++-- src/lib.rs | 105 +++++++++++++++++++--------------------------- src/main.rs | 2 +- src/metrics.rs | 12 +++--- src/routing.rs | 6 ++- tests/config.yaml | 2 +- 8 files changed, 60 insertions(+), 77 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e2da5c1..bc1fe11 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -282,7 +282,6 @@ dependencies = [ "axum", "axum-server", "env_logger", - "http", "hyper", "hyper-util", "log", diff --git a/Cargo.toml b/Cargo.toml index f785613..c2401d8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,6 +19,5 @@ serde_yaml = "0.9.34" tokio = { version = "1.43.0", features = ["full"] } [dev-dependencies] -http = "1.2.0" reqwest = { version = "0.12.12", features = ["rustls-tls"] } wiremock = "0.6" diff --git a/src/client.rs b/src/client.rs index 05d5b77..d3acdc9 100644 --- a/src/client.rs +++ b/src/client.rs @@ -13,23 +13,23 @@ use tokio::time::timeout; type HttpClient = hyper_util::client::legacy::Client; #[derive(Debug)] -pub struct Client { +pub(crate) struct Client { client: HttpClient, timeout: Duration, } impl Client { - pub fn new(timeout: Option) -> Client { + pub(crate) fn new(timeout_ms: Option) -> Client { let client: HttpClient = hyper_util::client::legacy::Client::<(), ()>::builder(TokioExecutor::new()) .build(HttpConnector::new()); Client { client, - timeout: Duration::from_millis(timeout.unwrap_or(60_000)), + timeout: Duration::from_millis(timeout_ms.unwrap_or(60_000)), } } - pub async fn make_request(&self, req: Request) -> Response { + pub(crate) async fn make_request(&self, req: Request) -> Response { match timeout(self.timeout, self.client.request(req)).await { Ok(result) => match result { Ok(response) => response.into_response(), diff --git a/src/lib.rs b/src/lib.rs index bc99a67..a6b6e24 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -27,7 +27,7 @@ use hyper::http::{ HeaderName, }; use log::{debug, info, warn}; -use serde::{Deserialize, Serialize}; +use serde::Deserialize; use tokio::signal; #[cfg(unix)] use tokio::signal::unix::SignalKind; @@ -54,25 +54,25 @@ const HOP_BY_HOP_HEADERS: [HeaderName; 9] = [ header::PROXY_AUTHENTICATE, ]; -#[derive(Debug, PartialEq, Serialize, Deserialize)] +#[derive(Debug, PartialEq, Deserialize)] #[serde(deny_unknown_fields)] struct TlsConfig { cert_path: String, key_path: String, } -#[derive(Debug, PartialEq, Serialize, Deserialize)] +#[derive(Debug, PartialEq, Deserialize)] #[serde(deny_unknown_fields)] struct Config { listen: SocketAddr, tls: TlsConfig, - timeout: Option, + timeout_ms: Option, backends: Vec, } -#[derive(Debug, Eq, PartialEq, Serialize, Deserialize, Clone)] +#[derive(Debug, Eq, PartialEq, Deserialize, Clone)] #[serde(tag = "backend_type", rename_all = "lowercase", deny_unknown_fields)] -pub enum Backend { +pub(crate) enum Backend { Single { name: String, location: String, @@ -115,12 +115,12 @@ impl ProxyConfig { } #[derive(Debug)] -pub struct BackendState { +pub(crate) struct BackendState { rr_count: AtomicUsize, } #[derive(Debug)] -pub struct RoutingState { +pub(crate) struct RoutingState { backends: HashMap, // keyed by name, LoadBalanced backends only } @@ -149,7 +149,7 @@ struct ProxyState { } #[derive(Debug, Clone)] -pub struct ResponseContext { +pub(crate) struct ResponseContext { backend_location: String, } @@ -246,13 +246,13 @@ fn get_host(req: &Request) -> Option { Some(host) } -fn error_response( - mut response: Response, - status: StatusCode, - message: String, -) -> Response { - *response.body_mut() = Body::from(message); +fn error_response(status: StatusCode, message: &str) -> Response { + let mut response = Response::new(Body::from(message.to_owned())); *response.status_mut() = status; + response.headers_mut().insert( + CONTENT_TYPE, + HeaderValue::from_static("text/plain; charset=utf-8"), + ); response } @@ -264,8 +264,6 @@ async fn proxy_handler( proxy_config, routing_state, } = state; - let mut response = Response::new(Body::empty()); - debug!( "Request version: {:?} method: {} uri: {} headers: {:?}", req.version(), @@ -279,9 +277,8 @@ async fn proxy_handler( Version::HTTP_10 | Version::HTTP_11 | Version::HTTP_2 => {} _ => { return Ok(error_response( - response, StatusCode::HTTP_VERSION_NOT_SUPPORTED, - format!("Unsupported HTTP version: {:?}", req.version()), + &format!("Unsupported HTTP version: {:?}", req.version()), )) } } @@ -296,36 +293,34 @@ async fn proxy_handler( no_proxy, host_authority ); - match (req.method(), req.uri().path(), no_proxy, host_authority) { + let response = match (req.method(), req.uri().path(), no_proxy, host_authority) { // Proxy internal endpoints - (&Method::GET, "/status", true, _) => { - *response.body_mut() = Body::from("The proxy is running"); - } + (&Method::GET, "/status", true, _) => Response::new(Body::from("The proxy is running")), (&Method::GET, "/metrics", true, _) => match encode_metrics() { Ok(encoded_metrics) => { - *response.body_mut() = Body::from(encoded_metrics); + let mut response = Response::new(Body::from(encoded_metrics)); response.headers_mut().insert( CONTENT_TYPE, HeaderValue::from_static("text/plain; charset=utf-8"), ); + response } Err(e) => { warn!("Error encoding metrics: {e}"); - *response.body_mut() = Body::from(format!("Error encoding metrics: {e}")); - *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR; + error_response( + StatusCode::INTERNAL_SERVER_ERROR, + &format!("Error encoding metrics: {e}"), + ) } }, // x-no-proxy request to an unknown internal path - (_, _, true, _) => { - *response.status_mut() = StatusCode::NOT_FOUND; - } + (_, _, true, _) => error_response(StatusCode::NOT_FOUND, ""), // A non internal request, but the host header has not been defined (_, _, false, None) => { debug!("Host header not defined"); - *response.body_mut() = Body::from("Host header not defined"); - *response.status_mut() = StatusCode::NOT_FOUND; + error_response(StatusCode::NOT_FOUND, "Host header not defined") } // Proxy the request @@ -338,13 +333,9 @@ async fn proxy_handler( ); match backend_location { - None => { - *response.status_mut() = StatusCode::NOT_FOUND; - } + None => error_response(StatusCode::NOT_FOUND, ""), Some(backend_location) => { - // Proxy to backend - - // Scheme currently hardcoded to http (given this is a TLS terminating proxy) + // Backend connections are plain HTTP — TLS is terminated at the proxy let scheme = "http"; // Default to "/" if the URI has no path component @@ -356,38 +347,35 @@ async fn proxy_handler( let uri = match Uri::builder() .scheme(scheme) - .authority(backend_location.clone()) + .authority(backend_location.as_str()) .path_and_query(path_and_query) .build() { Ok(uri) => uri, Err(e) => { warn!("Failed to build backend URI: {e}"); - *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR; - return Ok(response); + return Ok(error_response(StatusCode::INTERNAL_SERVER_ERROR, "")); } }; // Simply take the existing request and mutate the uri and headers - *req.uri_mut() = uri.clone(); + debug!("Proxying request to: {}", uri); + *req.uri_mut() = uri; adjust_backend_request_headers(&mut req, &host); - // If the backend scheme is http, adjust the original request HTTP version to 1 - // (It seems that the HTTP2 implementation requires TLS) - if scheme == "http" { - *req.version_mut() = Version::HTTP_11; - } - response = proxy_config.client.make_request(req).await; + // Downgrade to HTTP/1.1 for backend connections + *req.version_mut() = Version::HTTP_11; + let mut response = proxy_config.client.make_request(req).await; adjust_backend_response_headers(&mut response); debug!( - "Proxied response from: {} | Status: {} | Response headers: {:?}", - uri, + "Proxied response | Status: {} | Headers: {:?}", response.status(), response.headers() ); response .extensions_mut() .insert(ResponseContext { backend_location }); + response } } } @@ -428,17 +416,12 @@ pub async fn run_server(config_path: String) -> Result<()> { let listen_address = config.listen; - let client = client::Client::new(config.timeout); + let client = client::Client::new(config.timeout_ms); let routing_state = Arc::new(RoutingState::new(&config)); let proxy_config = Arc::new(ProxyConfig::new(config, client)); - let proxy_state = ProxyState { - proxy_config: proxy_config.clone(), - routing_state, - }; - let current_dir = env::current_dir().context("Unable to determine current directory")?; let tls_config = RustlsConfig::from_pem_file( current_dir.join(&proxy_config.config.tls.cert_path), @@ -451,6 +434,11 @@ pub async fn run_server(config_path: String) -> Result<()> { info!("backend: {backend}"); } + let proxy_state = ProxyState { + proxy_config, + routing_state, + }; + let app = Router::new() .route("/", any(proxy_handler)) .route("/{*wildcard}", any(proxy_handler)) @@ -674,12 +662,7 @@ backends: #[tokio::test] async fn test_error_response() { - let original_response = Response::new(Body::from("test")); - let response = error_response( - original_response, - StatusCode::BAD_REQUEST, - "test error".to_string(), - ); + let response = error_response(StatusCode::BAD_REQUEST, "test error"); assert_eq!(response.status(), 400); let body = axum::body::to_bytes(response.into_body(), 1024) .await diff --git a/src/main.rs b/src/main.rs index b757ab0..87d70eb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,7 +1,7 @@ use anyhow::Result; use brachyura::run_server; -#[tokio::main(worker_threads = 4)] +#[tokio::main] async fn main() -> Result<()> { let _ = rustls::crypto::ring::default_provider().install_default(); let config_path = String::from("./config.yaml"); diff --git a/src/metrics.rs b/src/metrics.rs index f7e4721..9ba2945 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -9,11 +9,11 @@ use prometheus::{ use crate::ResponseContext; -pub static METRICS: LazyLock = LazyLock::new(Metrics::new); +pub(crate) static METRICS: LazyLock = LazyLock::new(Metrics::new); -pub struct Metrics { - pub http_request_counter: IntCounterVec, - pub http_request_duration: HistogramVec, +pub(crate) struct Metrics { + pub(crate) http_request_counter: IntCounterVec, + pub(crate) http_request_duration: HistogramVec, } impl Metrics { @@ -36,7 +36,7 @@ impl Metrics { } } -pub fn encode_metrics() -> Result { +pub(crate) fn encode_metrics() -> Result { let mut buffer = Vec::new(); let encoder = TextEncoder::new(); let metric_families = prometheus::gather(); @@ -44,7 +44,7 @@ pub fn encode_metrics() -> Result { Ok(String::from_utf8(buffer)?) } -pub async fn record_metrics(req: Request, next: Next) -> impl IntoResponse { +pub(crate) async fn record_metrics(req: Request, next: Next) -> impl IntoResponse { let start = Instant::now(); let response = next.run(req).await; diff --git a/src/routing.rs b/src/routing.rs index c6e2af9..9c6f52a 100644 --- a/src/routing.rs +++ b/src/routing.rs @@ -5,7 +5,7 @@ use super::{Backend, BackendState, RoutingState}; pub fn router( backends_config: &[Backend], - proxy_state: Arc, + routing_state: Arc, host_authority: String, ) -> Option { // Matches a given host header or authority with a backend @@ -15,7 +15,7 @@ pub fn router( match backend { Backend::Single { location, .. } => Some(location.clone()), Backend::LoadBalanced { name, locations } => { - let backend_state = proxy_state.backends.get(name)?; + let backend_state = routing_state.backends.get(name)?; round_robin_select(locations, backend_state) } } @@ -40,6 +40,8 @@ fn round_robin_select( if backend_locations.is_empty() { return None; } + // Relaxed ordering is sufficient as fetch_add is an atomic read-modify-write + // so no two threads can observe the same counter value let idx = backend_state.rr_count.fetch_add(1, Ordering::Relaxed) % backend_locations.len(); Some(backend_locations[idx].clone()) } diff --git a/tests/config.yaml b/tests/config.yaml index c487532..5ba7a96 100644 --- a/tests/config.yaml +++ b/tests/config.yaml @@ -4,7 +4,7 @@ tls: key_path: "tests/self-signed-cert/test.key" cert_path: "tests/self-signed-cert/test.crt" -timeout: 500 +timeout_ms: 500 backends: - name: "test.home"