rocket/local/asynchronous/
response.rs

1use std::io;
2use std::future::Future;
3use std::{pin::Pin, task::{Context, Poll}};
4
5use tokio::io::{AsyncRead, ReadBuf};
6
7use crate::http::CookieJar;
8use crate::{Request, Response};
9
10/// An `async` response from a dispatched [`LocalRequest`](super::LocalRequest).
11///
12/// This `LocalResponse` implements [`tokio::io::AsyncRead`]. As such, if
13/// [`into_string()`](LocalResponse::into_string()) and
14/// [`into_bytes()`](LocalResponse::into_bytes()) do not suffice, the response's
15/// body can be read directly:
16///
17/// ```rust
18/// # #[macro_use] extern crate rocket;
19/// use std::io;
20///
21/// use rocket::local::asynchronous::Client;
22/// use rocket::tokio::io::AsyncReadExt;
23/// use rocket::http::Status;
24///
25/// #[get("/")]
26/// fn hello_world() -> &'static str {
27///     "Hello, world!"
28/// }
29///
30/// #[launch]
31/// fn rocket() -> _ {
32///     rocket::build().mount("/", routes![hello_world])
33///     #    .reconfigure(rocket::Config::debug_default())
34/// }
35///
36/// # async fn read_body_manually() -> io::Result<()> {
37/// // Dispatch a `GET /` request.
38/// let client = Client::tracked(rocket()).await.expect("valid rocket");
39/// let mut response = client.get("/").dispatch().await;
40///
41/// // Check metadata validity.
42/// assert_eq!(response.status(), Status::Ok);
43/// assert_eq!(response.body().preset_size(), Some(13));
44///
45/// // Read 10 bytes of the body. Note: in reality, we'd use `into_string()`.
46/// let mut buffer = [0; 10];
47/// response.read(&mut buffer).await?;
48/// assert_eq!(buffer, "Hello, wor".as_bytes());
49/// # Ok(())
50/// # }
51/// # rocket::async_test(read_body_manually()).expect("read okay");
52/// ```
53///
54/// For more, see [the top-level documentation](../index.html#localresponse).
55pub struct LocalResponse<'c> {
56    // XXX: SAFETY: This (dependent) field must come first due to drop order!
57    response: Response<'c>,
58    cookies: CookieJar<'c>,
59    _request: Box<Request<'c>>,
60}
61
62impl Drop for LocalResponse<'_> {
63    fn drop(&mut self) { }
64}
65
66impl<'c> LocalResponse<'c> {
67    pub(crate) fn new<F, O>(req: Request<'c>, f: F) -> impl Future<Output = LocalResponse<'c>>
68        where F: FnOnce(&'c Request<'c>) -> O + Send,
69              O: Future<Output = Response<'c>> + Send
70    {
71        // `LocalResponse` is a self-referential structure. In particular,
72        // `response` and `cookies` can refer to `_request` and its contents. As
73        // such, we must
74        //   1) Ensure `Request` has a stable address.
75        //
76        //      This is done by `Box`ing the `Request`, using only the stable
77        //      address thereafter.
78        //
79        //   2) Ensure no refs to `Request` or its contents leak with a lifetime
80        //      extending beyond that of `&self`.
81        //
82        //      We have no methods that return an `&Request`. However, we must
83        //      also ensure that `Response` doesn't leak any such references. To
84        //      do so, we don't expose the `Response` directly in any way;
85        //      otherwise, methods like `.headers()` could, in conjunction with
86        //      particular crafted `Responder`s, potentially be used to obtain a
87        //      reference to contents of `Request`. All methods, instead, return
88        //      references bounded by `self`. This is easily verified by noting
89        //      that 1) `LocalResponse` fields are private, and 2) all `impl`s
90        //      of `LocalResponse` aside from this method abstract the lifetime
91        //      away as `'_`, ensuring it is not used for any output value.
92        let boxed_req = Box::new(req);
93        let request: &'c Request<'c> = unsafe { &*(&*boxed_req as *const _) };
94
95        async move {
96            // NOTE: The cookie jar `secure` state will not reflect the last
97            // known value in `request.cookies()`. This is okay: new cookies
98            // should never be added to the resulting jar which is the only time
99            // the value is used to set cookie defaults.
100            let response: Response<'c> = f(request).await;
101            let mut cookies = CookieJar::new(None, request.rocket());
102            for cookie in response.cookies() {
103                cookies.add_original(cookie.into_owned());
104            }
105
106            LocalResponse { _request: boxed_req, cookies, response, }
107        }
108    }
109}
110
111impl LocalResponse<'_> {
112    pub(crate) fn _response(&self) -> &Response<'_> {
113        &self.response
114    }
115
116    pub(crate) fn _cookies(&self) -> &CookieJar<'_> {
117        &self.cookies
118    }
119
120    pub(crate) async fn _into_string(mut self) -> io::Result<String> {
121        self.response.body_mut().to_string().await
122    }
123
124    pub(crate) async fn _into_bytes(mut self) -> io::Result<Vec<u8>> {
125        self.response.body_mut().to_bytes().await
126    }
127
128    #[cfg(feature = "json")]
129    async fn _into_json<T>(self) -> Option<T>
130        where T: Send + serde::de::DeserializeOwned + 'static
131    {
132        self.blocking_read(|r| serde_json::from_reader(r)).await?.ok()
133    }
134
135    #[cfg(feature = "msgpack")]
136    async fn _into_msgpack<T>(self) -> Option<T>
137        where T: Send + serde::de::DeserializeOwned + 'static
138    {
139        self.blocking_read(|r| rmp_serde::from_read(r)).await?.ok()
140    }
141
142    #[cfg(any(feature = "json", feature = "msgpack"))]
143    async fn blocking_read<T, F>(mut self, f: F) -> Option<T>
144        where T: Send + 'static,
145              F: FnOnce(&mut dyn io::Read) -> T + Send + 'static
146    {
147        use tokio::sync::mpsc;
148        use tokio::io::AsyncReadExt;
149
150        struct ChanReader {
151            last: Option<io::Cursor<Vec<u8>>>,
152            rx: mpsc::Receiver<io::Result<Vec<u8>>>,
153        }
154
155        impl std::io::Read for ChanReader {
156            fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
157                loop {
158                    if let Some(ref mut cursor) = self.last {
159                        if cursor.position() < cursor.get_ref().len() as u64 {
160                            return std::io::Read::read(cursor, buf);
161                        }
162                    }
163
164                    if let Some(buf) = self.rx.blocking_recv() {
165                        self.last = Some(io::Cursor::new(buf?));
166                    } else {
167                        return Ok(0);
168                    }
169                }
170            }
171        }
172
173        let (tx, rx) = mpsc::channel(2);
174        let reader = tokio::task::spawn_blocking(move || {
175            let mut reader = ChanReader { last: None, rx };
176            f(&mut reader)
177        });
178
179        loop {
180            // TODO: Try to fill as much as the buffer before send it off?
181            let mut buf = Vec::with_capacity(1024);
182            match self.read_buf(&mut buf).await {
183                Ok(0) => break,
184                Ok(_) => tx.send(Ok(buf)).await.ok()?,
185                Err(e) => {
186                    tx.send(Err(e)).await.ok()?;
187                    break;
188                }
189            }
190        }
191
192        // NOTE: We _must_ drop tx now to prevent a deadlock!
193        drop(tx);
194
195        reader.await.ok()
196    }
197
198    // Generates the public API methods, which call the private methods above.
199    pub_response_impl!("# use rocket::local::asynchronous::Client;\n\
200        use rocket::local::asynchronous::LocalResponse;" async await);
201}
202
203impl AsyncRead for LocalResponse<'_> {
204    fn poll_read(
205        mut self: Pin<&mut Self>,
206        cx: &mut Context<'_>,
207        buf: &mut ReadBuf<'_>,
208    ) -> Poll<io::Result<()>> {
209        Pin::new(self.response.body_mut()).poll_read(cx, buf)
210    }
211}
212
213impl std::fmt::Debug for LocalResponse<'_> {
214    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
215        self._response().fmt(f)
216    }
217}