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