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 `seek()`ed to determine its 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 unsized_none() -> Self {
111 Body {
112 size: None,
113 inner: Inner::None,
114 max_chunk: Body::DEFAULT_MAX_CHUNK,
115 }
116 }
117
118 pub(crate) fn with_sized<T>(body: T, preset_size: Option<usize>) -> Self
119 where T: AsyncReadSeek + Send + 'r
120 {
121 Body {
122 size: preset_size,
123 inner: Inner::Seekable(Box::pin(body)),
124 max_chunk: Body::DEFAULT_MAX_CHUNK,
125 }
126 }
127
128 pub(crate) fn with_unsized<T>(body: T) -> Self
129 where T: AsyncRead + Send + 'r
130 {
131 Body {
132 size: None,
133 inner: Inner::Unsized(Box::pin(body)),
134 max_chunk: Body::DEFAULT_MAX_CHUNK,
135 }
136 }
137
138 pub(crate) fn set_max_chunk_size(&mut self, max_chunk: usize) {
139 self.max_chunk = max_chunk;
140 }
141
142 pub(crate) fn strip(&mut self) {
143 let body = std::mem::take(self);
144 *self = match body.inner {
145 Inner::Seekable(b) | Inner::Phantom(b) => Body {
146 size: body.size,
147 inner: Inner::Phantom(b),
148 max_chunk: body.max_chunk,
149 },
150 Inner::Unsized(_) | Inner::None => Body::default()
151 };
152 }
153
154 /// Returns `true` if the body is `None` or unset, the default.
155 ///
156 /// # Example
157 ///
158 /// ```rust
159 /// use rocket::response::Response;
160 ///
161 /// let r = Response::build().finalize();
162 /// assert!(r.body().is_none());
163 /// ```
164 #[inline(always)]
165 pub fn is_none(&self) -> bool {
166 matches!(self.inner, Inner::None)
167 }
168
169 /// Returns `true` if the body is _not_ `None`, anything other than the
170 /// default.
171 ///
172 /// # Example
173 ///
174 /// ```rust
175 /// use std::io::Cursor;
176 /// use rocket::response::Response;
177 ///
178 /// let body = "Brewing the best coffee!";
179 /// let r = Response::build()
180 /// .sized_body(body.len(), Cursor::new(body))
181 /// .finalize();
182 ///
183 /// assert!(r.body().is_some());
184 /// ```
185 #[inline(always)]
186 pub fn is_some(&self) -> bool {
187 !self.is_none()
188 }
189
190 /// A body's preset size, which may have been computed by a previous call to
191 /// [`Body::size()`].
192 ///
193 /// Unsized bodies _always_ return `None`, while sized bodies return `Some`
194 /// if the body size was supplied directly on creation or a call to
195 /// [`Body::size()`] successfully computed the size and `None` otherwise.
196 ///
197 /// # Example
198 ///
199 /// ```rust
200 /// use std::io::Cursor;
201 /// use rocket::response::Response;
202 ///
203 /// # rocket::async_test(async {
204 /// let body = "Brewing the best coffee!";
205 /// let r = Response::build()
206 /// .sized_body(body.len(), Cursor::new(body))
207 /// .finalize();
208 ///
209 /// // This will _always_ return `Some`.
210 /// assert_eq!(r.body().preset_size(), Some(body.len()));
211 ///
212 /// let r = Response::build()
213 /// .streamed_body(Cursor::new(body))
214 /// .finalize();
215 ///
216 /// // This will _never_ return `Some`.
217 /// assert_eq!(r.body().preset_size(), None);
218 ///
219 /// let mut r = Response::build()
220 /// .sized_body(None, Cursor::new(body))
221 /// .finalize();
222 ///
223 /// // This returns `Some` only after a call to `size()`.
224 /// assert_eq!(r.body().preset_size(), None);
225 /// assert_eq!(r.body_mut().size().await, Some(body.len()));
226 /// assert_eq!(r.body().preset_size(), Some(body.len()));
227 /// # });
228 /// ```
229 pub fn preset_size(&self) -> Option<usize> {
230 self.size
231 }
232
233 /// Returns the maximum chunk size for chunked transfers.
234 ///
235 /// If none is explicitly set, defaults to [`Body::DEFAULT_MAX_CHUNK`].
236 ///
237 /// # Example
238 ///
239 /// ```rust
240 /// use std::io::Cursor;
241 /// use rocket::response::{Response, Body};
242 ///
243 /// let body = "Brewing the best coffee!";
244 /// let r = Response::build()
245 /// .sized_body(body.len(), Cursor::new(body))
246 /// .finalize();
247 ///
248 /// assert_eq!(r.body().max_chunk_size(), Body::DEFAULT_MAX_CHUNK);
249 ///
250 /// let r = Response::build()
251 /// .sized_body(body.len(), Cursor::new(body))
252 /// .max_chunk_size(1024)
253 /// .finalize();
254 ///
255 /// assert_eq!(r.body().max_chunk_size(), 1024);
256 /// ```
257 pub fn max_chunk_size(&self) -> usize {
258 self.max_chunk
259 }
260
261 /// Attempts to compute the body's size and returns it if the body is sized.
262 ///
263 /// If the size was preset (see [`Body::preset_size()`]), the value is
264 /// returned immediately as `Some`. If the body is unsized or computing the
265 /// size fails, returns `None`. Otherwise, the size is computed by seeking,
266 /// and the `preset_size` is updated to reflect the known value.
267 ///
268 /// **Note:** the number of bytes read from the reader and/or written to the
269 /// network may differ from the value returned by this method. Some examples
270 /// include:
271 ///
272 /// * bodies in response to `HEAD` requests are never read or written
273 /// * the client may close the connection before the body is read fully
274 /// * reading the body may fail midway
275 /// * a preset size may differ from the actual body size
276 ///
277 /// # Example
278 ///
279 /// ```rust
280 /// use std::io::Cursor;
281 /// use rocket::response::Response;
282 ///
283 /// # rocket::async_test(async {
284 /// let body = "Hello, Rocketeers!";
285 /// let mut r = Response::build()
286 /// .sized_body(None, Cursor::new(body))
287 /// .finalize();
288 ///
289 /// assert_eq!(r.body().preset_size(), None);
290 /// assert_eq!(r.body_mut().size().await, Some(body.len()));
291 /// assert_eq!(r.body().preset_size(), Some(body.len()));
292 /// # });
293 /// ```
294 pub async fn size(&mut self) -> Option<usize> {
295 if let Some(size) = self.size {
296 return Some(size);
297 }
298
299 if let Inner::Seekable(ref mut body) | Inner::Phantom(ref mut body) = self.inner {
300 let pos = body.seek(io::SeekFrom::Current(0)).await.ok()?;
301 let end = body.seek(io::SeekFrom::End(0)).await.ok()?;
302 body.seek(io::SeekFrom::Start(pos)).await.ok()?;
303
304 let size = end as usize - pos as usize;
305 self.size = Some(size);
306 return Some(size);
307 }
308
309 None
310 }
311
312 /// Moves the body out of `self` and returns it, leaving a
313 /// [`Body::default()`] in its place.
314 ///
315 /// # Example
316 ///
317 /// ```rust
318 /// use std::io::Cursor;
319 /// use rocket::response::Response;
320 ///
321 /// let mut r = Response::build()
322 /// .sized_body(None, Cursor::new("Hi"))
323 /// .finalize();
324 ///
325 /// assert!(r.body().is_some());
326 ///
327 /// let body = r.body_mut().take();
328 /// assert!(body.is_some());
329 /// assert!(r.body().is_none());
330 /// ```
331 #[inline(always)]
332 pub fn take(&mut self) -> Self {
333 std::mem::take(self)
334 }
335
336 /// Reads all of `self` into a vector of bytes, consuming the contents.
337 ///
338 /// If reading fails, returns `Err`. Otherwise, returns `Ok`. Calling this
339 /// method may partially or fully consume the body's content. As such,
340 /// subsequent calls to `to_bytes()` will likely return different result.
341 ///
342 /// # Example
343 ///
344 /// ```rust
345 /// use std::io;
346 /// use rocket::response::Response;
347 ///
348 /// # let ok: io::Result<()> = rocket::async_test(async {
349 /// let mut r = Response::build()
350 /// .streamed_body(io::Cursor::new(&[1, 2, 3, 11, 13, 17]))
351 /// .finalize();
352 ///
353 /// assert_eq!(r.body_mut().to_bytes().await?, &[1, 2, 3, 11, 13, 17]);
354 /// # Ok(())
355 /// # });
356 /// # assert!(ok.is_ok());
357 /// ```
358 pub async fn to_bytes(&mut self) -> io::Result<Vec<u8>> {
359 let mut vec = Vec::new();
360 let n = match self.read_to_end(&mut vec).await {
361 Ok(n) => n,
362 Err(e) => {
363 error!("i/o error reading body: {:?}", e);
364 return Err(e);
365 }
366 };
367
368 if let Some(ref mut size) = self.size {
369 *size = size.checked_sub(n).unwrap_or(0);
370 }
371
372 Ok(vec)
373 }
374
375 /// Reads all of `self` into a string, consuming the contents.
376 ///
377 /// If reading fails, or the body contains invalid UTF-8 characters, returns
378 /// `Err`. Otherwise, returns `Ok`. Calling this method may partially or
379 /// fully consume the body's content. As such, subsequent calls to
380 /// `to_string()` will likely return different result.
381 ///
382 /// # Example
383 ///
384 /// ```rust
385 /// use std::io;
386 /// use rocket::response::Response;
387 ///
388 /// # let ok: io::Result<()> = rocket::async_test(async {
389 /// let mut r = Response::build()
390 /// .streamed_body(io::Cursor::new("Hello, Rocketeers!"))
391 /// .finalize();
392 ///
393 /// assert_eq!(r.body_mut().to_string().await?, "Hello, Rocketeers!");
394 /// # Ok(())
395 /// # });
396 /// # assert!(ok.is_ok());
397 /// ```
398 pub async fn to_string(&mut self) -> io::Result<String> {
399 String::from_utf8(self.to_bytes().await?).map_err(|e| {
400 error!("invalid body UTF-8: {}", e);
401 io::Error::new(io::ErrorKind::InvalidData, e)
402 })
403 }
404}
405
406impl AsyncRead for Body<'_> {
407 fn poll_read(
408 mut self: Pin<&mut Self>,
409 cx: &mut Context<'_>,
410 buf: &mut ReadBuf<'_>,
411 ) -> Poll<io::Result<()>> {
412 let reader = match self.inner {
413 Inner::Seekable(ref mut b) => b as &mut (dyn AsyncRead + Unpin),
414 Inner::Unsized(ref mut b) => b as &mut (dyn AsyncRead + Unpin),
415 Inner::Phantom(_) | Inner::None => return Poll::Ready(Ok(())),
416 };
417
418 Pin::new(reader).poll_read(cx, buf)
419 }
420}
421
422impl fmt::Debug for Inner<'_> {
423 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
424 match self {
425 Inner::Seekable(_) => "seekable".fmt(f),
426 Inner::Unsized(_) => "unsized".fmt(f),
427 Inner::Phantom(_) => "phantom".fmt(f),
428 Inner::None => "none".fmt(f),
429 }
430 }
431}