rocket/response/body.rs
1use std::{io, fmt};
2use std::task::{Context, Poll};
3use std::pin::Pin;
4
5use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt, ReadBuf};
6
7/// The body of a [`Response`].
8///
9/// A `Body` is never created directly, but instead, through the following
10/// methods on `Response` and `Builder`:
11///
12/// * [`Builder::sized_body()`]
13/// * [`Response::set_sized_body()`]
14/// * [`Builder::streamed_body()`]
15/// * [`Response::set_streamed_body()`]
16///
17/// [`Response`]: crate::Response
18/// [`Builder`]: crate::response::Builder
19/// [`Response::set_sized_body()`]: crate::Response::set_sized_body
20/// [`Response::set_streamed_body()`]: crate::Response::set_streamed_body
21/// [`Builder::sized_body()`]: crate::response::Builder::sized_body
22/// [`Builder::streamed_body()`]: crate::response::Builder::streamed_body
23///
24/// An unset body in a `Response` begins as a [`Body::default()`], a `None`
25/// body with a preset size of `0`.
26///
27/// # Sizing
28///
29/// A response body may be sized or unsized ("streamed"). A "sized" body is
30/// transferred with a `Content-Length` equal to its size while an "unsized"
31/// body is chunk-encoded. The body data is streamed in _all_ cases and is never
32/// buffered in memory beyond a minimal amount for efficient transmission.
33///
34/// ## Sized
35///
36/// A sized body may have a _preset_ size ([`Body::preset_size()`]) or may have
37/// its size computed on the fly by seeking ([`Body::size()`]). As such, sized
38/// bodies must implement [`AsyncSeek`]. If a body does not have a preset size
39/// and the fails to be computed dynamically, a sized body is treated as an
40/// unsized body when written out to the network.
41///
42/// ## Unsized
43///
44/// An unsized body's data is streamed as it arrives. In other words, as soon as
45/// the body's [`AsyncRead`] implementation returns bytes, the bytes are written
46/// to the network. Individual unsized bodies may use an internal buffer to
47/// curtail writes to the network.
48///
49/// The maximum number of bytes written to the network at once is controlled via
50/// the [`Body::max_chunk_size()`] parameter which can be set via
51/// [`Response::set_max_chunk_size()`] and [`Builder::max_chunk_size()`].
52///
53/// [`Response::set_max_chunk_size()`]: crate::Response::set_max_chunk_size
54/// [`Builder::max_chunk_size()`]: crate::response::Builder::max_chunk_size
55///
56/// # Reading
57///
58/// The contents of a body, decoded, can be read through [`Body::to_bytes()`],
59/// [`Body::to_string()`], or directly though `Body`'s [`AsyncRead`]
60/// implementation.
61#[derive(Debug)]
62pub struct Body<'r> {
63 /// The size of the body, if it is known.
64 size: Option<usize>,
65 /// The body itself.
66 inner: Inner<'r>,
67 /// The maximum chunk size.
68 max_chunk: usize,
69}
70
71/// A "trait alias" of sorts so we can use `AsyncRead + AsyncSeek` in `dyn`.
72pub trait AsyncReadSeek: AsyncRead + AsyncSeek { }
73
74/// Implemented for all `AsyncRead + AsyncSeek`, of course.
75impl<T: AsyncRead + AsyncSeek> AsyncReadSeek for T { }
76
77/// A pinned `AsyncRead + AsyncSeek` body type.
78type SizedBody<'r> = Pin<Box<dyn AsyncReadSeek + Send + 'r>>;
79
80/// A pinned `AsyncRead` (not `AsyncSeek`) body type.
81type UnsizedBody<'r> = Pin<Box<dyn AsyncRead + Send + 'r>>;
82
83enum Inner<'r> {
84 /// A body that can be seeked to determine it's size.
85 Seekable(SizedBody<'r>),
86 /// A body that has no known size.
87 Unsized(UnsizedBody<'r>),
88 /// A body that "exists" but only for metadata calculations.
89 Phantom(SizedBody<'r>),
90 /// An empty body: no body at all.
91 None,
92}
93
94impl Default for Body<'_> {
95 fn default() -> Self {
96 Body {
97 size: Some(0),
98 inner: Inner::None,
99 max_chunk: Body::DEFAULT_MAX_CHUNK,
100 }
101 }
102}
103
104impl<'r> Body<'r> {
105 /// The default max size, in bytes, of chunks for streamed responses.
106 ///
107 /// The present value is `4096`.
108 pub const DEFAULT_MAX_CHUNK: usize = 4096;
109
110 pub(crate) fn with_sized<T>(body: T, preset_size: Option<usize>) -> Self
111 where T: AsyncReadSeek + Send + 'r
112 {
113 Body {
114 size: preset_size,
115 inner: Inner::Seekable(Box::pin(body)),
116 max_chunk: Body::DEFAULT_MAX_CHUNK,
117 }
118 }
119
120 pub(crate) fn with_unsized<T>(body: T) -> Self
121 where T: AsyncRead + Send + 'r
122 {
123 Body {
124 size: None,
125 inner: Inner::Unsized(Box::pin(body)),
126 max_chunk: Body::DEFAULT_MAX_CHUNK,
127 }
128 }
129
130 pub(crate) fn set_max_chunk_size(&mut self, max_chunk: usize) {
131 self.max_chunk = max_chunk;
132 }
133
134 pub(crate) fn strip(&mut self) {
135 let body = std::mem::take(self);
136 *self = match body.inner {
137 Inner::Seekable(b) | Inner::Phantom(b) => Body {
138 size: body.size,
139 inner: Inner::Phantom(b),
140 max_chunk: body.max_chunk,
141 },
142 Inner::Unsized(_) | Inner::None => Body::default()
143 };
144 }
145
146 /// Returns `true` if the body is `None` or unset, the default.
147 ///
148 /// # Example
149 ///
150 /// ```rust
151 /// use rocket::response::Response;
152 ///
153 /// let r = Response::build().finalize();
154 /// assert!(r.body().is_none());
155 /// ```
156 #[inline(always)]
157 pub fn is_none(&self) -> bool {
158 matches!(self.inner, Inner::None)
159 }
160
161 /// Returns `true` if the body is _not_ `None`, anything other than the
162 /// default.
163 ///
164 /// # Example
165 ///
166 /// ```rust
167 /// use std::io::Cursor;
168 /// use rocket::response::Response;
169 ///
170 /// let body = "Brewing the best coffee!";
171 /// let r = Response::build()
172 /// .sized_body(body.len(), Cursor::new(body))
173 /// .finalize();
174 ///
175 /// assert!(r.body().is_some());
176 /// ```
177 #[inline(always)]
178 pub fn is_some(&self) -> bool {
179 !self.is_none()
180 }
181
182 /// A body's preset size, which may have been computed by a previous call to
183 /// [`Body::size()`].
184 ///
185 /// Unsized bodies _always_ return `None`, while sized bodies return `Some`
186 /// if the body size was supplied directly on creation or a call to
187 /// [`Body::size()`] successfully computed the size and `None` otherwise.
188 ///
189 /// # Example
190 ///
191 /// ```rust
192 /// use std::io::Cursor;
193 /// use rocket::response::Response;
194 ///
195 /// # rocket::async_test(async {
196 /// let body = "Brewing the best coffee!";
197 /// let r = Response::build()
198 /// .sized_body(body.len(), Cursor::new(body))
199 /// .finalize();
200 ///
201 /// // This will _always_ return `Some`.
202 /// assert_eq!(r.body().preset_size(), Some(body.len()));
203 ///
204 /// let r = Response::build()
205 /// .streamed_body(Cursor::new(body))
206 /// .finalize();
207 ///
208 /// // This will _never_ return `Some`.
209 /// assert_eq!(r.body().preset_size(), None);
210 ///
211 /// let mut r = Response::build()
212 /// .sized_body(None, Cursor::new(body))
213 /// .finalize();
214 ///
215 /// // This returns `Some` only after a call to `size()`.
216 /// assert_eq!(r.body().preset_size(), None);
217 /// assert_eq!(r.body_mut().size().await, Some(body.len()));
218 /// assert_eq!(r.body().preset_size(), Some(body.len()));
219 /// # });
220 /// ```
221 pub fn preset_size(&self) -> Option<usize> {
222 self.size
223 }
224
225 /// Returns the maximum chunk size for chunked transfers.
226 ///
227 /// If none is explicitly set, defaults to [`Body::DEFAULT_MAX_CHUNK`].
228 ///
229 /// # Example
230 ///
231 /// ```rust
232 /// use std::io::Cursor;
233 /// use rocket::response::{Response, Body};
234 ///
235 /// let body = "Brewing the best coffee!";
236 /// let r = Response::build()
237 /// .sized_body(body.len(), Cursor::new(body))
238 /// .finalize();
239 ///
240 /// assert_eq!(r.body().max_chunk_size(), Body::DEFAULT_MAX_CHUNK);
241 ///
242 /// let r = Response::build()
243 /// .sized_body(body.len(), Cursor::new(body))
244 /// .max_chunk_size(1024)
245 /// .finalize();
246 ///
247 /// assert_eq!(r.body().max_chunk_size(), 1024);
248 /// ```
249 pub fn max_chunk_size(&self) -> usize {
250 self.max_chunk
251 }
252
253 /// Attempts to compute the body's size and returns it if the body is sized.
254 ///
255 /// If the size was preset (see [`Body::preset_size()`]), the value is
256 /// returned immediately as `Some`. If the body is unsized or computing the
257 /// size fails, returns `None`. Otherwise, the size is computed by seeking,
258 /// and the `preset_size` is updated to reflect the known value.
259 ///
260 /// **Note:** the number of bytes read from the reader and/or written to the
261 /// network may differ from the value returned by this method. Some examples
262 /// include:
263 ///
264 /// * bodies in response to `HEAD` requests are never read or written
265 /// * the client may close the connection before the body is read fully
266 /// * reading the body may fail midway
267 /// * a preset size may differ from the actual body size
268 ///
269 /// # Example
270 ///
271 /// ```rust
272 /// use std::io::Cursor;
273 /// use rocket::response::Response;
274 ///
275 /// # rocket::async_test(async {
276 /// let body = "Hello, Rocketeers!";
277 /// let mut r = Response::build()
278 /// .sized_body(None, Cursor::new(body))
279 /// .finalize();
280 ///
281 /// assert_eq!(r.body().preset_size(), None);
282 /// assert_eq!(r.body_mut().size().await, Some(body.len()));
283 /// assert_eq!(r.body().preset_size(), Some(body.len()));
284 /// # });
285 /// ```
286 pub async fn size(&mut self) -> Option<usize> {
287 if let Some(size) = self.size {
288 return Some(size);
289 }
290
291 if let Inner::Seekable(ref mut body) | Inner::Phantom(ref mut body) = self.inner {
292 let pos = body.seek(io::SeekFrom::Current(0)).await.ok()?;
293 let end = body.seek(io::SeekFrom::End(0)).await.ok()?;
294 body.seek(io::SeekFrom::Start(pos)).await.ok()?;
295
296 let size = end as usize - pos as usize;
297 self.size = Some(size);
298 return Some(size);
299 }
300
301 None
302 }
303
304 /// Moves the body out of `self` and returns it, leaving a
305 /// [`Body::default()`] in its place.
306 ///
307 /// # Example
308 ///
309 /// ```rust
310 /// use std::io::Cursor;
311 /// use rocket::response::Response;
312 ///
313 /// let mut r = Response::build()
314 /// .sized_body(None, Cursor::new("Hi"))
315 /// .finalize();
316 ///
317 /// assert!(r.body().is_some());
318 ///
319 /// let body = r.body_mut().take();
320 /// assert!(body.is_some());
321 /// assert!(r.body().is_none());
322 /// ```
323 #[inline(always)]
324 pub fn take(&mut self) -> Self {
325 std::mem::take(self)
326 }
327
328 /// Reads all of `self` into a vector of bytes, consuming the contents.
329 ///
330 /// If reading fails, returns `Err`. Otherwise, returns `Ok`. Calling this
331 /// method may partially or fully consume the body's content. As such,
332 /// subsequent calls to `to_bytes()` will likely return different result.
333 ///
334 /// # Example
335 ///
336 /// ```rust
337 /// use std::io;
338 /// use rocket::response::Response;
339 ///
340 /// # let ok: io::Result<()> = rocket::async_test(async {
341 /// let mut r = Response::build()
342 /// .streamed_body(io::Cursor::new(&[1, 2, 3, 11, 13, 17]))
343 /// .finalize();
344 ///
345 /// assert_eq!(r.body_mut().to_bytes().await?, &[1, 2, 3, 11, 13, 17]);
346 /// # Ok(())
347 /// # });
348 /// # assert!(ok.is_ok());
349 /// ```
350 pub async fn to_bytes(&mut self) -> io::Result<Vec<u8>> {
351 let mut vec = Vec::new();
352 let n = match self.read_to_end(&mut vec).await {
353 Ok(n) => n,
354 Err(e) => {
355 error_!("Error reading body: {:?}", e);
356 return Err(e);
357 }
358 };
359
360 if let Some(ref mut size) = self.size {
361 *size = size.checked_sub(n).unwrap_or(0);
362 }
363
364 Ok(vec)
365 }
366
367 /// Reads all of `self` into a string, consuming the contents.
368 ///
369 /// If reading fails, or the body contains invalid UTF-8 characters, returns
370 /// `Err`. Otherwise, returns `Ok`. Calling this method may partially or
371 /// fully consume the body's content. As such, subsequent calls to
372 /// `to_string()` will likely return different result.
373 ///
374 /// # Example
375 ///
376 /// ```rust
377 /// use std::io;
378 /// use rocket::response::Response;
379 ///
380 /// # let ok: io::Result<()> = rocket::async_test(async {
381 /// let mut r = Response::build()
382 /// .streamed_body(io::Cursor::new("Hello, Rocketeers!"))
383 /// .finalize();
384 ///
385 /// assert_eq!(r.body_mut().to_string().await?, "Hello, Rocketeers!");
386 /// # Ok(())
387 /// # });
388 /// # assert!(ok.is_ok());
389 /// ```
390 pub async fn to_string(&mut self) -> io::Result<String> {
391 String::from_utf8(self.to_bytes().await?).map_err(|e| {
392 error_!("Body is invalid UTF-8: {}", e);
393 io::Error::new(io::ErrorKind::InvalidData, e)
394 })
395 }
396}
397
398impl AsyncRead for Body<'_> {
399 fn poll_read(
400 mut self: Pin<&mut Self>,
401 cx: &mut Context<'_>,
402 buf: &mut ReadBuf<'_>,
403 ) -> Poll<io::Result<()>> {
404 let reader = match self.inner {
405 Inner::Seekable(ref mut b) => b as &mut (dyn AsyncRead + Unpin),
406 Inner::Unsized(ref mut b) => b as &mut (dyn AsyncRead + Unpin),
407 Inner::Phantom(_) | Inner::None => return Poll::Ready(Ok(())),
408 };
409
410 Pin::new(reader).poll_read(cx, buf)
411 }
412}
413
414impl fmt::Debug for Inner<'_> {
415 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
416 match self {
417 Inner::Seekable(_) => "seekable".fmt(f),
418 Inner::Unsized(_) => "unsized".fmt(f),
419 Inner::Phantom(_) => "phantom".fmt(f),
420 Inner::None => "none".fmt(f),
421 }
422 }
423}