From 62734ac9ae9de86e6ab18f93c56f0a513a6d44e5 Mon Sep 17 00:00:00 2001 From: krabat-l Date: Mon, 9 Feb 2026 19:42:59 +0800 Subject: [PATCH] feat: add x-response-size header to all HTTP responses --- Cargo.lock | 3 + Cargo.toml | 1 + bin/debug-trace-server/Cargo.toml | 3 + bin/debug-trace-server/src/main.rs | 7 +- bin/debug-trace-server/src/response_size.rs | 210 ++++++++++++++++++++ 5 files changed, 223 insertions(+), 1 deletion(-) create mode 100644 bin/debug-trace-server/src/response_size.rs diff --git a/Cargo.lock b/Cargo.lock index d098742..f083852 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1899,11 +1899,14 @@ dependencies = [ "alloy-rpc-types-eth", "alloy-rpc-types-trace", "alloy-signer-local", + "bytes", "clap", "dashmap 5.5.3", "dotenvy", "eyre", "http", + "http-body", + "http-body-util", "jsonrpsee 0.24.10", "libc", "metrics", diff --git a/Cargo.toml b/Cargo.toml index c88183c..79f4f68 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -68,6 +68,7 @@ reth-trie-sparse = { git = "https://github.com/paradigmxyz/reth", tag = "v1.6.0" # revm http = "1" +http-body = "1" revm = { version = "27.1.0", default-features = false } revm-inspectors = { version = "0.27.1", features = ["std", "js-tracer"] } rustc-hash = "2.0" diff --git a/bin/debug-trace-server/Cargo.toml b/bin/debug-trace-server/Cargo.toml index 169cf73..83ec99b 100644 --- a/bin/debug-trace-server/Cargo.toml +++ b/bin/debug-trace-server/Cargo.toml @@ -20,6 +20,7 @@ dashmap = "5.5" # Error handling eyre.workspace = true http.workspace = true +http-body.workspace = true # RPC framework jsonrpsee = { version = "0.24", features = ["server", "macros"] } libc = "0.2" @@ -51,6 +52,8 @@ validator-core = { path = "../../crates/validator-core" } alloy-consensus.workspace = true alloy-network = "1.0.23" alloy-primitives.workspace = true +bytes = "1" +http-body-util = "0.1" # Transaction signing for integration tests alloy-signer-local = "1.0.23" diff --git a/bin/debug-trace-server/src/main.rs b/bin/debug-trace-server/src/main.rs index 1a39552..5ade8be 100644 --- a/bin/debug-trace-server/src/main.rs +++ b/bin/debug-trace-server/src/main.rs @@ -59,6 +59,7 @@ use validator_core::{ mod data_provider; mod metrics; mod response_cache; +mod response_size; mod rpc_service; mod timing; @@ -341,7 +342,11 @@ async fn main() -> Result<()> { // Start server let server = Server::builder() .max_response_body_size(u32::MAX) - .set_http_middleware(tower::ServiceBuilder::new().layer(timing::TimingHeaderLayer)) + .set_http_middleware( + tower::ServiceBuilder::new() + .layer(response_size::ResponseSizeLayer) + .layer(timing::TimingHeaderLayer), + ) .build(&args.addr) .await?; let addr = server.local_addr()?; diff --git a/bin/debug-trace-server/src/response_size.rs b/bin/debug-trace-server/src/response_size.rs new file mode 100644 index 0000000..bcfa62e --- /dev/null +++ b/bin/debug-trace-server/src/response_size.rs @@ -0,0 +1,210 @@ +//! Response size HTTP middleware. +//! +//! Provides [`ResponseSizeLayer`] which reads the response body size via +//! `http_body::Body::size_hint` and sets the `x-response-size` header (in bytes) +//! on every HTTP response. + +use std::{ + future::Future, + pin::Pin, + task::{Context, Poll}, +}; + +use http::{HeaderName, HeaderValue}; +use http_body::Body; +use pin_project_lite::pin_project; +use tower::{Layer, Service}; + +/// Header name for response size in bytes. +pub const RESPONSE_SIZE_HEADER_NAME: &str = "x-response-size"; + +/// Tower layer that creates [`ResponseSizeService`] instances. +#[derive(Debug, Clone, Copy, Default)] +pub struct ResponseSizeLayer; + +impl Layer for ResponseSizeLayer { + type Service = ResponseSizeService; + + #[inline] + fn layer(&self, inner: S) -> Self::Service { + ResponseSizeService { inner } + } +} + +/// HTTP middleware service that measures response body size and sets +/// the `x-response-size` header. +#[derive(Debug, Clone)] +pub struct ResponseSizeService { + inner: S, +} + +impl Service> for ResponseSizeService +where + S: Service, Response = http::Response> + Clone + Send + 'static, + S::Future: Send, + S::Error: Send, + ReqBody: Send + 'static, + ResBody: Body + Send + 'static, +{ + type Response = http::Response; + type Error = S::Error; + type Future = ResponseSizeFuture; + + #[inline] + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, req: http::Request) -> Self::Future { + ResponseSizeFuture::new(self.inner.call(req)) + } +} + +pin_project! { + /// Future that reads the response body size and injects the header. + pub struct ResponseSizeFuture { + #[pin] + inner: F, + _error: std::marker::PhantomData, + _body: std::marker::PhantomData, + } +} + +impl ResponseSizeFuture { + fn new(inner: F) -> Self { + Self { inner, _error: std::marker::PhantomData, _body: std::marker::PhantomData } + } +} + +impl Future for ResponseSizeFuture +where + F: Future, E>>, + ResBody: Body, +{ + type Output = Result, E>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.project(); + + match this.inner.poll(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(Ok(mut response)) => { + let size = response.body().size_hint().exact().unwrap_or(0); + let mut buf = [0u8; 24]; + if let Some(value_str) = format_u64(size, &mut buf) { + let header_name = HeaderName::from_static(RESPONSE_SIZE_HEADER_NAME); + if let Ok(header_value) = HeaderValue::from_str(value_str) { + response.headers_mut().insert(header_name, header_value); + } + } + Poll::Ready(Ok(response)) + } + Poll::Ready(Err(e)) => Poll::Ready(Err(e)), + } + } +} + +#[inline] +fn format_u64(value: u64, buf: &mut [u8; 24]) -> Option<&str> { + use std::io::Write; + let mut cursor = std::io::Cursor::new(&mut buf[..]); + write!(cursor, "{}", value).ok()?; + let len = cursor.position() as usize; + std::str::from_utf8(&buf[..len]).ok() +} + +#[cfg(test)] +mod tests { + use std::convert::Infallible; + + use bytes::Bytes; + use http::{Request as HttpRequest, Response as HttpResponse}; + use http_body_util::Full; + + use super::*; + + #[derive(Clone)] + struct MockService { + body: &'static str, + } + + impl Service>> for MockService { + type Response = HttpResponse>; + type Error = Infallible; + type Future = Pin> + Send>>; + + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, _req: HttpRequest>) -> Self::Future { + let body = self.body; + Box::pin(async move { Ok(HttpResponse::new(Full::new(Bytes::from(body)))) }) + } + } + + #[tokio::test] + async fn test_response_size_header_set() { + let mock = MockService { body: "hello world" }; + let layer = ResponseSizeLayer; + let mut service = layer.layer(mock); + + let req = HttpRequest::new(Full::new(Bytes::new())); + let response = service.call(req).await.unwrap(); + + let header = response + .headers() + .get(RESPONSE_SIZE_HEADER_NAME) + .expect("x-response-size header should exist"); + let size: u64 = header.to_str().unwrap().parse().unwrap(); + assert_eq!(size, 11); // "hello world" = 11 bytes + } + + #[tokio::test] + async fn test_response_size_empty_body() { + let mock = MockService { body: "" }; + let layer = ResponseSizeLayer; + let mut service = layer.layer(mock); + + let req = HttpRequest::new(Full::new(Bytes::new())); + let response = service.call(req).await.unwrap(); + + let header = response + .headers() + .get(RESPONSE_SIZE_HEADER_NAME) + .expect("x-response-size header should exist"); + let size: u64 = header.to_str().unwrap().parse().unwrap(); + assert_eq!(size, 0); + } + + #[tokio::test] + async fn test_response_size_large_body() { + let large = "x".repeat(1_000_000); + let large_static: &'static str = Box::leak(large.into_boxed_str()); + let mock = MockService { body: large_static }; + let layer = ResponseSizeLayer; + let mut service = layer.layer(mock); + + let req = HttpRequest::new(Full::new(Bytes::new())); + let response = service.call(req).await.unwrap(); + + let header = response + .headers() + .get(RESPONSE_SIZE_HEADER_NAME) + .expect("x-response-size header should exist"); + let size: u64 = header.to_str().unwrap().parse().unwrap(); + assert_eq!(size, 1_000_000); + } + + #[test] + fn test_format_u64_zero() { + let mut buf = [0u8; 24]; + assert_eq!(format_u64(0, &mut buf), Some("0")); + } + + #[test] + fn test_format_u64_max() { + let mut buf = [0u8; 24]; + assert_eq!(format_u64(u64::MAX, &mut buf), Some("18446744073709551615")); + } +}