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}