Limit body size in trillium

main.rs

use futures_lite::{AsyncRead, AsyncReadExt};
use std::{
    io::{Error, ErrorKind, Result},
    pin::Pin,
    task::{ready, Context, Poll},
};
use trillium::{Conn, Status};

pub struct Limit<T> {
    reader: T,
    length: usize,
    max: usize,
}

impl<T> Limit<T> {
    pub fn new(reader: T, max: usize) -> Self {
        Self {
            reader,
            length: 0,
            max,
        }
    }
}

impl<T> AsyncRead for Limit<T>
where
    T: AsyncRead + Unpin,
{
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8],
    ) -> Poll<Result<usize>> {
        let new_bytes = ready!(Pin::new(&mut self.reader).poll_read(cx, buf))?;
        self.length += new_bytes;
        if self.length >= self.max {
            Poll::Ready(Err(Error::new(
                ErrorKind::InvalidData, // or if on nightly, ErrorKind::FileTooLarge
                "max length exceeded",
            )))
        } else {
            Poll::Ready(Ok(new_bytes))
        }
    }
}

async fn app(mut conn: Conn) -> Conn {
    let mut request_body = vec![];
    let result = Limit::new(conn.request_body().await, 30 * 1024 * 1024)
        .read_to_end(&mut request_body) // this could alternatively stream directly to disk
        .await;
    match result {
        Err(e) if e.kind() == ErrorKind::InvalidData => conn
            .with_status(Status::NotAcceptable)
            .with_body(e.to_string()),
        Err(e) => conn
            .with_status(Status::InternalServerError)
            .with_body(e.to_string()),
        Ok(bytes) => conn.ok(format!("read {bytes}")),
    }
}

fn main() {
    trillium_smol::run(app)
}

$ wc -c big-file
31457280 big-file

$ curl -v -T big-file -X POST http://localhost:8080
*   Trying 127.0.0.1:8080...
* connect to 127.0.0.1 port 8080 failed: Connection refused
*   Trying [::1]:8080...
* Connected to localhost (::1) port 8080 (#0)
> POST /big-file HTTP/1.1
> Host: localhost:8080
> User-Agent: curl/8.1.2
> Accept: */*
> Content-Length: 51457280
> Expect: 100-continue
>
< HTTP/1.1 100 Continue
< HTTP/1.1 406 Not Acceptable
< Date: Wed, 02 Aug 2023 23:50:49 GMT
< Server: trillium/0.3.2
< Content-Length: 19
< Connection: keep-alive
* HTTP error before end of send, stop sending
<
* Closing connection 0
max length exceeded

$ wc -c small-file
31457279 small-file

$ curl -v -T small-file -X POST http://localhost:8080
*   Trying 127.0.0.1:8080...
* connect to 127.0.0.1 port 8080 failed: Connection refused
*   Trying [::1]:8080...
* Connected to localhost (::1) port 8080 (#0)
> POST /small-file HTTP/1.1
> Host: localhost:8080
> User-Agent: curl/8.1.2
> Accept: */*
> Content-Length: 31457279
> Expect: 100-continue
>
< HTTP/1.1 100 Continue
* We are completely uploaded and fine
< HTTP/1.1 200 OK
< Date: Thu, 03 Aug 2023 00:01:39 GMT
< Server: trillium/0.3.2
< Content-Length: 13
< Connection: keep-alive
<
* Connection #0 to host localhost left intact
read 31457279