unb_runtime/
body.rs

1use std::future::poll_fn;
2use std::pin::Pin;
3
4use bytes::{Bytes, BytesMut};
5use futures_util::Stream;
6use unb_core::CoreError;
7
8pub type BodyStream = Pin<Box<dyn Stream<Item = Result<Bytes, CoreError>> + Send>>;
9
10pub enum WireBody {
11    Bytes(Bytes),
12    Stream(BodyStream),
13}
14
15impl WireBody {
16    pub async fn next_chunk(&mut self) -> Option<Result<Bytes, CoreError>> {
17        match self {
18            WireBody::Bytes(bytes) if bytes.is_empty() => None,
19            WireBody::Bytes(bytes) => Some(Ok(std::mem::take(bytes))),
20            WireBody::Stream(stream) => poll_fn(|cx| stream.as_mut().poll_next(cx)).await,
21        }
22    }
23
24    pub async fn collect_to(self, ceiling: usize) -> Result<Bytes, CoreError> {
25        match self {
26            WireBody::Bytes(bytes) => {
27                if bytes.len() > ceiling {
28                    return Err(CoreError::BodyTooLarge(ceiling));
29                }
30                Ok(bytes)
31            }
32            WireBody::Stream(mut stream) => {
33                let mut collected = BytesMut::new();
34                loop {
35                    match poll_fn(|cx| stream.as_mut().poll_next(cx)).await {
36                        Some(Ok(chunk)) => {
37                            if collected.len() + chunk.len() > ceiling {
38                                return Err(CoreError::BodyTooLarge(ceiling));
39                            }
40                            collected.extend_from_slice(&chunk);
41                        }
42                        Some(Err(error)) => return Err(error),
43                        None => return Ok(collected.freeze()),
44                    }
45                }
46            }
47        }
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54
55    fn chunked(chunks: Vec<Result<Bytes, CoreError>>) -> WireBody {
56        WireBody::Stream(Box::pin(futures_util::stream::iter(chunks)))
57    }
58
59    #[test]
60    fn collected_bytes_below_the_ceiling_are_identity() {
61        let body = WireBody::Bytes(Bytes::from_static(b"hello"));
62        let collected = futures_executor::block_on(body.collect_to(16)).unwrap();
63        assert_eq!(collected, Bytes::from_static(b"hello"));
64    }
65
66    #[test]
67    fn a_stream_collects_chunks_in_order_below_the_ceiling() {
68        let body = chunked(vec![
69            Ok(Bytes::from_static(b"hel")),
70            Ok(Bytes::from_static(b"lo")),
71        ]);
72        let collected = futures_executor::block_on(body.collect_to(16)).unwrap();
73        assert_eq!(collected, Bytes::from_static(b"hello"));
74    }
75
76    #[test]
77    fn unary_and_streaming_bodies_share_the_chunk_api() {
78        let mut unary = WireBody::Bytes(Bytes::from_static(b"one"));
79        assert_eq!(
80            futures_executor::block_on(unary.next_chunk())
81                .unwrap()
82                .unwrap(),
83            Bytes::from_static(b"one")
84        );
85        assert!(futures_executor::block_on(unary.next_chunk()).is_none());
86
87        let mut streaming = chunked(vec![
88            Ok(Bytes::from_static(b"two")),
89            Ok(Bytes::from_static(b"three")),
90        ]);
91        assert_eq!(
92            futures_executor::block_on(streaming.next_chunk())
93                .unwrap()
94                .unwrap(),
95            Bytes::from_static(b"two")
96        );
97        assert_eq!(
98            futures_executor::block_on(streaming.next_chunk())
99                .unwrap()
100                .unwrap(),
101            Bytes::from_static(b"three")
102        );
103        assert!(futures_executor::block_on(streaming.next_chunk()).is_none());
104    }
105
106    #[test]
107    fn collection_above_the_ceiling_fails_body_too_large() {
108        let oversized = WireBody::Bytes(Bytes::from_static(b"toolarge"));
109        assert!(matches!(
110            futures_executor::block_on(oversized.collect_to(4)),
111            Err(CoreError::BodyTooLarge(4))
112        ));
113        let streamed = chunked(vec![
114            Ok(Bytes::from_static(b"too")),
115            Ok(Bytes::from_static(b"large")),
116        ]);
117        assert!(matches!(
118            futures_executor::block_on(streamed.collect_to(4)),
119            Err(CoreError::BodyTooLarge(4))
120        ));
121    }
122}
unb_runtime/body.rs

unb_runtime/
body.rs
