Skip to main content

churust_core/
body.rs

1//! The response [`Body`]: either a fully-buffered `Bytes` payload or a lazy
2//! stream of byte chunks (for files, large/dynamic responses, SSE, etc.).
3
4use crate::error::{Error, Result};
5use bytes::{Bytes, BytesMut};
6use futures_util::stream::{Stream, StreamExt};
7use std::pin::Pin;
8
9/// A response body.
10///
11/// Construct buffered bodies with `From` (`Body::from(bytes)`,
12/// `"text".into()`), an empty body with [`Body::empty`], or a streaming body
13/// with [`Body::from_stream`]. The engine sends buffered bodies in one shot and
14/// streamed bodies frame-by-frame.
15pub enum Body {
16    /// A fully-buffered, in-memory body.
17    Bytes(Bytes),
18    /// A lazily-produced stream of byte chunks.
19    Stream(Pin<Box<dyn Stream<Item = Result<Bytes>> + Send + 'static>>),
20}
21
22impl Body {
23    /// An empty buffered body.
24    pub fn empty() -> Self {
25        Body::Bytes(Bytes::new())
26    }
27
28    /// Build a streaming body from any `Stream` of byte chunks. Chunk errors are
29    /// converted into [`Error`] and surface when the body is read/collected.
30    pub fn from_stream<S, E>(stream: S) -> Self
31    where
32        S: Stream<Item = std::result::Result<Bytes, E>> + Send + 'static,
33        E: std::fmt::Display,
34    {
35        let mapped =
36            stream.map(|chunk| chunk.map_err(|e| Error::internal(format!("body stream: {e}"))));
37        Body::Stream(Box::pin(mapped))
38    }
39
40    /// Borrow the buffered bytes, or `None` if this body is a stream. Used by
41    /// middleware that only post-processes already-materialized bodies.
42    pub fn as_bytes(&self) -> Option<&Bytes> {
43        match self {
44            Body::Bytes(b) => Some(b),
45            Body::Stream(_) => None,
46        }
47    }
48
49    /// Borrow the buffered bytes as a slice, or `None` for a stream.
50    pub fn as_slice(&self) -> Option<&[u8]> {
51        self.as_bytes().map(|b| b.as_ref())
52    }
53
54    /// True if this is a buffered, empty body. A stream is never reported empty
55    /// (its length is unknown until read).
56    pub fn is_empty(&self) -> bool {
57        matches!(self, Body::Bytes(b) if b.is_empty())
58    }
59
60    /// Collect the whole body into `Bytes` (returns the buffer directly when
61    /// already buffered, or drains the stream).
62    pub async fn into_bytes(self) -> Result<Bytes> {
63        match self {
64            Body::Bytes(b) => Ok(b),
65            Body::Stream(mut s) => {
66                let mut buf = BytesMut::new();
67                while let Some(chunk) = s.next().await {
68                    buf.extend_from_slice(&chunk?);
69                }
70                Ok(buf.freeze())
71            }
72        }
73    }
74}
75
76impl Default for Body {
77    fn default() -> Self {
78        Body::empty()
79    }
80}
81
82impl From<Bytes> for Body {
83    fn from(b: Bytes) -> Self {
84        Body::Bytes(b)
85    }
86}
87impl From<Vec<u8>> for Body {
88    fn from(b: Vec<u8>) -> Self {
89        Body::Bytes(Bytes::from(b))
90    }
91}
92impl From<String> for Body {
93    fn from(s: String) -> Self {
94        Body::Bytes(Bytes::from(s))
95    }
96}
97impl From<&'static str> for Body {
98    fn from(s: &'static str) -> Self {
99        Body::Bytes(Bytes::from_static(s.as_bytes()))
100    }
101}
102
103/// Compare a body to `Bytes` (true only when buffered and equal). Lets tests and
104/// callers assert on buffered bodies ergonomically.
105impl PartialEq<Bytes> for Body {
106    fn eq(&self, other: &Bytes) -> bool {
107        matches!(self, Body::Bytes(b) if b == other)
108    }
109}
110
111impl std::fmt::Debug for Body {
112    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
113        match self {
114            Body::Bytes(b) => f.debug_tuple("Body::Bytes").field(b).finish(),
115            Body::Stream(_) => f.write_str("Body::Stream(..)"),
116        }
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    #[tokio::test]
125    async fn buffered_round_trips() {
126        let body = Body::from(Bytes::from("hello"));
127        assert_eq!(body.as_slice(), Some(&b"hello"[..]));
128        assert!(!body.is_empty());
129        assert_eq!(body.into_bytes().await.unwrap(), Bytes::from("hello"));
130    }
131
132    #[tokio::test]
133    async fn empty_is_empty() {
134        assert!(Body::empty().is_empty());
135        assert_eq!(Body::empty().into_bytes().await.unwrap(), Bytes::new());
136    }
137
138    #[tokio::test]
139    async fn stream_collects_and_has_no_bytes_view() {
140        let chunks = futures_util::stream::iter(vec![
141            Ok::<_, std::io::Error>(Bytes::from("ab")),
142            Ok(Bytes::from("cd")),
143        ]);
144        let body = Body::from_stream(chunks);
145        assert!(body.as_bytes().is_none()); // streams expose no borrowed view
146        assert!(!body.is_empty());
147        assert_eq!(body.into_bytes().await.unwrap(), Bytes::from("abcd"));
148    }
149
150    #[test]
151    #[allow(clippy::cmp_owned)] // exercises From<String>; the owned String is the input under test, not a throwaway
152    fn partial_eq_bytes() {
153        assert!(Body::from("x".to_string()) == Bytes::from("x"));
154        let s = Body::from_stream(futures_util::stream::iter(vec![Ok::<_, std::io::Error>(
155            Bytes::new(),
156        )]));
157        assert!(!(s == Bytes::new())); // a stream never equals bytes
158    }
159}