simple_hyper_client/blocking/
body.rs

1/* Copyright (c) Fortanix, Inc.
2 *
3 * This Source Code Form is subject to the terms of the Mozilla Public
4 * License, v. 2.0. If a copy of the MPL was not distributed with this
5 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
6
7use super::client::KeepClientAlive;
8
9use futures_util::StreamExt;
10use http_body_util::BodyExt;
11use hyper::body::Body as HyperBody;
12use hyper::body::{Buf, Bytes};
13use tokio::sync::mpsc;
14
15use std::future::Future;
16use std::{fmt, io};
17
18/// A body type for HTTP responses that implement `std::io::Read`
19pub struct Body {
20    pub(super) keep_client_alive: KeepClientAlive,
21    bytes: Bytes,
22    rx: mpsc::Receiver<io::Result<Bytes>>,
23}
24
25impl fmt::Debug for Body {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        f.debug_struct("Body").finish()
28    }
29}
30
31impl Body {
32    pub(super) fn new<T>(hyper_body: T) -> (impl Future<Output = ()> + Send + 'static, Self)
33    where
34        T: HyperBody + Send + 'static,
35        T::Data: Send + Into<Bytes>,
36        T::Error: std::error::Error + Send + Sync + 'static,
37    {
38        let (tx, rx) = mpsc::channel(1);
39
40        let fut = async move {
41            let mut stream = std::pin::pin!(hyper_body.into_data_stream());
42
43            loop {
44                tokio::select! {
45                    _ = tx.closed() => {
46                        break; // body has been dropped.
47                    }
48
49                    res = stream.next() => {
50                        let res = match res {
51                            None => break, // EOF
52
53                            Some(Ok(chunk)) => {
54                                let chunk = chunk.into();
55
56                                if chunk.is_empty() {
57                                    continue;
58                                } else {
59                                    Ok(chunk)
60                                }
61                            }
62
63                            Some(Err(e)) => Err(io::Error::new(io::ErrorKind::Other, e)),
64                        };
65
66                        if tx.send(res).await.is_err() {
67                            break; // body has been dropped.
68                        }
69                    }
70                }
71            }
72        };
73
74        let body = Body {
75            keep_client_alive: KeepClientAlive::empty(),
76            bytes: Bytes::new(),
77            rx,
78        };
79
80        (fut, body)
81    }
82}
83
84impl io::Read for Body {
85    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
86        if self.bytes.is_empty() {
87            match self.rx.blocking_recv() {
88                Some(Ok(bytes)) => {
89                    self.bytes = bytes;
90                }
91                Some(Err(e)) => return Err(e),
92                None => return Ok(0),
93            }
94        }
95        (&mut self.bytes).reader().read(buf)
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use http_body_util::{Channel, Full, StreamBody};
103    use hyper::body::Frame;
104    use std::io::Read;
105    use std::{future::Future, thread};
106    use tokio::time::{self, Duration};
107
108    fn run_future<F: Future<Output = ()> + Send + 'static>(fut: F) {
109        thread::spawn(move || {
110            let rt = tokio::runtime::Builder::new_current_thread()
111                .enable_all()
112                .build()
113                .unwrap();
114            rt.block_on(fut);
115        });
116    }
117
118    #[test]
119    fn single_chunk() {
120        let body = Full::new(&b"hello, world!"[..]);
121        let (fut, mut reader) = Body::new(body);
122        run_future(fut);
123
124        let mut bytes = Vec::<u8>::new();
125        reader.read_to_end(&mut bytes).unwrap();
126        assert_eq!(bytes, b"hello, world!");
127    }
128
129    #[test]
130    fn multiple_chunks() {
131        let (mut sender, body) = Channel::<&[u8]>::new(5);
132        let (fut, mut reader) = Body::new(body);
133
134        run_future(async move {
135            let h = tokio::spawn(fut);
136
137            sender.send_data(b"hello").await.unwrap();
138            time::sleep(Duration::from_millis(10)).await;
139            sender.send_data(b", ").await.unwrap();
140            sender.send_data(b"world!").await.unwrap();
141
142            drop(sender);
143            h.await.unwrap();
144        });
145
146        let mut bytes = Vec::<u8>::new();
147        reader.read_to_end(&mut bytes).unwrap();
148        assert_eq!(bytes, b"hello, world!");
149    }
150
151    #[test]
152    fn with_empty_chunk() {
153        let (mut sender, body) = Channel::<&[u8]>::new(5);
154        let (fut, mut reader) = Body::new(body);
155
156        run_future(async move {
157            let h = tokio::spawn(fut);
158
159            sender.send_data(b"hello").await.unwrap();
160            time::sleep(Duration::from_millis(10)).await;
161            sender.send_data(b"").await.unwrap();
162            sender.send_data(b", world!").await.unwrap();
163
164            drop(sender);
165            h.await.unwrap();
166        });
167
168        let mut bytes = Vec::<u8>::new();
169        reader.read_to_end(&mut bytes).unwrap();
170        assert_eq!(bytes, b"hello, world!");
171    }
172
173    #[test]
174    fn hyper_error() {
175        eprintln!("test");
176        let chunks: Vec<Result<&[u8], io::Error>> = vec![
177            Ok(b"hello"),
178            Ok(b" "),
179            Ok(b"world"),
180            Err(io::ErrorKind::BrokenPipe.into()),
181        ];
182        let chunks = chunks.into_iter().map(|res| res.map(Frame::data));
183        let stream = futures_util::stream::iter(chunks);
184        let body = StreamBody::new(stream);
185        let (fut, mut reader) = Body::new(body);
186
187        run_future(fut);
188
189        let mut bytes = Vec::<u8>::new();
190        let err = reader.read_to_end(&mut bytes).unwrap_err();
191        assert_eq!(err.kind(), io::ErrorKind::Other);
192        assert_eq!(bytes, b"hello world");
193
194        let mut buf = [0u8; 8];
195        let n = reader.read(&mut buf).unwrap();
196        assert_eq!(n, 0);
197    }
198}
simple_hyper_client/blocking/body.rs

simple_hyper_client/blocking/
body.rs
