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