1use 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#[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 let accounting = BodyMemoryAccounting::default();
90 let retained_accounting = accounting.clone();
95 let (mut parts, body) = req.into_parts();
96 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 quota_exceeded_response(&retained_accounting, response)
110}
111
112fn quota_exceeded_response(accounting: &BodyMemoryAccounting, response: Response) -> Response {
118 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
130pub 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 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 quota_exceeded_response(&retained_accounting, response)
169}
170
171#[derive(Clone, Default)]
177struct BodyMemoryAccounting {
178 guards: Arc<Mutex<Vec<MemoryGuard<RequestMemoryMetrics>>>>,
179 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
207struct AccountedBody {
210 inner: Body,
211 limiter: ServerMemoryLimiter,
212 accounting: BodyMemoryAccounting,
213 streamed: u64,
215 incrementally_charged: u64,
217 pre_charged: u64,
219 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 fn uncharged(&self) -> u64 {
246 self.streamed
247 .saturating_sub(self.pre_charged)
248 .saturating_sub(self.incrementally_charged)
249 }
250
251 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 drop(external);
596 tokio::time::sleep(Duration::from_millis(300)).await;
597 tx.send(Ok(Bytes::from(vec![b'b'; 8192]))).await.unwrap();
598 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 let limiter = ServerMemoryLimiter::new(16 * 1024, OnExhaustedPolicy::Fail);
626 let app = decompressing_app(limiter.clone());
627
628 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 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}