actix_http/body/
sized_stream.rs

1use std::{
2    error::Error as StdError,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use bytes::Bytes;
8use futures_core::{ready, Stream};
9use pin_project_lite::pin_project;
10
11use super::{BodySize, MessageBody};
12
13pin_project! {
14    /// Known sized streaming response wrapper.
15    ///
16    /// This body implementation should be used if total size of stream is known. Data is sent as-is
17    /// without using chunked transfer encoding.
18    pub struct SizedStream<S> {
19        size: u64,
20        #[pin]
21        stream: S,
22    }
23}
24
25impl<S, E> SizedStream<S>
26where
27    S: Stream<Item = Result<Bytes, E>>,
28    E: Into<Box<dyn StdError>> + 'static,
29{
30    #[inline]
31    pub fn new(size: u64, stream: S) -> Self {
32        SizedStream { size, stream }
33    }
34}
35
36// TODO: from_infallible method
37
38impl<S, E> MessageBody for SizedStream<S>
39where
40    S: Stream<Item = Result<Bytes, E>>,
41    E: Into<Box<dyn StdError>> + 'static,
42{
43    type Error = E;
44
45    #[inline]
46    fn size(&self) -> BodySize {
47        BodySize::Sized(self.size)
48    }
49
50    /// Attempts to pull out the next value of the underlying [`Stream`].
51    ///
52    /// Empty values are skipped to prevent [`SizedStream`]'s transmission being
53    /// ended on a zero-length chunk, but rather proceed until the underlying
54    /// [`Stream`] ends.
55    fn poll_next(
56        mut self: Pin<&mut Self>,
57        cx: &mut Context<'_>,
58    ) -> Poll<Option<Result<Bytes, Self::Error>>> {
59        loop {
60            let stream = self.as_mut().project().stream;
61
62            let chunk = match ready!(stream.poll_next(cx)) {
63                Some(Ok(ref bytes)) if bytes.is_empty() => continue,
64                val => val,
65            };
66
67            return Poll::Ready(chunk);
68        }
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use std::{convert::Infallible, pin::pin};
75
76    use actix_utils::future::poll_fn;
77    use futures_util::stream;
78    use static_assertions::{assert_impl_all, assert_not_impl_any};
79
80    use super::*;
81    use crate::body::to_bytes;
82
83    assert_impl_all!(SizedStream<stream::Empty<Result<Bytes, crate::Error>>>: MessageBody);
84    assert_impl_all!(SizedStream<stream::Empty<Result<Bytes, &'static str>>>: MessageBody);
85    assert_impl_all!(SizedStream<stream::Repeat<Result<Bytes, &'static str>>>: MessageBody);
86    assert_impl_all!(SizedStream<stream::Empty<Result<Bytes, Infallible>>>: MessageBody);
87    assert_impl_all!(SizedStream<stream::Repeat<Result<Bytes, Infallible>>>: MessageBody);
88
89    assert_not_impl_any!(SizedStream<stream::Empty<Bytes>>: MessageBody);
90    assert_not_impl_any!(SizedStream<stream::Repeat<Bytes>>: MessageBody);
91    // crate::Error is not Clone
92    assert_not_impl_any!(SizedStream<stream::Repeat<Result<Bytes, crate::Error>>>: MessageBody);
93
94    #[actix_rt::test]
95    async fn skips_empty_chunks() {
96        let body = SizedStream::new(
97            2,
98            stream::iter(
99                ["1", "", "2"]
100                    .iter()
101                    .map(|&v| Ok::<_, Infallible>(Bytes::from(v))),
102            ),
103        );
104
105        let mut body = pin!(body);
106
107        assert_eq!(
108            poll_fn(|cx| body.as_mut().poll_next(cx))
109                .await
110                .unwrap()
111                .ok(),
112            Some(Bytes::from("1")),
113        );
114
115        assert_eq!(
116            poll_fn(|cx| body.as_mut().poll_next(cx))
117                .await
118                .unwrap()
119                .ok(),
120            Some(Bytes::from("2")),
121        );
122    }
123
124    #[actix_rt::test]
125    async fn read_to_bytes() {
126        let body = SizedStream::new(
127            2,
128            stream::iter(
129                ["1", "", "2"]
130                    .iter()
131                    .map(|&v| Ok::<_, Infallible>(Bytes::from(v))),
132            ),
133        );
134
135        assert_eq!(to_bytes(body).await.ok(), Some(Bytes::from("12")));
136    }
137
138    #[actix_rt::test]
139    async fn stream_string_error() {
140        // `&'static str` does not impl `Error`
141        // but it does impl `Into<Box<dyn Error>>`
142
143        let body = SizedStream::new(0, stream::once(async { Err("stringy error") }));
144        assert_eq!(to_bytes(body).await, Ok(Bytes::new()));
145
146        let body = SizedStream::new(1, stream::once(async { Err("stringy error") }));
147        assert!(matches!(to_bytes(body).await, Err("stringy error")));
148    }
149
150    #[actix_rt::test]
151    async fn stream_boxed_error() {
152        // `Box<dyn Error>` does not impl `Error`
153        // but it does impl `Into<Box<dyn Error>>`
154
155        let body = SizedStream::new(
156            0,
157            stream::once(async { Err(Box::<dyn StdError>::from("stringy error")) }),
158        );
159        assert_eq!(to_bytes(body).await.unwrap(), Bytes::new());
160
161        let body = SizedStream::new(
162            1,
163            stream::once(async { Err(Box::<dyn StdError>::from("stringy error")) }),
164        );
165        assert_eq!(
166            to_bytes(body).await.unwrap_err().to_string(),
167            "stringy error"
168        );
169    }
170}
actix_http/body/sized_stream.rs

actix_http/body/
sized_stream.rs
