Skip to main content

mini_static/
handler.rs

1use std::pin::Pin;
2use std::task::{Context, Poll};
3
4use bytes::{BufMut, Bytes, BytesMut};
5use http_body::{Body, Frame};
6use http_body_util::Full;
7use tokio::fs::File;
8use tokio::io::{AsyncRead, ReadBuf};
9
10use crate::error::StaticError;
11
12/// Chunk size for streaming file reads — 64 KB per frame.
13pub(crate) const FILE_CHUNK_SIZE: usize = 65_536;
14
15/// The response body type used by mini-static.
16///
17/// This is a concrete sum type, not a boxed trait object: `Buffered` covers redirects,
18/// errors, HEAD responses, and 304s; `Streamed` covers file `GET` responses. Keeping it
19/// concrete (rather than erasing into `BoxBody` here) lets an embedding crate — one that
20/// needs to remap the error type before erasing into its own body type, e.g.
21/// `mini-unified` bridging into `mini-serve`'s `ResponseBody` — erase exactly once at
22/// that boundary instead of erasing here and then again there.
23pub enum ResponseBody {
24    /// A body already fully in memory.
25    Buffered(Full<Bytes>),
26    /// A body streamed from disk one chunk at a time.
27    Streamed(FileBody),
28}
29
30impl Body for ResponseBody {
31    type Data = Bytes;
32    type Error = StaticError;
33
34    fn poll_frame(
35        self: Pin<&mut Self>,
36        cx: &mut Context<'_>,
37    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
38        match self.get_mut() {
39            ResponseBody::Buffered(body) => match Pin::new(body).poll_frame(cx) {
40                Poll::Ready(Some(Ok(frame))) => Poll::Ready(Some(Ok(frame))),
41                // `Full<Bytes>`'s error type is `Infallible` — this arm can never run.
42                Poll::Ready(Some(Err(never))) => match never {},
43                Poll::Ready(None) => Poll::Ready(None),
44                Poll::Pending => Poll::Pending,
45            },
46            ResponseBody::Streamed(body) => Pin::new(body).poll_frame(cx),
47        }
48    }
49}
50
51/// An `http_body::Body` that streams a `tokio::fs::File` to the client one chunk at a
52/// time, instead of buffering the whole file before the response body is polled.
53///
54/// Each `poll_frame` call reads directly into `buf`'s spare (uninitialized) capacity via
55/// `ReadBuf::uninit` and marks only the bytes the read syscall actually wrote as
56/// initialized via `advance_mut` — there's no `resize`-driven zero-fill and no extra
57/// copy: `split_to(n).freeze()` hands the just-filled bytes to the caller and leaves
58/// `buf`'s already-reserved spare capacity in place for the next read.
59pub struct FileBody {
60    file: File,
61    buf: BytesMut,
62}
63
64impl FileBody {
65    pub(crate) fn new(file: File) -> Self {
66        FileBody {
67            file,
68            buf: BytesMut::new(),
69        }
70    }
71}
72
73impl Body for FileBody {
74    type Data = Bytes;
75    type Error = StaticError;
76
77    fn poll_frame(
78        self: Pin<&mut Self>,
79        cx: &mut Context<'_>,
80    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
81        let this = self.get_mut();
82
83        if this.buf.capacity() - this.buf.len() < FILE_CHUNK_SIZE {
84            this.buf.reserve(FILE_CHUNK_SIZE);
85        }
86
87        let mut read_buf = ReadBuf::uninit(this.buf.spare_capacity_mut());
88        let file = Pin::new(&mut this.file);
89
90        match file.poll_read(cx, &mut read_buf) {
91            Poll::Ready(Ok(())) => {
92                let n = read_buf.filled().len();
93                if n == 0 {
94                    return Poll::Ready(None);
95                }
96                // Safety: `poll_read` reported exactly `n` bytes filled into the spare
97                // capacity we handed it via `ReadBuf::uninit`; advancing by that same
98                // `n` only marks bytes the reader actually initialized.
99                unsafe { this.buf.advance_mut(n) };
100                let chunk = this.buf.split_to(n).freeze();
101                Poll::Ready(Some(Ok(Frame::data(chunk))))
102            }
103            Poll::Ready(Err(e)) => Poll::Ready(Some(Err(StaticError::Io(e)))),
104            Poll::Pending => Poll::Pending,
105        }
106    }
107}