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    remaining: Option<u64>,
68}
69
70impl FileBody {
71    pub(crate) fn new(file: File) -> Self {
72        FileBody {
73            file,
74            buf: BytesMut::new(),
75            remaining: None,
76        }
77    }
78
79    pub(crate) fn new_ranged(file: File, len: u64) -> Self {
80        FileBody {
81            file,
82            buf: BytesMut::new(),
83            remaining: Some(len),
84        }
85    }
86}
87
88impl Body for FileBody {
89    type Data = Bytes;
90    type Error = StaticError;
91
92    fn poll_frame(
93        self: Pin<&mut Self>,
94        cx: &mut Context<'_>,
95    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
96        let this = self.get_mut();
97
98        if let Some(remaining) = this.remaining {
99            if remaining == 0 {
100                return Poll::Ready(None);
101            }
102            let chunk_size = FILE_CHUNK_SIZE.min(remaining as usize);
103            if this.buf.capacity() - this.buf.len() < chunk_size {
104                this.buf.reserve(chunk_size);
105            }
106        } else {
107            if this.buf.capacity() - this.buf.len() < FILE_CHUNK_SIZE {
108                this.buf.reserve(FILE_CHUNK_SIZE);
109            }
110        }
111
112        let mut read_buf = ReadBuf::uninit(this.buf.spare_capacity_mut());
113        let file = Pin::new(&mut this.file);
114
115        match file.poll_read(cx, &mut read_buf) {
116            Poll::Ready(Ok(())) => {
117                let n = read_buf.filled().len();
118                if n == 0 {
119                    return Poll::Ready(None);
120                }
121
122                if let Some(ref mut remaining) = this.remaining {
123                    *remaining = remaining.saturating_sub(n as u64);
124                }
125
126                // Safety: `poll_read` reported exactly `n` bytes filled into the spare
127                // capacity we handed it via `ReadBuf::uninit`; advancing by that same
128                // `n` only marks bytes the reader actually initialized.
129                unsafe { this.buf.advance_mut(n) };
130                let chunk = this.buf.split_to(n).freeze();
131                Poll::Ready(Some(Ok(Frame::data(chunk))))
132            }
133            Poll::Ready(Err(e)) => Poll::Ready(Some(Err(StaticError::Io(e)))),
134            Poll::Pending => Poll::Pending,
135        }
136    }
137}