Skip to main content

trillium_client/
response_body.rs

1use crate::Error;
2use futures_lite::AsyncRead;
3use std::pin::Pin;
4use trillium_http::ReceivedBody;
5use trillium_macros::AsyncRead;
6use trillium_server_common::Transport;
7
8/// A response body received from a server
9///
10/// This type represents a body that will be read from the underlying transport
11///
12/// ```rust
13/// use trillium_client::Client;
14/// use trillium_testing::{client_config, with_server};
15///
16/// with_server("hello from trillium", |url| async move {
17///     let client = Client::new(client_config());
18///     let mut conn = client.get(url).await?;
19///     let body = conn.response_body(); //<-
20///     assert_eq!(Some(19), body.content_length());
21///     assert_eq!("hello from trillium", body.read_string().await?);
22///     Ok(())
23/// });
24/// ```
25///
26/// ## Bounds checking
27///
28/// Every `ResponseBody` has a maximum length beyond which it will return an error, expressed as a
29/// u64. To override this on the specific `ResponseBody`, use [`ResponseBody::with_max_len`] or
30/// [`ResponseBody::set_max_len`]
31#[derive(AsyncRead, Debug)]
32pub struct ResponseBody<'a>(ReceivedBody<'a, Box<dyn Transport>>);
33
34impl ResponseBody<'_> {
35    /// Similar to [`ResponseBody::read_string`], but returns the raw bytes. This is useful for
36    /// bodies that are not text.
37    ///
38    /// You can use this in conjunction with `encoding` if you need different handling of malformed
39    /// character encoding than the lossy conversion provided by [`ResponseBody::read_string`].
40    ///
41    /// An empty or nonexistent body will yield an empty Vec, not an error.
42    ///
43    /// # Errors
44    ///
45    /// This will return an error if there is an IO error on the underlying transport such as a
46    /// disconnect
47    ///
48    /// This will also return an error if the length exceeds the maximum length. To configure the
49    /// value on this specific request body, use [`ResponseBody::with_max_len`] or
50    /// [`ResponseBody::set_max_len`]
51    pub async fn read_bytes(self) -> Result<Vec<u8>, Error> {
52        self.0.read_bytes().await
53    }
54
55    /// # Reads the entire body to `String`.
56    ///
57    /// This uses the encoding determined by the content-type (mime) charset. If an encoding problem
58    /// is encountered, the String returned by [`ResponseBody::read_string`] will contain utf8
59    /// replacement characters.
60    ///
61    /// Note that this can only be performed once per Conn, as the underlying data is not cached
62    /// anywhere. This is the only copy of the body contents.
63    ///
64    /// An empty or nonexistent body will yield an empty String, not an error
65    ///
66    /// # Errors
67    ///
68    /// This will return an error if there is an IO error on the
69    /// underlying transport such as a disconnect
70    ///
71    ///
72    /// This will also return an error if the length exceeds the maximum length. To configure the
73    /// value on this specific response body, use [`ResponseBody::with_max_len`] or
74    /// [`ResponseBody::set_max_len`].
75    pub async fn read_string(self) -> Result<String, Error> {
76        self.0.read_string().await
77    }
78
79    /// Set the maximum content length to read, returning self
80    ///
81    /// This protects against an memory-use denial-of-service attack wherein an untrusted peer sends
82    /// an unbounded request body. This is especially important when using
83    /// [`ResponseBody::read_string`] and [`ResponseBody::read_bytes`] instead of streaming with
84    /// `AsyncRead`.
85    ///
86    /// The default value can be found documented [in the trillium-http
87    /// crate](https://docs.trillium.rs/trillium_http/struct.httpconfig#received_body_max_len)
88    #[must_use]
89    pub fn with_max_len(mut self, max_len: u64) -> Self {
90        self.0.set_max_len(max_len);
91        self
92    }
93
94    /// Set the maximum content length to read
95    ///
96    /// This protects against an memory-use denial-of-service attack wherein an untrusted peer sends
97    /// an unbounded request body. This is especially important when using
98    /// [`ResponseBody::read_string`] and [`ResponseBody::read_bytes`] instead of streaming with
99    /// `AsyncRead`.
100    ///
101    /// The default value can be found documented [in the trillium-http
102    /// crate](https://docs.trillium.rs/trillium_http/struct.httpconfig#received_body_max_len)
103    pub fn set_max_len(&mut self, max_len: u64) -> &mut Self {
104        self.0.set_max_len(max_len);
105        self
106    }
107
108    /// The content-length of this body, if available.
109    ///
110    /// This value usually is derived from the content-length header. If the request that this body
111    /// is attached to uses transfer-encoding chunked, this will be None.
112    pub fn content_length(&self) -> Option<u64> {
113        self.0.content_length()
114    }
115
116    pub(crate) async fn drain(self) -> std::io::Result<u64> {
117        self.0.drain().await
118    }
119}
120
121impl<'a> From<ReceivedBody<'a, Box<dyn Transport>>> for ResponseBody<'a> {
122    fn from(received_body: ReceivedBody<'a, Box<dyn Transport>>) -> Self {
123        Self(received_body)
124    }
125}
126
127impl<'a> From<ResponseBody<'a>> for ReceivedBody<'a, Box<dyn Transport>> {
128    fn from(value: ResponseBody<'a>) -> Self {
129        value.0
130    }
131}
132
133impl<'a> IntoFuture for ResponseBody<'a> {
134    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>>;
135    type Output = trillium_http::Result<String>;
136
137    fn into_future(self) -> Self::IntoFuture {
138        Box::pin(async move { self.read_string().await })
139    }
140}