Skip to main content

wasmtime_wasi_http/p3/
response.rs

1use crate::p3::bindings::http::types::ErrorCode;
2use crate::p3::body::{Body, BodyExt as _, GuestBody};
3use crate::{Error, FieldMap, WasiHttpCtxView, WasiHttpHooks, WasiHttpView, get_content_length};
4use bytes::Bytes;
5use http::StatusCode;
6use http_body_util::BodyExt as _;
7use http_body_util::combinators::UnsyncBoxBody;
8use tokio::sync::oneshot;
9use wasmtime::AsContextMut;
10use wasmtime::error::Context as _;
11
12/// The concrete type behind a `wasi:http/types.response` resource.
13pub struct Response {
14    /// The status of the response.
15    pub status: StatusCode,
16    /// The headers of the response.
17    pub headers: FieldMap,
18    /// Response body.
19    pub(crate) body: Body,
20}
21
22impl TryFrom<Response> for http::Response<Body> {
23    type Error = http::Error;
24
25    fn try_from(
26        Response {
27            status,
28            headers,
29            body,
30        }: Response,
31    ) -> Result<Self, Self::Error> {
32        let mut res = http::Response::builder().status(status);
33        *res.headers_mut().unwrap() = headers.into();
34        res.body(body)
35    }
36}
37
38impl Response {
39    /// Convert [Response] into [http::Response].
40    ///
41    /// The specified [Future] `fut` can be used to communicate
42    /// a response processing error, if any, to the constructor of the response.
43    /// For example, if the response was constructed via `wasi:http/types.response#new`,
44    /// a result sent on `fut` will be forwarded to the guest on the future handle returned.
45    pub fn into_http<T: WasiHttpView + 'static>(
46        self,
47        store: impl AsContextMut<Data = T>,
48        fut: impl Future<Output = Result<(), Error>> + Send + 'static,
49    ) -> wasmtime::Result<http::Response<UnsyncBoxBody<Bytes, Error>>> {
50        self.into_http_with_getter(store, fut, T::http)
51    }
52
53    /// Like [`Self::into_http`], but with a custom function for converting `T`
54    /// to a [`WasiHttpCtxView`].
55    pub fn into_http_with_getter<T: 'static>(
56        self,
57        mut store: impl AsContextMut<Data = T>,
58        fut: impl Future<Output = Result<(), Error>> + Send + 'static,
59        mut getter: impl FnMut(&mut T) -> WasiHttpCtxView<'_> + Clone + Unpin + Send + 'static,
60    ) -> wasmtime::Result<http::Response<UnsyncBoxBody<Bytes, Error>>> {
61        let res = http::Response::try_from(self)?;
62        let (res, body) = res.into_parts();
63        // `Content-Length` header value is validated in `fields` implementation
64        let content_length =
65            get_content_length(&res.headers).context("failed to parse `content-length`")?;
66        let body = match body {
67            Body::Guest {
68                contents_rx,
69                trailers_rx,
70                result_tx,
71            } => GuestBody::new(
72                &mut store,
73                contents_rx,
74                trailers_rx,
75                result_tx,
76                fut,
77                content_length,
78                ErrorCode::HttpResponseBodySize,
79                getter,
80            )?
81            .boxed_unsync(),
82            Body::Host { body, result_tx } => {
83                let is_forbidden = getter(store.as_context_mut().data_mut())
84                    .hooks
85                    .forbidden_header_filter();
86                let body = body.with_filtered_trailers(is_forbidden);
87                if let Some(limit) = content_length {
88                    let (http_result_tx, http_result_rx) = oneshot::channel();
89                    _ = result_tx.send(Box::new(async move {
90                        if let Ok(err) = http_result_rx.await {
91                            return Err(err);
92                        };
93                        fut.await
94                    }));
95                    body.with_content_length(limit, http_result_tx, Error::HttpResponseBodySize)
96                        .boxed_unsync()
97                } else {
98                    _ = result_tx.send(Box::new(fut));
99                    body.boxed_unsync()
100                }
101            }
102        };
103        Ok(http::Response::from_parts(res, body))
104    }
105
106    /// Convert [http::Response] into [Response].
107    pub fn from_http<T>(
108        hooks: &mut dyn WasiHttpHooks,
109        res: http::Response<T>,
110    ) -> (
111        Self,
112        impl Future<Output = Result<(), Error>> + Send + 'static + use<T>,
113    )
114    where
115        T: http_body::Body<Data = Bytes> + Send + 'static,
116        T::Error: Into<Error>,
117    {
118        let (parts, body) = res.into_parts();
119        let (result_tx, result_rx) = tokio::sync::oneshot::channel();
120
121        let wasi_response = Response {
122            status: parts.status,
123            headers: FieldMap::new_immutable(hooks, parts.headers),
124            body: Body::Host {
125                body: body.map_err(Into::into).boxed_unsync(),
126                result_tx,
127            },
128        };
129
130        let io_future = async {
131            let Ok(fut) = result_rx.await else {
132                return Ok(());
133            };
134            Box::into_pin(fut).await
135        };
136
137        (wasi_response, io_future)
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use crate::default_hooks;
145    use core::future::Future;
146    use core::pin::pin;
147    use core::task::{Context, Poll, Waker};
148    use http_body_util::Full;
149
150    #[tokio::test]
151    async fn test_response_from_http() {
152        let http_response = http::Response::builder()
153            .status(StatusCode::OK)
154            .header("x-custom-header", "value123")
155            .body(Full::new(Bytes::from_static(b"hello wasm")))
156            .unwrap();
157
158        let (wasi_resp, io_future) = Response::from_http(default_hooks(), http_response);
159        assert_eq!(wasi_resp.status, StatusCode::OK);
160        assert_eq!(
161            wasi_resp.headers.get("x-custom-header").unwrap(),
162            "value123"
163        );
164        match wasi_resp.body {
165            Body::Host { body, result_tx } => {
166                let collected = body.collect().await;
167                assert!(collected.is_ok(), "Body stream failed unexpectedly");
168                let chunks = collected.unwrap().to_bytes();
169                assert_eq!(chunks, &b"hello wasm"[..]);
170                _ = result_tx.send(Box::new(async { Ok(()) }));
171            }
172            _ => panic!("Response body should be of type Host"),
173        }
174
175        let mut cx = Context::from_waker(Waker::noop());
176        let result = pin!(io_future).poll(&mut cx);
177        assert!(matches!(result, Poll::Ready(Ok(_))));
178    }
179}