Skip to main content

servers/http/
memory_limit.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Middleware for limiting total memory usage of concurrent HTTP request bodies.
16//!
17//! Admission works in two stages:
18//!
19//! 1. Upfront, the request's `Content-Length` (if present) is charged to the
20//!    shared [`ServerMemoryLimiter`], preserving the configured wait/fail
21//!    policy for regular requests.
22//! 2. While the body is streamed, [`AccountedBody`] charges every byte beyond
23//!    the upfront reservation. This closes two gaps of a header-only
24//!    reservation: HTTP/1.1 chunked requests carry no `Content-Length` (and
25//!    would otherwise be admitted for free), and a client can understate the
26//!    header while sending a much larger body.
27//!
28//! Permits acquired while streaming are held by [`BodyMemoryAccounting`],
29//! which is also inserted into the request extensions so they stay alive until
30//! the whole request is finished (extractors collect the body into owned
31//! buffers that outlive the body stream itself).
32
33use std::pin::Pin;
34use std::sync::atomic::{AtomicBool, Ordering};
35use std::sync::{Arc, Mutex};
36use std::task::{Context, Poll};
37
38use axum::body::Body;
39use axum::extract::{Request, State};
40use axum::middleware::Next;
41use axum::response::{IntoResponse, Response};
42use bytes::Bytes;
43use common_memory_manager::MemoryGuard;
44use http::StatusCode;
45use http_body::{Body as HttpBody, Frame};
46
47use crate::error::Result;
48use crate::request_memory_limiter::ServerMemoryLimiter;
49use crate::request_memory_metrics::RequestMemoryMetrics;
50
51/// Marker inserted by [`memory_limit_middleware`] when the request arrives
52/// with a non-identity `Content-Encoding` header. Route-local
53/// [`decoded_body_accounting_middleware`] uses it to account the decompressed
54/// body: only such requests expand during request decompression.
55#[derive(Clone, Copy)]
56pub(crate) struct ContentEncoded;
57
58pub async fn memory_limit_middleware(
59    State(limiter): State<ServerMemoryLimiter>,
60    req: Request,
61    next: Next,
62) -> Response {
63    let content_length = req
64        .headers()
65        .get(http::header::CONTENT_LENGTH)
66        .and_then(|v| v.to_str().ok())
67        .and_then(|v| v.parse::<u64>().ok())
68        .unwrap_or(0);
69
70    let _guard = match limiter.acquire(content_length).await {
71        Ok(guard) => guard,
72        Err(e) => {
73            return (
74                StatusCode::TOO_MANY_REQUESTS,
75                format!("Request body memory limit exceeded: {}", e),
76            )
77                .into_response();
78        }
79    };
80
81    let content_encoded = req
82        .headers()
83        .get(http::header::CONTENT_ENCODING)
84        .and_then(|v| v.to_str().ok())
85        .is_some_and(|v| !v.eq_ignore_ascii_case("identity"));
86
87    // Account whatever is actually streamed beyond the upfront reservation
88    // (chunked transfer has none).
89    let accounting = BodyMemoryAccounting::default();
90    // The middleware keeps its own handle so the acquired permits stay held
91    // for the whole request: the body wrapper and the extensions are dropped
92    // once the extractors finish collecting the body, before the handler is
93    // done with the decoded data.
94    let retained_accounting = accounting.clone();
95    let (mut parts, body) = req.into_parts();
96    // Expose the limiter to handlers that decompress request bodies on their
97    // own (e.g. the Loki protobuf push) so they can charge the decoded size.
98    parts.extensions.insert(limiter.clone());
99    parts.extensions.insert(accounting.clone());
100    if content_encoded {
101        parts.extensions.insert(ContentEncoded);
102    }
103    let accounted = AccountedBody::new(body, limiter, content_length, accounting);
104    let req = Request::from_parts(parts, Body::new(accounted));
105
106    let response = next.run(req).await;
107    // `retained_accounting` stays alive until this function returns, holding
108    // the permits across the handler's use of the collected body.
109    quota_exceeded_response(&retained_accounting, response)
110}
111
112/// Rewrites the response of a request whose body accounting hit the quota:
113/// the failure aborts the body mid-stream, so the extractors reject the
114/// request with a generic body error (mapped to 400). Quota exhaustion must
115/// surface as 429, matching the upfront admission path, so clients can tell
116/// backpressure apart from malformed input.
117fn quota_exceeded_response(accounting: &BodyMemoryAccounting, response: Response) -> Response {
118    // The flag can only be set when the body failed before the handler ran,
119    // so the response is the extractor rejection and can be replaced.
120    if accounting.take_quota_exceeded() {
121        return (
122            StatusCode::TOO_MANY_REQUESTS,
123            "Request body memory limit exceeded",
124        )
125            .into_response();
126    }
127    response
128}
129
130/// Route-local accounting for routes that install `RequestDecompressionLayer`:
131/// it must be layered *inside* the
132/// decompression layer, so the body it wraps is the **decompressed** stream.
133/// It requires the global accounting installed by [`crate::http::HttpServer`]
134/// and must use the same [`ServerMemoryLimiter`].
135///
136/// Only encoded requests marked by the global middleware are
137/// accounted; plain (uncompressed) requests are already covered by the global
138/// wire-byte accounting, and charging them here as well would double-count.
139/// For decompressed requests the `Content-Length` header has been stripped by
140/// the decompression layer, so the full decoded size is charged as it streams.
141pub async fn decoded_body_accounting_middleware(
142    State(limiter): State<ServerMemoryLimiter>,
143    req: Request,
144    next: Next,
145) -> Response {
146    if req.extensions().get::<ContentEncoded>().is_none() {
147        return next.run(req).await;
148    }
149
150    let content_length = req
151        .headers()
152        .get(http::header::CONTENT_LENGTH)
153        .and_then(|v| v.to_str().ok())
154        .and_then(|v| v.parse::<u64>().ok())
155        .unwrap_or(0);
156
157    let accounting = BodyMemoryAccounting::default();
158    // Retain the permits for the whole request, like the global middleware.
159    let retained_accounting = accounting.clone();
160    let (mut parts, body) = req.into_parts();
161    parts.extensions.insert(accounting.clone());
162    let accounted = AccountedBody::new(body, limiter, content_length, accounting);
163    let req = Request::from_parts(parts, Body::new(accounted));
164
165    let response = next.run(req).await;
166    // `retained_accounting` stays alive until this function returns, holding
167    // the permits across the handler's use of the collected body.
168    quota_exceeded_response(&retained_accounting, response)
169}
170
171/// Holds the memory guards acquired while a request body is streamed. Shared
172/// between the [`AccountedBody`] wrapper, the request extensions and the
173/// middleware itself, so the permits are only released when the request
174/// (including any collected body buffers and the handler still using them)
175/// is finished.
176#[derive(Clone, Default)]
177struct BodyMemoryAccounting {
178    guards: Arc<Mutex<Vec<MemoryGuard<RequestMemoryMetrics>>>>,
179    /// Set when an incremental charge hit the quota, so the rejection can be
180    /// rewritten into a 429 after the middleware regains control.
181    quota_exceeded: Arc<AtomicBool>,
182}
183
184impl BodyMemoryAccounting {
185    fn hold(&self, guard: MemoryGuard<RequestMemoryMetrics>) {
186        self.guards.lock().unwrap().push(guard);
187    }
188
189    fn mark_quota_exceeded(&self) {
190        self.quota_exceeded.store(true, Ordering::Release);
191    }
192
193    fn take_quota_exceeded(&self) -> bool {
194        self.quota_exceeded.swap(false, Ordering::AcqRel)
195    }
196}
197
198type AcquireFuture =
199    Pin<Box<dyn std::future::Future<Output = Result<MemoryGuard<RequestMemoryMetrics>>> + Send>>;
200
201enum ChargeOutcome {
202    Charged,
203    Pending,
204    Failed,
205}
206
207/// A request body wrapper that charges the shared limiter for every byte
208/// streamed beyond `pre_charged` (the upfront `Content-Length` reservation).
209struct AccountedBody {
210    inner: Body,
211    limiter: ServerMemoryLimiter,
212    accounting: BodyMemoryAccounting,
213    /// Bytes of body data streamed so far.
214    streamed: u64,
215    /// Bytes of the streamed body already covered by incremental permits.
216    incrementally_charged: u64,
217    /// Bytes covered by the upfront `Content-Length` reservation.
218    pre_charged: u64,
219    /// A data frame held back while its permit is being acquired, together
220    /// with the number of extra bytes to charge for it.
221    pending: Option<(Frame<Bytes>, u64)>,
222    pending_acquire: Option<AcquireFuture>,
223}
224
225impl AccountedBody {
226    fn new(
227        inner: Body,
228        limiter: ServerMemoryLimiter,
229        pre_charged: u64,
230        accounting: BodyMemoryAccounting,
231    ) -> Self {
232        Self {
233            inner,
234            limiter,
235            accounting,
236            streamed: 0,
237            incrementally_charged: 0,
238            pre_charged,
239            pending: None,
240            pending_acquire: None,
241        }
242    }
243
244    /// Bytes that still need an incremental permit for the body streamed so far.
245    fn uncharged(&self) -> u64 {
246        self.streamed
247            .saturating_sub(self.pre_charged)
248            .saturating_sub(self.incrementally_charged)
249    }
250
251    /// Acquires a permit for `bytes`; only one acquisition may be in flight.
252    ///
253    /// A parked acquisition always belongs to the currently buffered frame, so
254    /// it must be resolved *before* any new acquisition is tried: letting the
255    /// fast path succeed while a waiter is parked would leave the waiter alive,
256    /// and its late completion would then be credited with another frame's
257    /// byte count (under-reserving memory relative to what is marked charged).
258    fn charge(&mut self, bytes: u64, cx: &mut Context<'_>) -> ChargeOutcome {
259        debug_assert!(bytes > 0);
260        if let Some(fut) = self.pending_acquire.as_mut() {
261            return match fut.as_mut().poll(cx) {
262                Poll::Ready(Ok(guard)) => {
263                    self.accounting.hold(guard);
264                    self.pending_acquire = None;
265                    self.incrementally_charged += bytes;
266                    ChargeOutcome::Charged
267                }
268                Poll::Ready(Err(_)) => {
269                    self.pending_acquire = None;
270                    ChargeOutcome::Failed
271                }
272                Poll::Pending => ChargeOutcome::Pending,
273            };
274        }
275        // Fast path: the permit is immediately available.
276        if let Some(guard) = self.limiter.try_acquire(bytes) {
277            self.accounting.hold(guard);
278            self.incrementally_charged += bytes;
279            return ChargeOutcome::Charged;
280        }
281        let limiter = self.limiter.clone();
282        let mut fut = Box::pin(async move { limiter.acquire(bytes).await });
283        // Poll immediately so a ready result is not missed (returning
284        // `Pending` without registering progress would park the task forever).
285        match fut.as_mut().poll(cx) {
286            Poll::Ready(Ok(guard)) => {
287                self.accounting.hold(guard);
288                self.incrementally_charged += bytes;
289                ChargeOutcome::Charged
290            }
291            Poll::Ready(Err(_)) => ChargeOutcome::Failed,
292            Poll::Pending => {
293                self.pending_acquire = Some(fut);
294                ChargeOutcome::Pending
295            }
296        }
297    }
298
299    /// Hands out `frame` once `bytes` of it are charged, buffering the frame
300    /// while the acquisition is pending.
301    fn hand_out(
302        &mut self,
303        frame: Frame<Bytes>,
304        bytes: u64,
305        cx: &mut Context<'_>,
306    ) -> Poll<Option<Result<Frame<Bytes>, axum::Error>>> {
307        if bytes == 0 {
308            return Poll::Ready(Some(Ok(frame)));
309        }
310        match self.charge(bytes, cx) {
311            ChargeOutcome::Charged => Poll::Ready(Some(Ok(frame))),
312            ChargeOutcome::Pending => {
313                self.pending = Some((frame, bytes));
314                Poll::Pending
315            }
316            ChargeOutcome::Failed => {
317                self.accounting.mark_quota_exceeded();
318                Poll::Ready(Some(Err(limit_exceeded_error())))
319            }
320        }
321    }
322}
323
324impl HttpBody for AccountedBody {
325    type Data = Bytes;
326    type Error = axum::Error;
327
328    fn poll_frame(
329        self: Pin<&mut Self>,
330        cx: &mut Context<'_>,
331    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
332        let this = self.get_mut();
333
334        // Finish a pending acquisition for a buffered frame first.
335        if let Some((frame, bytes)) = this.pending.take() {
336            return this.hand_out(frame, bytes, cx);
337        }
338
339        match Pin::new(&mut this.inner).poll_frame(cx) {
340            Poll::Pending => Poll::Pending,
341            Poll::Ready(None) => Poll::Ready(None),
342            Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))),
343            Poll::Ready(Some(Ok(frame))) => {
344                if let Some(data) = frame.data_ref() {
345                    this.streamed += data.len() as u64;
346                    let uncharged = this.uncharged();
347                    this.hand_out(frame, uncharged, cx)
348                } else {
349                    // Trailers pass through uncharged.
350                    Poll::Ready(Some(Ok(frame)))
351                }
352            }
353        }
354    }
355
356    fn is_end_stream(&self) -> bool {
357        self.pending.is_none() && self.inner.is_end_stream()
358    }
359
360    fn size_hint(&self) -> http_body::SizeHint {
361        self.inner.size_hint()
362    }
363}
364
365fn limit_exceeded_error() -> axum::Error {
366    axum::Error::new(std::io::Error::other(
367        "request body exceeded the aggregate memory limit",
368    ))
369}
370
371#[cfg(test)]
372mod tests {
373    use axum::body::{Body, Bytes};
374    use axum::http::{Request, StatusCode, header};
375    use axum::routing::post;
376    use axum::{Router, middleware};
377    use common_memory_manager::OnExhaustedPolicy;
378    use tower::ServiceExt;
379
380    use super::memory_limit_middleware;
381    use crate::request_memory_limiter::ServerMemoryLimiter;
382
383    /// Streams `body` in `chunk_size` frames without a Content-Length header
384    /// (like HTTP/1.1 chunked transfer).
385    fn chunked_request(uri: &str, body: Vec<u8>, chunk_size: usize) -> Request<Body> {
386        let stream = futures_util::stream::iter(
387            body.chunks(chunk_size)
388                .map(|c| Ok::<_, std::io::Error>(Bytes::copy_from_slice(c)))
389                .collect::<Vec<_>>(),
390        );
391        Request::builder()
392            .uri(uri)
393            .method("POST")
394            .body(Body::from_stream(stream))
395            .unwrap()
396    }
397
398    fn counted_body_app(limiter: ServerMemoryLimiter) -> Router {
399        Router::new()
400            .route(
401                "/echo",
402                post(|body: Bytes| async move { (StatusCode::OK, body.len().to_string()) }),
403            )
404            .layer(middleware::from_fn_with_state(
405                limiter,
406                memory_limit_middleware,
407            ))
408    }
409
410    #[tokio::test]
411    async fn test_chunked_body_larger_than_quota_is_rejected() {
412        // The quota cannot cover a fully streamed 64 KiB chunked body.
413        let limiter = ServerMemoryLimiter::new(8 * 1024, OnExhaustedPolicy::Fail);
414        let app = counted_body_app(limiter.clone());
415
416        let res = app
417            .oneshot(chunked_request("/echo", vec![b'x'; 64 * 1024], 1024))
418            .await
419            .unwrap();
420        assert_eq!(
421            res.status(),
422            StatusCode::TOO_MANY_REQUESTS,
423            "a chunked body larger than the quota must be rejected with 429, not 400"
424        );
425        assert_eq!(0, limiter.used_bytes(), "guards must be released");
426    }
427
428    #[tokio::test]
429    async fn test_small_chunked_body_fits() {
430        let limiter = ServerMemoryLimiter::new(8 * 1024, OnExhaustedPolicy::Fail);
431        let app = counted_body_app(limiter.clone());
432
433        let res = app
434            .oneshot(chunked_request("/echo", vec![b'x'; 1024], 512))
435            .await
436            .unwrap();
437        assert_eq!(res.status(), StatusCode::OK);
438        let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap();
439        assert_eq!(&body[..], b"1024");
440        assert_eq!(0, limiter.used_bytes(), "guards must be released");
441    }
442
443    #[tokio::test]
444    async fn test_content_length_still_admitted_upfront() {
445        // A Content-Length larger than the quota is rejected before the body
446        // is read (429), unchanged from the previous behavior.
447        let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail);
448        let app = counted_body_app(limiter);
449
450        let req = Request::builder()
451            .uri("/echo")
452            .method("POST")
453            .header(header::CONTENT_LENGTH, "4096")
454            .body(Body::from(vec![b'x'; 4096]))
455            .unwrap();
456        let res = app.oneshot(req).await.unwrap();
457        assert_eq!(res.status(), StatusCode::TOO_MANY_REQUESTS);
458    }
459
460    #[tokio::test]
461    async fn test_understated_content_length_is_charged_the_difference() {
462        // Claim 1 KiB in Content-Length but stream 16 KiB: the extra 15 KiB
463        // must be charged while streaming and rejected by the small quota.
464        let limiter = ServerMemoryLimiter::new(2 * 1024, OnExhaustedPolicy::Fail);
465        let app = counted_body_app(limiter.clone());
466
467        let stream = futures_util::stream::iter(vec![Ok::<_, std::io::Error>(
468            Bytes::copy_from_slice(&vec![b'x'; 16 * 1024]),
469        )]);
470        let req = Request::builder()
471            .uri("/echo")
472            .method("POST")
473            .header(header::CONTENT_LENGTH, "1024")
474            .body(Body::from_stream(stream))
475            .unwrap();
476        let res = app.oneshot(req).await.unwrap();
477        assert_eq!(res.status(), StatusCode::TOO_MANY_REQUESTS);
478        assert_eq!(0, limiter.used_bytes());
479    }
480
481    #[tokio::test]
482    async fn test_honest_content_length_is_not_double_charged() {
483        // With a quota large enough for the whole body, an accurate
484        // Content-Length must admit the request without extra charges.
485        let limiter = ServerMemoryLimiter::new(8 * 1024, OnExhaustedPolicy::Fail);
486        let app = counted_body_app(limiter.clone());
487
488        let req = Request::builder()
489            .uri("/echo")
490            .method("POST")
491            .header(header::CONTENT_LENGTH, "4096")
492            .body(Body::from(vec![b'x'; 4096]))
493            .unwrap();
494        let res = app.oneshot(req).await.unwrap();
495        assert_eq!(res.status(), StatusCode::OK);
496        assert_eq!(0, limiter.used_bytes());
497    }
498
499    /// Mirrors the route-local stack: decoded-body accounting *inside* a
500    /// request-decompression layer, like the ingestion routes in `http.rs`.
501    fn decompressing_app(limiter: ServerMemoryLimiter) -> Router {
502        use tower_http::decompression::RequestDecompressionLayer;
503
504        Router::new()
505            .route(
506                "/echo",
507                post(|body: Bytes| async move { (StatusCode::OK, body.len().to_string()) }),
508            )
509            .layer(middleware::from_fn_with_state(
510                limiter,
511                super::decoded_body_accounting_middleware,
512            ))
513            .layer(RequestDecompressionLayer::new().pass_through_unaccepted(true))
514    }
515
516    fn zstd_body(decoded: &[u8]) -> Body {
517        let compressed = zstd::stream::encode_all(decoded, 3).unwrap();
518        Body::from(compressed)
519    }
520
521    #[tokio::test]
522    async fn test_decompressed_body_is_charged_for_content_encoded_requests() {
523        // The decoded size (64 KiB) exceeds the 16 KiB quota: the accounting
524        // layer must see the decompressed bytes and reject mid-stream.
525        let limiter = ServerMemoryLimiter::new(16 * 1024, OnExhaustedPolicy::Fail);
526        let app = decompressing_app(limiter.clone());
527
528        let req = Request::builder()
529            .uri("/echo")
530            .method("POST")
531            .header("content-encoding", "zstd")
532            .extension(super::ContentEncoded)
533            .body(zstd_body(&vec![b'x'; 64 * 1024]))
534            .unwrap();
535        let res = app.oneshot(req).await.unwrap();
536        assert_eq!(
537            res.status(),
538            StatusCode::TOO_MANY_REQUESTS,
539            "decoded body larger than the quota must be rejected with 429"
540        );
541        assert_eq!(0, limiter.used_bytes(), "guards must be released");
542    }
543
544    #[tokio::test]
545    async fn test_small_decompressed_body_fits() {
546        let limiter = ServerMemoryLimiter::new(16 * 1024, OnExhaustedPolicy::Fail);
547        let app = decompressing_app(limiter.clone());
548
549        let req = Request::builder()
550            .uri("/echo")
551            .method("POST")
552            .header("content-encoding", "zstd")
553            .extension(super::ContentEncoded)
554            .body(zstd_body(&vec![b'x'; 1024]))
555            .unwrap();
556        let res = app.oneshot(req).await.unwrap();
557        assert_eq!(res.status(), StatusCode::OK);
558        let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap();
559        assert_eq!(&body[..], b"1024", "the handler must see the decoded body");
560        assert_eq!(0, limiter.used_bytes(), "guards must be released");
561    }
562
563    #[tokio::test]
564    async fn test_pending_acquisition_is_attributed_to_its_own_bytes() {
565        use std::time::Duration;
566
567        use tokio_stream::wrappers::ReceiverStream;
568
569        // Regression test: a parked acquisition must be resolved before any
570        // new acquisition is tried; otherwise a late waiter completion can be
571        // credited with a later frame's byte count.
572        //
573        // Scenario: quota 8 KiB fully held externally; frame A (4 KiB) parks
574        // its acquisition; the external guard is then released so A drains;
575        // frame B (8 KiB) arrives. Correct behavior: B waits for its own 8 KiB
576        // acquisition and is rejected when it times out (only 4 KiB are free).
577        // Buggy behavior: the stale waiter (4 KiB) completes and is credited
578        // with B's 8 KiB, so B is delivered although the quota cannot cover it.
579        let limiter = ServerMemoryLimiter::new(
580            8 * 1024,
581            OnExhaustedPolicy::Wait {
582                timeout: Duration::from_millis(500),
583            },
584        );
585        let app = counted_body_app(limiter.clone());
586
587        let (tx, rx) = tokio::sync::mpsc::channel::<Result<Bytes, std::io::Error>>(2);
588        let driver_limiter = limiter.clone();
589        let driver = tokio::spawn(async move {
590            // Hold the whole quota so frame A's acquisition parks.
591            let external = driver_limiter.acquire(8 * 1024).await.unwrap();
592            tx.send(Ok(Bytes::from(vec![b'a'; 4096]))).await.unwrap();
593            tokio::time::sleep(Duration::from_millis(300)).await;
594            // Free the quota: frame A's waiter drains and A is delivered.
595            drop(external);
596            tokio::time::sleep(Duration::from_millis(300)).await;
597            tx.send(Ok(Bytes::from(vec![b'b'; 8192]))).await.unwrap();
598            // Keep the channel open; frame B's own acquisition must time out.
599            tokio::time::sleep(Duration::from_millis(1500)).await;
600        });
601
602        let res = app
603            .oneshot(
604                Request::builder()
605                    .uri("/echo")
606                    .method("POST")
607                    .body(Body::from_stream(ReceiverStream::new(rx)))
608                    .unwrap(),
609            )
610            .await
611            .unwrap();
612        driver.await.unwrap();
613
614        assert_eq!(
615            res.status(),
616            StatusCode::TOO_MANY_REQUESTS,
617            "frame B must wait for its own 8 KiB and time out, not borrow A's reservation"
618        );
619    }
620
621    #[tokio::test]
622    async fn test_plain_requests_are_not_charged_by_the_decoded_layer() {
623        // Without the ContentEncoded marker the decoded layer must be a
624        // no-op: plain bodies are accounted by the global wire accounting.
625        let limiter = ServerMemoryLimiter::new(16 * 1024, OnExhaustedPolicy::Fail);
626        let app = decompressing_app(limiter.clone());
627
628        // Plain body, no marker: passes through.
629        let req = Request::builder()
630            .uri("/echo")
631            .method("POST")
632            .body(Body::from(vec![b'x'; 1024]))
633            .unwrap();
634        let res = app.clone().oneshot(req).await.unwrap();
635        assert_eq!(res.status(), StatusCode::OK);
636
637        // Compressed body WITHOUT the marker (e.g. a route that has the
638        // decompression layer but no accounting): still passes through.
639        let req = Request::builder()
640            .uri("/echo")
641            .method("POST")
642            .header("content-encoding", "zstd")
643            .body(zstd_body(&vec![b'x'; 1024]))
644            .unwrap();
645        let res = app.oneshot(req).await.unwrap();
646        assert_eq!(res.status(), StatusCode::OK);
647    }
648}