wasi_hyperium/
incoming.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use bytes::Bytes;
8use wasi::http::types;
9
10use crate::{
11    poll::PollableRegistry,
12    wasi::{FieldEntries, FutureTrailers, IncomingBody},
13    Error,
14};
15
16pub struct IncomingHttpBody<Registry>
17where
18    Registry: PollableRegistry,
19{
20    pub(crate) state: IncomingState<Registry>,
21}
22
23pub(crate) enum IncomingState<Registry>
24where
25    Registry: PollableRegistry,
26{
27    Empty,
28    Body(IncomingBody<Registry>),
29    Trailers(FutureTrailers<Registry>),
30}
31
32const READ_FRAME_SIZE: usize = 16 * 1024;
33
34impl<Registry> IncomingHttpBody<Registry>
35where
36    Registry: PollableRegistry,
37{
38    pub fn new(body: types::IncomingBody, registry: Registry) -> Result<Self, Error> {
39        Ok(IncomingBody::new(body, registry)?.into())
40    }
41
42    pub fn poll_incoming_body(&mut self, cx: &mut Context) -> Poll<Option<Result<Bytes, Error>>> {
43        let IncomingState::Body(incoming_body) = &mut self.state else {
44            panic!("poll_incoming_body called on non-body state")
45        };
46
47        match incoming_body.stream().poll_read(cx, READ_FRAME_SIZE) {
48            Poll::Ready(Ok(data)) => Poll::Ready(Some(Ok(data.into()))),
49            Poll::Ready(Err(Error::WasiStreamClosed)) => {
50                self.state = IncomingState::Trailers(self.take_body().finish());
51                Poll::Ready(None)
52            }
53            Poll::Ready(Err(err)) => Poll::Ready(Some(Err(err))),
54            Poll::Pending => Poll::Pending,
55        }
56    }
57
58    pub fn poll_incoming_trailers(
59        &mut self,
60        cx: &mut std::task::Context<'_>,
61    ) -> Poll<Result<Option<FieldEntries>, Error>> {
62        match &mut self.state {
63            IncomingState::Empty => Poll::Ready(Ok(None)),
64            IncomingState::Body { .. } => panic!("poll_trailers called before body completion"),
65            IncomingState::Trailers(trailers) => match Pin::new(trailers).poll(cx) {
66                Poll::Ready(Ok(Some(trailers))) => {
67                    self.state = IncomingState::Empty;
68                    Poll::Ready(Ok(Some(trailers)))
69                }
70                Poll::Ready(Ok(None)) => {
71                    self.state = IncomingState::Empty;
72                    Poll::Ready(Ok(None))
73                }
74                // TODO: figure out why this is happening
75                Poll::Ready(Err(Error::WasiErrorCode(s))) if s.contains("ConnectionTerminated") => {
76                    self.state = IncomingState::Empty;
77                    Poll::Ready(Ok(None))
78                }
79                Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
80                Poll::Pending => Poll::Pending,
81            },
82        }
83    }
84
85    pub(crate) fn take_body(&mut self) -> IncomingBody<Registry> {
86        match std::mem::replace(&mut self.state, IncomingState::Empty) {
87            IncomingState::Body(body) => body,
88            _ => panic!("called take_body on non-body state"),
89        }
90    }
91}
92
93impl<Registry> From<IncomingBody<Registry>> for IncomingHttpBody<Registry>
94where
95    Registry: PollableRegistry,
96{
97    fn from(body: IncomingBody<Registry>) -> Self {
98        Self {
99            state: IncomingState::Body(body),
100        }
101    }
102}
wasi_hyperium/incoming.rs

wasi_hyperium/
incoming.rs
