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}