wasmtime_wasi_http/p3/
response.rs1use 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
12pub struct Response {
14 pub status: StatusCode,
16 pub headers: FieldMap,
18 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 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 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 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 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}