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}