Skip to main content

serverkit/
body.rs

1use crate::{RequestStream, StreamError, stream::LimitedRequestStream};
2
3pub struct Body {
4    stream: Box<dyn RequestStream>,
5}
6
7impl Body {
8    pub(crate) fn new(stream: Box<dyn RequestStream>, limit: Option<usize>) -> Self {
9        let stream = match limit {
10            Some(limit) => {
11                Box::new(LimitedRequestStream::new(stream, limit)) as Box<dyn RequestStream>
12            }
13            None => stream,
14        };
15
16        Self { stream }
17    }
18
19    pub async fn next(&mut self) -> Option<Result<&[u8], StreamError>> {
20        let result = std::future::poll_fn(|context| self.stream.poll_next(context)).await?;
21
22        match result {
23            Ok(()) => Some(Ok(self.stream.chunk())),
24            Err(error) => Some(Err(error)),
25        }
26    }
27}
28
29#[cfg(test)]
30mod tests {
31    use std::{
32        future::Future,
33        task::{Context, Poll, Waker},
34    };
35
36    use super::Body;
37    use crate::{RequestStream, StreamError};
38
39    struct OneChunk {
40        bytes: Vec<u8>,
41        sent: bool,
42    }
43
44    impl RequestStream for OneChunk {
45        fn poll_next(
46            &mut self,
47            _context: &mut Context<'_>,
48        ) -> Poll<Option<Result<(), StreamError>>> {
49            if self.sent {
50                Poll::Ready(None)
51            } else {
52                self.sent = true;
53                Poll::Ready(Some(Ok(())))
54            }
55        }
56
57        fn chunk(&self) -> &[u8] {
58            &self.bytes
59        }
60    }
61
62    fn block_on<F: Future>(future: F) -> F::Output {
63        let mut future = std::pin::pin!(future);
64        let waker = Waker::noop();
65        let mut context = Context::from_waker(waker);
66
67        loop {
68            match future.as_mut().poll(&mut context) {
69                Poll::Ready(output) => return output,
70                Poll::Pending => std::thread::yield_now(),
71            }
72        }
73    }
74
75    #[test]
76    fn borrows_the_adapter_chunk_without_copying() {
77        let bytes = b"chunk".to_vec();
78        let pointer = bytes.as_ptr();
79        let mut body = Body::new(Box::new(OneChunk { bytes, sent: false }), None);
80        let chunk = block_on(body.next()).unwrap().unwrap();
81
82        assert_eq!(chunk, b"chunk");
83        assert_eq!(chunk.as_ptr(), pointer);
84    }
85}