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}