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
13pub(crate) const FILE_CHUNK_SIZE: usize = 65_536;
15
16pub enum ResponseBody {
26 Buffered(Full<Bytes>),
28 Streamed(FileBody),
30 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 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
56pub 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 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}