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(rx: oneshot::Receiver<Result<(), Error>>) -> Result<(), Error> {
31    rx.await
32        .unwrap_or_else(|_| Err(Error::InternalError(Some(DROPPED_FUTURE_ERROR.to_string()))))
33}
34
35fn send_dummy_io_err<T, D>(
36    store: &Accessor<T, D>,
37    mut getter: impl FnMut(&mut T) -> WasiHttpCtxView<'_>,
38    e: Error,
39    io_result_tx: oneshot::Sender<Result<(), Error>>,
40) -> HttpError
41where
42    D: HasData,
43{
44    let err_code =
45        store.with(|mut store| getter(store.as_context_mut().data_mut()).error_to_p3(&e));
46    _ = io_result_tx.send(Err(e));
47    err_code.into()
48}
49
50impl<T> HostWithStore<T> for WasiHttp {
51    async fn send(
52        store: &Accessor<T, Self>,
53        req: Resource<Request>,
54    ) -> HttpResult<Resource<Response>> {
55        let getter = store.getter();
56        send(store, getter, req).await
57    }
58}
59
60async fn send<T, D>(
61    store: &Accessor<T, D>,
62    mut getter: impl FnMut(&mut T) -> WasiHttpCtxView<'_> + Copy + Unpin + Send + 'static,
63    req: Resource<Request>,
64) -> HttpResult<Resource<Response>>
65where
66    D: HasData,
67    T: 'static,
68{
69    // The I/O result will be sent on this channel
70    let (io_result_tx, io_result_rx) = oneshot::channel();
71
72    // Response processing result will be sent on this channel
73    let (res_result_tx, res_result_rx) = oneshot::channel();
74
75    let fut = store.with(|mut store| {
76        let WasiHttpCtxView { table, .. } = getter(store.data_mut());
77        let req = table
78            .delete(req)
79            .context("failed to delete request from table")
80            .map_err(HttpError::trap)?;
81        let (req, options) =
82            req.into_http_with_getter(&mut store, io_task_result(io_result_rx), getter)?;
83        HttpResult::Ok(getter(store.data_mut()).hooks.send_request(
84            req,
85            options.as_deref().copied(),
86            Box::new(async {
87                // Forward the response processing result to `WasiHttpCtx` implementation
88                let Ok(fut) = res_result_rx.await else {
89                    return Ok(());
90                };
91                Box::into_pin(fut).await
92            }),
93        ))
94    });
95    let fut = match fut {
96        Ok(fut) => fut,
97        Err(e) => match e.downcast() {
98            Ok(err_code) => {
99                _ = io_result_tx.send(Err(err_code.clone().into()));
100                return Err(err_code.into());
101            }
102            Err(e) => {
103                let e = Error::InternalError(Some(format!("{e}")));
104                return Err(send_dummy_io_err(store, getter, e, io_result_tx));
105            }
106        },
107    };
108    let (res, io) = match Box::into_pin(fut).await {
109        Ok(r) => r,
110        Err(e) => {
111            return Err(send_dummy_io_err(store, getter, e, io_result_tx));
112        }
113    };
114    let (
115        http::response::Parts {
116            status, headers, ..
117        },
118        body,
119    ) = res.into_parts();
120
121    let mut io = Box::into_pin(io);
122    let body = match io.as_mut().poll(&mut Context::from_waker(Waker::noop())) {
123        Poll::Ready(Ok(())) => {
124            _ = io_result_tx.send(Ok(()));
125            body
126        }
127        Poll::Ready(Err(e)) => {
128            return Err(send_dummy_io_err(store, getter, e, io_result_tx));
129        }
130        Poll::Pending => {
131            // I/O driver still needs to be polled, spawn a task
132            let (done_tx, done_rx) = oneshot::channel::<()>();
133            let io = Arc::new(AbortOnDropJoinHandle(task::spawn(async move {
134                let _done_tx = done_tx;
135                let res = io.await;
136                debug!(?res, "`send_request` I/O future finished");
137                _ = io_result_tx.send(res);
138            })));
139            // The store holds a reference to the task until it finishes, so
140            // that dropping the store cancels any I/O not referenced by
141            // anything outside the store, e.g. a response body handed to the
142            // host.
143            let store_io = Arc::clone(&io);
144            store
145                .spawn(async move |_| {
146                    let _io = store_io;
147                    _ = done_rx.await;
148                    Ok(())
149                })
150                .map_err(HttpError::trap)?;
151            // Attach a reference to the io task to the response body so that
152            // it isn't cancelled while the body is being read.
153            body.with_state(io).boxed_unsync()
154        }
155    };
156    store.with(|mut store| {
157        let view = getter(store.data_mut());
158        let res = Response {
159            status,
160            headers: FieldMap::new_immutable(view.hooks, headers),
161            body: Body::Host {
162                body,
163                result_tx: res_result_tx,
164            },
165        };
166        view.table
167            .push(res)
168            .context("failed to push response to table")
169            .map_err(HttpError::trap)
170    })
171}
172
173impl Host for WasiHttpCtxView<'_> {}
174
175mod named {
176    use crate::p3::bindings::http::types::{ErrorCode, Request, Response};
177    use crate::p3::bindings::named_imports::wasi::http::client::{Host, HostWithStore};
178    use crate::p3::{HttpError, HttpResult};
179    use crate::{WasiHttpNamed, WasiHttpNamedView};
180    use wasmtime::component::{Accessor, Resource};
181    use wasmtime_wasi::{NamedId, WasiCtxNamedView};
182
183    impl<T, U> HostWithStore<U> for WasiHttpNamed<T>
184    where
185        T: WasiHttpNamedView,
186        U: 'static,
187    {
188        async fn send(
189            store: &Accessor<U, Self>,
190            id: NamedId,
191            req: Resource<Request>,
192        ) -> HttpResult<Resource<Response>> {
193            let getter = store.getter();
194            super::send(store, move |data| getter(data).0.http(id), req).await
195        }
196    }
197
198    impl<T> Host for WasiCtxNamedView<'_, T>
199    where
200        T: WasiHttpNamedView,
201    {
202        fn convert_error_code(&mut self, error: HttpError) -> wasmtime::Result<ErrorCode> {
203            error.downcast()
204        }
205    }
206}
wasmtime_wasi_http/p3/host/handler.rs

wasmtime_wasi_http/p3/host/
handler.rs
