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

servers/http/
memory_limit.rs
