Skip to main content

wasi_hyperium/
wasi.rs

1use std::{
2    future::{Future, IntoFuture},
3    task::{Context, Poll},
4};
5
6use wasi::http::types;
7
8use crate::{
9    poll::{PollableRegistry, WasiSubscribe},
10    Error,
11};
12
13struct Subscribable<T, Registry: PollableRegistry> {
14    // NOTE: order matters; handle must be dropped before inner
15    handle: Option<Registry::RegisteredPollable>,
16    inner: T,
17    registry: Registry,
18}
19
20impl<T, Registry> Subscribable<T, Registry>
21where
22    T: WasiSubscribe,
23    Registry: PollableRegistry,
24{
25    fn new(inner: T, registry: Registry) -> Self {
26        Self {
27            handle: None,
28            inner,
29            registry,
30        }
31    }
32
33    fn register_subscribe(&mut self, cx: &mut Context) {
34        let pollable = self.inner.subscribe();
35        self.handle = Some(self.registry.register_pollable(cx, pollable));
36    }
37
38    fn maybe_subscribe(&mut self, cx: &mut Context) -> Poll<()> {
39        let pollable = self.inner.subscribe();
40        if pollable.ready() {
41            Poll::Ready(())
42        } else {
43            self.handle = Some(self.registry.register_pollable(cx, pollable));
44            Poll::Pending
45        }
46    }
47
48    fn registry(&self) -> &Registry {
49        &self.registry
50    }
51}
52
53impl<T, Registry: PollableRegistry> std::ops::Deref for Subscribable<T, Registry> {
54    type Target = T;
55
56    fn deref(&self) -> &Self::Target {
57        &self.inner
58    }
59}
60
61pub struct InputStream<Registry: PollableRegistry> {
62    stream: Subscribable<types::InputStream, Registry>,
63}
64
65impl<Registry> InputStream<Registry>
66where
67    Registry: PollableRegistry,
68{
69    pub fn new(stream: types::InputStream, registry: Registry) -> Self {
70        let stream = Subscribable::new(stream, registry);
71        Self { stream }
72    }
73
74    pub fn poll_read(&mut self, cx: &mut Context, len: usize) -> Poll<Result<Vec<u8>, Error>> {
75        let data = self
76            .stream
77            .read(len.try_into().unwrap())
78            .map_err(Error::wasi_stream_error)?;
79        if data.is_empty() {
80            self.stream.register_subscribe(cx);
81            Poll::Pending
82        } else {
83            Poll::Ready(Ok(data))
84        }
85    }
86
87    fn registry(&self) -> &Registry {
88        self.stream.registry()
89    }
90}
91
92pub struct OutputStream<Registry: PollableRegistry> {
93    stream: Subscribable<types::OutputStream, Registry>,
94}
95
96impl<Registry> OutputStream<Registry>
97where
98    Registry: PollableRegistry,
99{
100    pub fn new(stream: types::OutputStream, registry: Registry) -> Self {
101        let stream = Subscribable::new(stream, registry);
102        Self { stream }
103    }
104
105    pub fn poll_check_write(
106        &mut self,
107        cx: &mut Context,
108    ) -> Poll<Result<OutputStreamPermit, Error>> {
109        let size = self
110            .stream
111            .check_write()
112            .map_err(Error::wasi_stream_error)?;
113        if size == 0 {
114            self.stream.register_subscribe(cx);
115            Poll::Pending
116        } else {
117            Poll::Ready(Ok(OutputStreamPermit {
118                stream: &self.stream.inner,
119                size,
120            }))
121        }
122    }
123
124    pub fn poll_splice(
125        &mut self,
126        cx: &mut Context,
127        src: &InputStream<Registry>,
128        len: u64,
129    ) -> Poll<Result<u64, Error>> {
130        if len == 0 {
131            return Poll::Ready(Ok(0));
132        }
133        let size = self
134            .stream
135            .splice(&src.stream.inner, len)
136            .map_err(Error::wasi_stream_error)?;
137        if size == 0 {
138            self.stream.register_subscribe(cx);
139            Poll::Pending
140        } else {
141            Poll::Ready(Ok(size))
142        }
143    }
144
145    pub fn poll_flush(&mut self, cx: &mut Context) -> Poll<Result<(), Error>> {
146        self.stream.flush().map_err(Error::wasi_stream_error)?;
147        self.stream.maybe_subscribe(cx).map(|()| Ok(()))
148    }
149
150    fn registry(&self) -> &Registry {
151        self.stream.registry()
152    }
153}
154
155pub struct OutputStreamPermit<'a> {
156    stream: &'a types::OutputStream,
157    size: u64,
158}
159
160impl OutputStreamPermit<'_> {
161    pub fn write(self, contents: &[u8]) -> Result<usize, Error> {
162        let len = self
163            .size
164            .min(contents.len().try_into().unwrap())
165            .try_into()
166            .unwrap();
167        self.stream
168            .write(&contents[..len])
169            .map_err(Error::wasi_stream_error)?;
170        Ok(len)
171    }
172
173    pub fn size(&self) -> u64 {
174        self.size
175    }
176}
177
178pub struct IncomingBody<Registry: PollableRegistry> {
179    // NOTE: order matters; stream must be dropped before body
180    stream: InputStream<Registry>,
181    body: types::IncomingBody,
182}
183
184impl<Registry> IncomingBody<Registry>
185where
186    Registry: PollableRegistry,
187{
188    pub fn new(body: types::IncomingBody, registry: Registry) -> Result<Self, Error> {
189        let stream = InputStream::new(
190            body.stream()
191                .map_err(|()| Error::WasiInvalidState("incoming-body.stream already called"))?,
192            registry,
193        );
194        Ok(Self { stream, body })
195    }
196
197    pub fn stream(&mut self) -> &mut InputStream<Registry> {
198        &mut self.stream
199    }
200
201    pub fn finish(self) -> FutureTrailers<Registry> {
202        let Self { stream, body } = self;
203        let registry = stream.registry().clone();
204        drop(stream);
205        let wasi_trailers = types::IncomingBody::finish(body);
206        let trailers = Subscribable::new(wasi_trailers, registry);
207        FutureTrailers { trailers }
208    }
209}
210
211pub struct FutureTrailers<Registry: PollableRegistry> {
212    trailers: Subscribable<types::FutureTrailers, Registry>,
213}
214
215impl<Registry> Future for FutureTrailers<Registry>
216where
217    Registry: PollableRegistry,
218{
219    type Output = Result<Option<FieldEntries>, Error>;
220
221    fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
222        match self.trailers.get() {
223            Some(Ok(Ok(Some(fields)))) => Poll::Ready(Ok(Some(fields.into()))),
224            Some(Ok(Ok(None))) => Poll::Ready(Ok(None)),
225            Some(Ok(Err(err))) => Poll::Ready(Err(Error::wasi_error_code(err))),
226            Some(Err(())) => Poll::Ready(Err(Error::WasiInvalidState(
227                "future-trailers.get already consumed",
228            ))),
229            None => {
230                self.trailers.register_subscribe(cx);
231                Poll::Pending
232            }
233        }
234    }
235}
236
237#[derive(Debug)]
238pub enum Method {
239    Get,
240    Head,
241    Post,
242    Put,
243    Delete,
244    Connect,
245    Options,
246    Trace,
247    Patch,
248    Other(String),
249}
250
251impl From<Method> for types::Method {
252    fn from(method: Method) -> Self {
253        match method {
254            Method::Get => Self::Get,
255            Method::Head => Self::Head,
256            Method::Post => Self::Post,
257            Method::Put => Self::Put,
258            Method::Delete => Self::Delete,
259            Method::Connect => Self::Connect,
260            Method::Options => Self::Options,
261            Method::Trace => Self::Trace,
262            Method::Patch => Self::Patch,
263            Method::Other(other) => Self::Other(other),
264        }
265    }
266}
267
268impl From<types::Method> for Method {
269    fn from(method: types::Method) -> Self {
270        match method {
271            types::Method::Get => Self::Get,
272            types::Method::Head => Self::Head,
273            types::Method::Post => Self::Post,
274            types::Method::Put => Self::Put,
275            types::Method::Delete => Self::Delete,
276            types::Method::Connect => Self::Connect,
277            types::Method::Options => Self::Options,
278            types::Method::Trace => Self::Trace,
279            types::Method::Patch => Self::Patch,
280            types::Method::Other(other) => Self::Other(other),
281        }
282    }
283}
284
285#[derive(Debug)]
286pub enum Scheme {
287    Http,
288    Https,
289    Other(String),
290}
291
292impl From<Scheme> for types::Scheme {
293    fn from(scheme: Scheme) -> Self {
294        match scheme {
295            Scheme::Http => Self::Http,
296            Scheme::Https => Self::Https,
297            Scheme::Other(other) => Self::Other(other),
298        }
299    }
300}
301
302impl From<types::Scheme> for Scheme {
303    fn from(scheme: types::Scheme) -> Self {
304        match scheme {
305            types::Scheme::Http => Self::Http,
306            types::Scheme::Https => Self::Https,
307            types::Scheme::Other(other) => Self::Other(other),
308        }
309    }
310}
311
312pub struct IncomingRequest<Registry: PollableRegistry> {
313    request: types::IncomingRequest,
314    body: IncomingBody<Registry>,
315}
316
317impl<Registry> IncomingRequest<Registry>
318where
319    Registry: PollableRegistry,
320{
321    pub fn new(request: types::IncomingRequest, registry: Registry) -> Result<Self, Error> {
322        let body = request
323            .consume()
324            .map_err(|()| Error::WasiInvalidState("incoming-request.consume already called"))?;
325        let body = IncomingBody::new(body, registry)?;
326        Ok(Self { request, body })
327    }
328
329    pub fn method(&self) -> Method {
330        self.request.method().into()
331    }
332
333    pub fn path_with_query(&self) -> Option<String> {
334        self.request.path_with_query()
335    }
336
337    pub fn scheme(&self) -> Option<Scheme> {
338        self.request.scheme().map(Into::into)
339    }
340
341    pub fn authority(&self) -> Option<String> {
342        self.request.authority()
343    }
344
345    pub fn headers(&self) -> FieldEntries {
346        self.request.headers().into()
347    }
348
349    pub fn body(&mut self) -> &mut IncomingBody<Registry> {
350        &mut self.body
351    }
352
353    pub fn into_body(self) -> IncomingBody<Registry> {
354        self.body
355    }
356}
357
358pub struct IncomingResponse<Registry: PollableRegistry> {
359    response: types::IncomingResponse,
360    body: IncomingBody<Registry>,
361}
362
363impl<Registry> IncomingResponse<Registry>
364where
365    Registry: PollableRegistry,
366{
367    pub fn new(response: types::IncomingResponse, registry: Registry) -> Result<Self, Error> {
368        let body = response
369            .consume()
370            .map_err(|()| Error::WasiInvalidState("incoming-response.consume already called"))?;
371        let body = IncomingBody::new(body, registry)?;
372        Ok(Self { response, body })
373    }
374
375    pub fn status(&self) -> u16 {
376        self.response.status()
377    }
378
379    pub fn headers(&self) -> FieldEntries {
380        self.response.headers().into()
381    }
382
383    pub fn body(&mut self) -> &mut IncomingBody<Registry> {
384        &mut self.body
385    }
386
387    pub fn into_body(self) -> IncomingBody<Registry> {
388        self.body
389    }
390}
391
392pub struct OutgoingBody<Registry: PollableRegistry> {
393    // NOTE: order matters; stream must be dropped before body
394    stream: OutputStream<Registry>,
395    body: types::OutgoingBody,
396}
397
398impl<Registry> OutgoingBody<Registry>
399where
400    Registry: PollableRegistry,
401{
402    pub fn new(body: types::OutgoingBody, registry: Registry) -> Result<Self, Error> {
403        let stream = OutputStream::new(
404            body.write()
405                .map_err(|()| Error::WasiInvalidState("outgoing-body.write already called"))?,
406            registry,
407        );
408        Ok(Self { stream, body })
409    }
410
411    pub fn stream(&mut self) -> &mut OutputStream<Registry> {
412        &mut self.stream
413    }
414
415    pub fn finish(self, trailers: Option<FieldEntries>) -> Result<(), Error> {
416        let trailers = match trailers {
417            Some(trailers) => Some(trailers.try_into_fields()?),
418            None => None,
419        };
420        drop(self.stream);
421        types::OutgoingBody::finish(self.body, trailers).map_err(Error::wasi_error_code)
422    }
423
424    fn registry(&self) -> &Registry {
425        self.stream.registry()
426    }
427}
428
429pub struct OutgoingRequest<Registry: PollableRegistry> {
430    request: types::OutgoingRequest,
431    body: OutgoingBody<Registry>,
432}
433
434impl<Registry> OutgoingRequest<Registry>
435where
436    Registry: PollableRegistry,
437{
438    pub fn new(request: types::OutgoingRequest, registry: Registry) -> Result<Self, Error> {
439        let body = request
440            .body()
441            .map_err(|()| Error::WasiInvalidState("outgoing-request.body already called"))?;
442        let body = OutgoingBody::new(body, registry)?;
443        Ok(Self { request, body })
444    }
445
446    pub fn from_headers(headers: &FieldEntries, registry: Registry) -> Result<Self, Error> {
447        let fields = headers.try_into_fields()?;
448        let response = types::OutgoingRequest::new(fields);
449        Self::new(response, registry)
450    }
451
452    pub fn set_method(&mut self, method: Method) -> Result<(), Error> {
453        self.request
454            .set_method(&method.into())
455            .map_err(|()| Error::WasiInvalidValue("invalid method"))
456    }
457
458    pub fn set_path_with_query(&mut self, path_with_query: Option<&str>) -> Result<(), Error> {
459        self.request
460            .set_path_with_query(path_with_query)
461            .map_err(|()| Error::WasiInvalidValue("invalid path_with_query"))
462    }
463
464    pub fn set_scheme(&mut self, scheme: Option<Scheme>) -> Result<(), Error> {
465        self.request
466            .set_scheme(scheme.map(|scheme| scheme.into()).as_ref())
467            .map_err(|()| Error::WasiInvalidValue("invalid scheme"))
468    }
469
470    pub fn set_authority(&mut self, authority: Option<&str>) -> Result<(), Error> {
471        self.request
472            .set_authority(authority)
473            .map_err(|()| Error::WasiInvalidValue("invalid authority"))
474    }
475}
476
477impl<Registry> OutgoingRequest<Registry>
478where
479    Registry: PollableRegistry,
480{
481    pub fn send(
482        self,
483        options: Option<types::RequestOptions>,
484    ) -> Result<ActiveOutgoingRequest<Registry>, Error> {
485        let Self { request, body } = self;
486        let response = wasi::http::outgoing_handler::handle(request, options)
487            .map_err(Error::wasi_error_code)?;
488        let inner = Subscribable::new(response, body.registry().clone());
489        let future_response = FutureIncomingResponse { inner };
490        Ok(ActiveOutgoingRequest {
491            body,
492            future_response,
493        })
494    }
495}
496
497pub struct ActiveOutgoingRequest<Registry>
498where
499    Registry: PollableRegistry,
500{
501    body: OutgoingBody<Registry>,
502    future_response: FutureIncomingResponse<Registry>,
503}
504
505impl<Registry> ActiveOutgoingRequest<Registry>
506where
507    Registry: PollableRegistry,
508{
509    pub fn body(&mut self) -> &mut OutgoingBody<Registry> {
510        &mut self.body
511    }
512
513    pub fn into_parts(self) -> (OutgoingBody<Registry>, FutureIncomingResponse<Registry>) {
514        (self.body, self.future_response)
515    }
516}
517
518impl<Registry> IntoFuture for ActiveOutgoingRequest<Registry>
519where
520    Registry: PollableRegistry,
521{
522    type Output = Result<IncomingResponse<Registry>, Error>;
523    type IntoFuture = FutureIncomingResponse<Registry>;
524
525    fn into_future(self) -> Self::IntoFuture {
526        self.future_response
527    }
528}
529
530pub struct FutureIncomingResponse<Registry: PollableRegistry> {
531    inner: Subscribable<types::FutureIncomingResponse, Registry>,
532}
533
534impl<Registry> Future for FutureIncomingResponse<Registry>
535where
536    Registry: PollableRegistry,
537{
538    type Output = Result<IncomingResponse<Registry>, Error>;
539
540    fn poll(mut self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
541        match self.inner.get() {
542            Some(Ok(res)) => {
543                let response = res.map_err(Error::wasi_error_code)?;
544                // FIXME: figure out proper type contraints to avoid this
545                let registry = self.inner.registry().clone();
546                Poll::Ready(Ok(IncomingResponse::new(response, registry.clone())?))
547            }
548            Some(Err(())) => Poll::Ready(Err(Error::WasiInvalidState(
549                "FutureIncomingResponse polled after completion",
550            ))),
551            None => {
552                self.inner.register_subscribe(cx);
553                Poll::Pending
554            }
555        }
556    }
557}
558
559pub struct OutgoingResponse<Registry: PollableRegistry> {
560    response: types::OutgoingResponse,
561    body: OutgoingBody<Registry>,
562}
563
564impl<Registry> OutgoingResponse<Registry>
565where
566    Registry: PollableRegistry,
567{
568    pub fn new(response: types::OutgoingResponse, registry: Registry) -> Result<Self, Error> {
569        let body = response
570            .body()
571            .map_err(|()| Error::WasiInvalidState("outgoing-response.body already called"))?;
572        let body = OutgoingBody::new(body, registry)?;
573        Ok(Self { response, body })
574    }
575
576    pub fn from_headers(headers: &FieldEntries, registry: Registry) -> Result<Self, Error> {
577        let fields = headers.try_into_fields()?;
578        let response = types::OutgoingResponse::new(fields);
579        Self::new(response, registry)
580    }
581
582    pub fn set_status_code(&mut self, status_code: u16) -> Result<(), Error> {
583        self.response
584            .set_status_code(status_code)
585            .map_err(|()| Error::WasiInvalidValue("invalid status code"))
586    }
587
588    pub fn body(&mut self) -> &mut OutgoingBody<Registry> {
589        &mut self.body
590    }
591
592    pub fn into_body(self) -> OutgoingBody<Registry> {
593        self.body
594    }
595
596    fn into_parts(self) -> (types::OutgoingResponse, OutgoingBody<Registry>) {
597        (self.response, self.body)
598    }
599}
600
601pub struct ResponseOutparam {
602    outparam: types::ResponseOutparam,
603}
604
605impl ResponseOutparam {
606    pub fn new(outparam: types::ResponseOutparam) -> Self {
607        Self { outparam }
608    }
609
610    pub fn set_response<Registry>(
611        self,
612        response: OutgoingResponse<Registry>,
613    ) -> OutgoingBody<Registry>
614    where
615        Registry: PollableRegistry,
616    {
617        let (wasi_response, body) = response.into_parts();
618        types::ResponseOutparam::set(self.outparam, Ok(wasi_response));
619        body
620    }
621
622    pub fn set_error(self, err: types::ErrorCode) {
623        types::ResponseOutparam::set(self.outparam, Err(err));
624    }
625}
626
627#[derive(Debug)]
628pub struct FieldEntries(Vec<(String, Vec<u8>)>);
629
630impl FieldEntries {
631    pub fn try_into_fields(&self) -> Result<types::Fields, Error> {
632        types::Fields::from_list(&self.0).map_err(|err| Error::WasiFieldsError(err.to_string()))
633    }
634}
635
636impl From<Vec<(String, Vec<u8>)>> for FieldEntries {
637    fn from(value: Vec<(String, Vec<u8>)>) -> Self {
638        Self(value)
639    }
640}
641
642impl IntoIterator for FieldEntries {
643    type Item = (String, Vec<u8>);
644
645    type IntoIter = std::vec::IntoIter<Self::Item>;
646
647    fn into_iter(self) -> Self::IntoIter {
648        self.0.into_iter()
649    }
650}
651
652impl From<types::Fields> for FieldEntries {
653    fn from(fields: types::Fields) -> Self {
654        Self(fields.entries())
655    }
656}