body 大小限制

impl<ReqBody, ResBody, S> Service<Request<ReqBody>> for RequestBodyLimit<S>
where
    ResBody: Body,
    S: Service<Request<Limited<ReqBody>>, Response = Response<ResBody>>,
{
    type Response = Response<ResponseBody<ResBody>>;

    enum ResponseBodyInner<B> {
        PayloadTooLarge {
            #[pin]
            body: Full<Bytes>,
        },
        Body {
            #[pin]
            body: B
        }
    }

tower-http = { version = "=0.7.1", features = ["limit", "decompression-gzip"] }
tower = { version = "=0.5.3", features = ["util"] }
tokio = { version = "=1.54.1", features = ["macros", "rt"] }
http = "1"
http-body = "1"
http-body-util = "0.1"
bytes = "1"
futures = "0.3"
flate2 = "1"

use bytes::Bytes;
use futures::stream;
use http::{header, Request, Response};
use http_body::Frame;
use http_body_util::{combinators::UnsyncBoxBody, BodyExt, Full, LengthLimitError, StreamBody};
use std::{convert::Infallible, io::Write};
use tower::{BoxError, ServiceBuilder, ServiceExt};
use tower_http::{decompression::RequestDecompressionLayer, limit::RequestBodyLimitLayer};

type Body = UnsyncBoxBody<Bytes, Infallible>;

// a streaming body of `len` bytes in 1 KiB frames, no size hint
fn stream_body(len: usize) -> Body {
  let frames: Vec<Result<Frame<Bytes>, Infallible>> =
    vec![b'a'; len].chunks(1024).map(|c| Ok(Frame::data(Bytes::copy_from_slice(c)))).collect();
  StreamBody::new(stream::iter(frames)).boxed_unsync()
}

// reads the whole body and reports what it saw
async fn handler<B>(req: Request<B>) -> Result<Response<Full<Bytes>>, Infallible>
where
  B: http_body::Body,
  B::Error: Into<BoxError>,
{
  let msg = match req.into_body().collect().await {
    Ok(c) => format!("read {} bytes", c.to_bytes().len()),
    Err(e) => {
      let e: BoxError = e.into();
      format!("error {:?}, LengthLimitError? {}", e.to_string(), e.is::<LengthLimitError>())
    }
  };
  Ok(Response::new(Full::from(msg)))
}

async fn send(name: &str, content_length: Option<&str>, body: Body) {
  let svc = ServiceBuilder::new().layer(RequestBodyLimitLayer::new(4096)).service_fn(handler);
  let mut req = Request::post("/");
  if let Some(cl) = content_length {
    req = req.header(header::CONTENT_LENGTH, cl);
  }
  let res = svc.oneshot(req.body(body).unwrap()).await.unwrap();
  let status = res.status();
  let ct = res.headers().get(header::CONTENT_TYPE).map(|v| v.to_str().unwrap().to_owned());
  let body = res.into_body().collect().await.unwrap().to_bytes();
  println!("{name:<32} {} {ct:?} {body:?}", status.as_u16());
}

fn gzip(data: &[u8]) -> Vec<u8> {
  let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::best());
  e.write_all(data).unwrap();
  e.finish().unwrap()
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  println!("limit 4096");
  send("Content-Length: 5000", Some("5000"), stream_body(5000)).await;
  send("Content-Length: abc, 5000 bytes", Some("abc"), stream_body(5000)).await;
  send("Content-Length: 10, 100 bytes", Some("10"), stream_body(100)).await;
  send("no Content-Length, 3000 bytes", None, stream_body(3000)).await;
  send("no Content-Length, 5000 bytes", None, stream_body(5000)).await;

  let bomb = Bytes::from(gzip(&vec![0u8; 10 << 20]));
  println!("gzip of 10 MiB zeros: {} bytes, limit 1 MiB", bomb.len());
  let req = || {
    Request::post("/")
      .header(header::CONTENT_ENCODING, "gzip")
      .header(header::CONTENT_LENGTH, bomb.len())
      .body(Full::new(bomb.clone()))
      .unwrap()
  };
  let outside = ServiceBuilder::new()
    .layer(RequestBodyLimitLayer::new(1 << 20))
    .layer(RequestDecompressionLayer::new())
    .service_fn(handler);
  let res = outside.oneshot(req()).await.unwrap();
  println!("limit outside decompression: {}", String::from_utf8_lossy(&res.into_body().collect().await.unwrap().to_bytes()));
  let inside = ServiceBuilder::new()
    .layer(RequestDecompressionLayer::new())
    .layer(RequestBodyLimitLayer::new(1 << 20))
    .service_fn(handler);
  let res = inside.oneshot(req()).await.unwrap();
  println!("limit inside decompression:  {}", String::from_utf8_lossy(&res.into_body().collect().await.unwrap().to_bytes()));
}

limit 4096
Content-Length: 5000             413 Some("text/plain; charset=utf-8") b"length limit exceeded"
Content-Length: abc, 5000 bytes  200 None b"error \"length limit exceeded\", LengthLimitError? true"
Content-Length: 10, 100 bytes    200 None b"error \"length limit exceeded\", LengthLimitError? true"
no Content-Length, 3000 bytes    200 None b"read 3000 bytes"
no Content-Length, 5000 bytes    200 None b"error \"length limit exceeded\", LengthLimitError? true"
gzip of 10 MiB zeros: 10208 bytes, limit 1 MiB
limit outside decompression: read 10485760 bytes
limit inside decompression:  error "length limit exceeded", LengthLimitError? true