reth_rpc_layer/
compression_layer.rs

1use jsonrpsee_http_client::{HttpBody, HttpRequest, HttpResponse};
2use std::{
3    future::Future,
4    pin::Pin,
5    task::{Context, Poll},
6};
7use tower::{Layer, Service};
8use tower_http::compression::{
9    predicate::SizeAbove, Compression, CompressionLayer as TowerCompressionLayer, CompressionLevel,
10};
11
12/// Responses smaller than this many bytes are sent uncompressed.
13///
14/// They fit in a single TCP segment anyway, and setting up a compressor costs far more than
15/// sending them as they are.
16const MIN_COMPRESSION_SIZE: u16 = 1024;
17
18/// This layer is a wrapper around [`tower_http::compression::CompressionLayer`] that integrates
19/// with jsonrpsee's HTTP types. It automatically compresses responses based on the client's
20/// `Accept-Encoding` header.
21#[expect(missing_debug_implementations)]
22#[derive(Clone)]
23pub struct CompressionLayer {
24    inner_layer: TowerCompressionLayer<SizeAbove>,
25}
26
27impl CompressionLayer {
28    /// Creates a new compression layer with zstd, gzip, brotli and deflate enabled.
29    ///
30    /// Responses are compressed at the fastest level, and only if they are at least 1 KiB.
31    pub fn new() -> Self {
32        Self::from_layer(TowerCompressionLayer::new().gzip(true).br(true).deflate(true).zstd(true))
33    }
34
35    /// Creates a compression layer allowing only the listed algorithms.
36    ///
37    /// The client's `Accept-Encoding` quality values determine selection among allowed algorithms;
38    /// ties use `tower_http`'s fixed **zstd > br > gzip > deflate** order. The order of `algos` is
39    /// ignored. Omitted quality values default to 1; if no allowed algorithm is acceptable, the
40    /// response is uncompressed.
41    pub fn with_algorithms(algos: &[impl AsRef<str>]) -> Self {
42        // Clear tower_http's enabled-by-default algorithms to make `algos` an allowlist.
43        let mut layer = TowerCompressionLayer::new().no_zstd().no_gzip().no_deflate().no_br();
44
45        for algo in algos {
46            match algo.as_ref() {
47                "zstd" => layer = layer.zstd(true),
48                "deflate" => layer = layer.deflate(true),
49                "gzip" => layer = layer.gzip(true),
50                "br" | "brotli" => layer = layer.br(true),
51                _ => {}
52            }
53        }
54
55        Self::from_layer(layer)
56    }
57
58    /// Applies the compression level and the minimum response size to `layer`.
59    fn from_layer(layer: TowerCompressionLayer) -> Self {
60        let inner_layer = layer
61            .quality(CompressionLevel::Fastest)
62            .compress_when(SizeAbove::new(MIN_COMPRESSION_SIZE));
63        Self { inner_layer }
64    }
65}
66
67impl Default for CompressionLayer {
68    /// Creates a new compression layer with default settings.
69    /// See [`CompressionLayer::new`] for details.
70    fn default() -> Self {
71        Self::new()
72    }
73}
74
75impl<S> Layer<S> for CompressionLayer {
76    type Service = CompressionService<S>;
77
78    fn layer(&self, inner: S) -> Self::Service {
79        CompressionService { compression: self.inner_layer.layer(inner) }
80    }
81}
82
83/// Service that performs response compression.
84///
85/// Created by [`CompressionLayer`].
86#[expect(missing_debug_implementations)]
87#[derive(Clone)]
88pub struct CompressionService<S> {
89    compression: Compression<S, SizeAbove>,
90}
91
92impl<S> Service<HttpRequest> for CompressionService<S>
93where
94    S: Service<HttpRequest, Response = HttpResponse>,
95    S::Future: Send + 'static,
96{
97    type Response = HttpResponse;
98    type Error = S::Error;
99    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
100
101    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
102        self.compression.poll_ready(cx)
103    }
104
105    fn call(&mut self, req: HttpRequest) -> Self::Future {
106        let fut = self.compression.call(req);
107
108        Box::pin(async move {
109            let resp = fut.await?;
110            let (parts, compressed_body) = resp.into_parts();
111            let http_body = HttpBody::new(compressed_body);
112
113            Ok(Self::Response::from_parts(parts, http_body))
114        })
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121    use http::header::{ACCEPT_ENCODING, CONTENT_ENCODING};
122    use http_body_util::BodyExt;
123    use jsonrpsee_http_client::{HttpRequest, HttpResponse};
124    use std::{convert::Infallible, future::ready};
125
126    const TEST_DATA: &str = "compress test data ";
127    const REPEAT_COUNT: usize = 1000;
128
129    #[derive(Clone)]
130    struct MockRequestService(String);
131
132    impl Service<HttpRequest> for MockRequestService {
133        type Response = HttpResponse;
134        type Error = Infallible;
135        type Future = std::future::Ready<Result<Self::Response, Self::Error>>;
136
137        fn poll_ready(
138            &mut self,
139            _: &mut std::task::Context<'_>,
140        ) -> std::task::Poll<Result<(), Self::Error>> {
141            std::task::Poll::Ready(Ok(()))
142        }
143
144        fn call(&mut self, _: HttpRequest) -> Self::Future {
145            let body = HttpBody::from(self.0.clone());
146            let response = HttpResponse::builder().body(body).unwrap();
147            ready(Ok(response))
148        }
149    }
150
151    fn setup_compression_service(
152        body: String,
153    ) -> impl Service<HttpRequest, Response = HttpResponse, Error = Infallible> {
154        CompressionLayer::new().layer(MockRequestService(body))
155    }
156
157    async fn get_response_size(response: HttpResponse) -> usize {
158        // Get the total size of the response body
159        response.into_body().collect().await.unwrap().to_bytes().len()
160    }
161
162    #[tokio::test]
163    async fn test_gzip_compression() {
164        let mut service = setup_compression_service(TEST_DATA.repeat(REPEAT_COUNT));
165        let request =
166            HttpRequest::builder().header(ACCEPT_ENCODING, "gzip").body(HttpBody::empty()).unwrap();
167
168        let uncompressed_len = TEST_DATA.repeat(REPEAT_COUNT).len();
169
170        // Make the request
171        let response = service.call(request).await.unwrap();
172
173        // Verify the response has gzip content-encoding
174        assert_eq!(
175            response.headers().get(CONTENT_ENCODING).unwrap(),
176            "gzip",
177            "Response should be gzip encoded"
178        );
179
180        // Verify the response body is actually compressed (should be smaller than original)
181        let compressed_size = get_response_size(response).await;
182        assert!(
183            compressed_size < uncompressed_len,
184            "Compressed size ({compressed_size}) should be smaller than original size ({uncompressed_len})"
185        );
186    }
187
188    #[tokio::test]
189    async fn test_no_compression_when_not_requested() {
190        // Create a service with compression
191        let mut service = setup_compression_service(TEST_DATA.repeat(REPEAT_COUNT));
192        let request = HttpRequest::builder().body(HttpBody::empty()).unwrap();
193
194        let response = service.call(request).await.unwrap();
195        assert!(
196            response.headers().get(CONTENT_ENCODING).is_none(),
197            "Response should not be compressed when not requested"
198        );
199
200        let uncompressed_len = TEST_DATA.repeat(REPEAT_COUNT).len();
201
202        // Verify the response body matches the original size
203        let response_size = get_response_size(response).await;
204        assert!(
205            response_size == uncompressed_len,
206            "Response size ({response_size}) should equal original size ({uncompressed_len})"
207        );
208    }
209
210    #[tokio::test]
211    async fn test_no_compression_for_small_responses() {
212        let body = "a".repeat(MIN_COMPRESSION_SIZE as usize - 1);
213        let mut service = setup_compression_service(body.clone());
214        let request =
215            HttpRequest::builder().header(ACCEPT_ENCODING, "br").body(HttpBody::empty()).unwrap();
216
217        let response = service.call(request).await.unwrap();
218        assert!(response.headers().get(CONTENT_ENCODING).is_none());
219        assert_eq!(get_response_size(response).await, body.len());
220    }
221}
reth_rpc_layer/compression_layer.rs

reth_rpc_layer/
compression_layer.rs
