wasmtime_wasi_http/p3/host/
handler.rs

1use crate::FieldMap;
2use crate::p3::bindings::http::client::{Host, HostWithStore};
3use crate::p3::bindings::http::types::{Request, Response};
4use crate::p3::body::{Body, BodyExt as _};
5use crate::p3::{HttpError, HttpResult};
6use crate::{Error, WasiHttp, WasiHttpCtxView};
7use core::task::{Context, Poll, Waker};
8use http_body_util::BodyExt as _;
9use std::sync::Arc;
10use tokio::sync::oneshot;
11use tokio::task::{self, JoinHandle};
12use tracing::debug;
13use wasmtime::AsContextMut as _;
14use wasmtime::component::{Accessor, HasData, Resource};
15use wasmtime::error::Context as _;
16
17/// A wrapper around [`JoinHandle`], which will [`JoinHandle::abort`] the task
18/// when dropped
19struct AbortOnDropJoinHandle(JoinHandle<()>);
20
21impl Drop for AbortOnDropJoinHandle {
22    fn drop(&mut self) {
23        self.0.abort();
24    }
25}
26
27const DROPPED_FUTURE_ERROR: &str =
28    "Future indicating transmission result dropped without being resolved.";
29
30async fn io_task_result(
31    rx: oneshot::Receiver<(
32        Option<Arc<AbortOnDropJoinHandle>>,
33        oneshot::Receiver<Result<(), Error>>,
34    )>,
35) -> Result<(), Error> {
36    let Ok((_io, io_result_rx)) = rx.await else {
37        return Err(Error::InternalError(Some(DROPPED_FUTURE_ERROR.to_string())));
38    };
39    io_result_rx
40        .await
41        .unwrap_or_else(|_| Err(Error::InternalError(Some(DROPPED_FUTURE_ERROR.to_string()))))
42}
43
44fn send_dummy_io(
45    result: Result<(), Error>,
46    io_result_tx: oneshot::Sender<(
47        Option<Arc<AbortOnDropJoinHandle>>,
48        oneshot::Receiver<Result<(), Error>>,
49    )>,
50) {
51    let (tx, rx) = oneshot::channel();
52    let _ = tx.send(result);
53    let _ = io_result_tx.send((None, rx));
54}
55
56fn send_dummy_io_err<T, D>(
57    store: &Accessor<T, D>,
58    mut getter: impl FnMut(&mut T) -> WasiHttpCtxView<'_>,
59    e: Error,
60    io_result_tx: oneshot::Sender<(
61        Option<Arc<AbortOnDropJoinHandle>>,
62        oneshot::Receiver<Result<(), Error>>,
63    )>,
64) -> HttpError
65where
66    D: HasData,
67{
68    let err_code =
69        store.with(|mut store| getter(store.as_context_mut().data_mut()).error_to_p3(&e));
70    send_dummy_io(Err(e), io_result_tx);
71    err_code.into()
72}
73
74impl<T> HostWithStore<T> for WasiHttp {
75    async fn send(
76        store: &Accessor<T, Self>,
77        req: Resource<Request>,
78    ) -> HttpResult<Resource<Response>> {
79        let getter = store.getter();
80        send(store, getter, req).await
81    }
82}
83
84async fn send<T, D>(
85    store: &Accessor<T, D>,
86    mut getter: impl FnMut(&mut T) -> WasiHttpCtxView<'_> + Copy + Unpin + Send + 'static,
87    req: Resource<Request>,
88) -> HttpResult<Resource<Response>>
89where
90    D: HasData,
91    T: 'static,
92{
93    // A handle to the I/O task, if spawned, will be sent on this channel
94    // and kept as part of request body state
95    let (io_task_tx, io_task_rx) = oneshot::channel();
96
97    // A handle to the I/O task, if spawned, will be sent on this channel
98    // along with the result receiver
99    let (io_result_tx, io_result_rx) = oneshot::channel();
100
101    // Response processing result will be sent on this channel
102    let (res_result_tx, res_result_rx) = oneshot::channel();
103
104    let fut = store.with(|mut store| {
105        let WasiHttpCtxView { table, .. } = getter(store.data_mut());
106        let req = table
107            .delete(req)
108            .context("failed to delete request from table")
109            .map_err(HttpError::trap)?;
110        let (req, options) =
111            req.into_http_with_getter(&mut store, io_task_result(io_result_rx), getter)?;
112        HttpResult::Ok(getter(store.data_mut()).hooks.send_request(
113            // Attach a reference to the io task to the body so that it
114            // isn't cancelled if the body is dropped.
115            req.map(|body| body.with_state(io_task_rx).boxed_unsync()),
116            options.as_deref().copied(),
117            Box::new(async {
118                // Forward the response processing result to `WasiHttpCtx` implementation
119                let Ok(fut) = res_result_rx.await else {
120                    return Ok(());
121                };
122                Box::into_pin(fut).await
123            }),
124        ))
125    });
126    let fut = match fut {
127        Ok(fut) => fut,
128        Err(e) => match e.downcast() {
129            Ok(err_code) => {
130                send_dummy_io(Err(err_code.clone().into()), io_result_tx);
131                return Err(err_code.into());
132            }
133            Err(e) => {
134                let e = Error::InternalError(Some(format!("{e}")));
135                return Err(send_dummy_io_err(store, getter, e, io_result_tx));
136            }
137        },
138    };
139    let (res, io) = match Box::into_pin(fut).await {
140        Ok(r) => r,
141        Err(e) => {
142            return Err(send_dummy_io_err(store, getter, e, io_result_tx));
143        }
144    };
145    let (
146        http::response::Parts {
147            status, headers, ..
148        },
149        body,
150    ) = res.into_parts();
151
152    let mut io = Box::into_pin(io);
153    let body = match io.as_mut().poll(&mut Context::from_waker(Waker::noop())) {
154        Poll::Ready(Ok(())) => {
155            send_dummy_io(Ok(()), io_result_tx);
156            body
157        }
158        Poll::Ready(Err(e)) => {
159            return Err(send_dummy_io_err(store, getter, e, io_result_tx));
160        }
161        Poll::Pending => {
162            // I/O driver still needs to be polled, spawn a task and send handles to it
163            let (tx, rx) = oneshot::channel();
164            let io = Arc::new(AbortOnDropJoinHandle(task::spawn(async move {
165                let res = io.await;
166                debug!(?res, "`send_request` I/O future finished");
167                _ = tx.send(res);
168            })));
169            _ = io_result_tx.send((Some(Arc::clone(&io)), rx));
170            _ = io_task_tx.send(Arc::clone(&io));
171            // Attach a reference to the io task to the body so that it
172            // isn't cancelled if the body is dropped.
173            body.with_state(io).boxed_unsync()
174        }
175    };
176    store.with(|mut store| {
177        let view = getter(store.data_mut());
178        let res = Response {
179            status,
180            headers: FieldMap::new_immutable(view.hooks, headers),
181            body: Body::Host {
182                body,
183                result_tx: res_result_tx,
184            },
185        };
186        view.table
187            .push(res)
188            .context("failed to push response to table")
189            .map_err(HttpError::trap)
190    })
191}
192
193impl Host for WasiHttpCtxView<'_> {}
194
195mod named {
196    use crate::p3::bindings::http::types::{ErrorCode, Request, Response};
197    use crate::p3::bindings::named_imports::wasi::http::client::{Host, HostWithStore};
198    use crate::p3::{HttpError, HttpResult};
199    use crate::{WasiHttpNamed, WasiHttpNamedView};
200    use wasmtime::component::{Accessor, Resource};
201    use wasmtime_wasi::{NamedId, WasiCtxNamedView};
202
203    impl<T, U> HostWithStore<U> for WasiHttpNamed<T>
204    where
205        T: WasiHttpNamedView,
206        U: 'static,
207    {
208        async fn send(
209            store: &Accessor<U, Self>,
210            id: NamedId,
211            req: Resource<Request>,
212        ) -> HttpResult<Resource<Response>> {
213            let getter = store.getter();
214            super::send(store, move |data| getter(data).0.http(id), req).await
215        }
216    }
217
218    impl<T> Host for WasiCtxNamedView<'_, T>
219    where
220        T: WasiHttpNamedView,
221    {
222        fn convert_error_code(&mut self, error: HttpError) -> wasmtime::Result<ErrorCode> {
223            error.downcast()
224        }
225    }
226}
wasmtime_wasi_http/p3/host/handler.rs

wasmtime_wasi_http/p3/host/
handler.rs
