net/
decoder.rs

1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
4
5//! Adapted from an implementation in reqwest.
6
7/*!
8A non-blocking response decoder.
9
10The decoder wraps a stream of bytes and produces a new stream of decompressed bytes.
11The decompressed bytes aren't guaranteed to align to the compressed ones.
12
13If the response is plaintext then no additional work is carried out.
14Bytes are just passed along.
15
16If the response is gzip, deflate or brotli then the bytes are decompressed.
17*/
18
19use std::error::Error;
20use std::fmt;
21use std::io::{self};
22use std::pin::Pin;
23
24use async_compression::tokio::bufread::{BrotliDecoder, GzipDecoder, ZlibDecoder, ZstdDecoder};
25use bytes::Bytes;
26use futures::stream::Peekable;
27use futures::task::{Context, Poll};
28use futures::{Future, Stream};
29use futures_util::StreamExt;
30use headers::{ContentLength, HeaderMapExt};
31use http_body_util::BodyExt;
32use hyper::Response;
33use hyper::body::Body;
34use hyper::header::{CONTENT_ENCODING, HeaderValue, TRANSFER_ENCODING};
35use tokio_util::codec::{BytesCodec, FramedRead};
36use tokio_util::io::StreamReader;
37
38use crate::connector::BoxedBody;
39
40pub const DECODER_BUFFER_SIZE: usize = 8192;
41
42/// Marker wrapper for errors that originate from the network body stream
43#[derive(Debug)]
44pub struct BodyStreamError(pub Box<dyn Error + Send + Sync>);
45
46impl fmt::Display for BodyStreamError {
47    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
48        self.0.fmt(f)
49    }
50}
51
52impl Error for BodyStreamError {
53    fn source(&self) -> Option<&(dyn Error + 'static)> {
54        Some(self.0.as_ref())
55    }
56}
57
58/// Normalize errors produced by the decompressors to `ErrorKind::InvalidData`
59/// so that `http_loader` reports them as `NetworkError::DecompressionError`.
60pub fn map_decode_error(err: io::Error) -> io::Error {
61    if err.kind() == io::ErrorKind::InvalidData {
62        return err;
63    }
64
65    let mut source: Option<&(dyn Error + 'static)> = err.get_ref().map(|e| e as _);
66    while let Some(e) = source {
67        if e.is::<BodyStreamError>() {
68            return err;
69        }
70
71        source = match e.downcast_ref::<io::Error>() {
72            Some(io_error) => io_error.get_ref().map(|e| e as _),
73            None => e.source(),
74        };
75    }
76    io::Error::new(io::ErrorKind::InvalidData, err)
77}
78
79/// A response decompressor over a non-blocking stream of bytes.
80///
81/// The inner decoder may be constructed asynchronously.
82pub struct Decoder {
83    inner: Inner,
84}
85
86#[derive(PartialEq)]
87enum DecoderType {
88    Gzip,
89    Brotli,
90    Deflate,
91    Zstd,
92}
93
94enum Inner {
95    /// A `PlainText` decoder just returns the response content as is.
96    PlainText(BodyStream),
97    /// A `Gzip` decoder will uncompress the gzipped response content before returning it.
98    Gzip(FramedRead<GzipDecoder<StreamReader<Peekable<BodyStream>, Bytes>>, BytesCodec>),
99    /// A `Delfate` decoder will uncompress the inflated response content before returning it.
100    Deflate(FramedRead<ZlibDecoder<StreamReader<Peekable<BodyStream>, Bytes>>, BytesCodec>),
101    /// A `Brotli` decoder will uncompress the brotli-encoded response content before returning it.
102    Brotli(FramedRead<BrotliDecoder<StreamReader<Peekable<BodyStream>, Bytes>>, BytesCodec>),
103    /// A `Zstd` decoder will uncompress the zstd-encoded response content before returning it.
104    Zstd(FramedRead<ZstdDecoder<StreamReader<Peekable<BodyStream>, Bytes>>, BytesCodec>),
105    /// A decoder that doesn't have a value yet.
106    Pending(Pending),
107}
108
109/// A future attempt to poll the response body for EOF so we know whether to use gzip or not.
110struct Pending {
111    body: Peekable<BodyStream>,
112    type_: DecoderType,
113}
114
115impl fmt::Debug for Decoder {
116    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
117        f.debug_struct("Decoder").finish()
118    }
119}
120
121impl Decoder {
122    /// A plain text decoder.
123    ///
124    /// This decoder will emit the underlying bytes as-is.
125    #[inline]
126    fn plain_text(
127        body: BoxedBody,
128        is_secure_scheme: bool,
129        content_length: Option<ContentLength>,
130    ) -> Decoder {
131        Decoder {
132            inner: Inner::PlainText(BodyStream::new(body, is_secure_scheme, content_length)),
133        }
134    }
135
136    /// A pending decoder.
137    ///
138    /// This decoder will buffer and decompress bytes that are encoded in the expected format.
139    #[inline]
140    fn pending(
141        body: BoxedBody,
142        type_: DecoderType,
143        is_secure_scheme: bool,
144        content_length: Option<ContentLength>,
145    ) -> Decoder {
146        Decoder {
147            inner: Inner::Pending(Pending {
148                body: BodyStream::new(body, is_secure_scheme, content_length).peekable(),
149                type_,
150            }),
151        }
152    }
153
154    /// Constructs a Decoder from a hyper response.
155    ///
156    /// A decoder is just a wrapper around the hyper response that knows
157    /// how to decode the content body of the response.
158    ///
159    /// Uses the correct variant by inspecting the Content-Encoding header.
160    pub fn detect(response: Response<BoxedBody>, is_secure_scheme: bool) -> Response<Decoder> {
161        let values = response
162            .headers()
163            .get_all(CONTENT_ENCODING)
164            .iter()
165            .chain(response.headers().get_all(TRANSFER_ENCODING).iter());
166        let decoder = values.fold(None, |acc, enc| {
167            acc.or_else(|| {
168                if enc == HeaderValue::from_static("gzip") {
169                    Some(DecoderType::Gzip)
170                } else if enc == HeaderValue::from_static("br") {
171                    Some(DecoderType::Brotli)
172                } else if enc == HeaderValue::from_static("deflate") {
173                    Some(DecoderType::Deflate)
174                } else if enc == HeaderValue::from_static("zstd") {
175                    Some(DecoderType::Zstd)
176                } else {
177                    None
178                }
179            })
180        });
181        let content_length = response.headers().typed_get::<ContentLength>();
182        match decoder {
183            Some(type_) => {
184                response.map(|r| Decoder::pending(r, type_, is_secure_scheme, content_length))
185            },
186            None => response.map(|r| Decoder::plain_text(r, is_secure_scheme, content_length)),
187        }
188    }
189}
190
191impl Stream for Decoder {
192    type Item = Result<Bytes, io::Error>;
193
194    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
195        // Do a read or poll for a pending decoder value.
196        match self.inner {
197            Inner::Pending(ref mut future) => match futures_core::ready!(Pin::new(future).poll(cx))
198            {
199                Ok(inner) => {
200                    self.inner = inner;
201                    self.poll_next(cx)
202                },
203                Err(e) => Poll::Ready(Some(Err(e))),
204            },
205            Inner::PlainText(ref mut body) => Pin::new(body).poll_next(cx),
206            Inner::Gzip(ref mut decoder) => {
207                match futures_core::ready!(Pin::new(decoder).poll_next(cx)) {
208                    Some(Ok(bytes)) => Poll::Ready(Some(Ok(bytes.freeze()))),
209                    Some(Err(err)) => Poll::Ready(Some(Err(map_decode_error(err)))),
210                    None => Poll::Ready(None),
211                }
212            },
213            Inner::Brotli(ref mut decoder) => {
214                match futures_core::ready!(Pin::new(decoder).poll_next(cx)) {
215                    Some(Ok(bytes)) => Poll::Ready(Some(Ok(bytes.freeze()))),
216                    Some(Err(err)) => Poll::Ready(Some(Err(map_decode_error(err)))),
217                    None => Poll::Ready(None),
218                }
219            },
220            Inner::Deflate(ref mut decoder) => {
221                match futures_core::ready!(Pin::new(decoder).poll_next(cx)) {
222                    Some(Ok(bytes)) => Poll::Ready(Some(Ok(bytes.freeze()))),
223                    Some(Err(err)) => Poll::Ready(Some(Err(map_decode_error(err)))),
224                    None => Poll::Ready(None),
225                }
226            },
227            Inner::Zstd(ref mut decoder) => {
228                match futures_core::ready!(Pin::new(decoder).poll_next(cx)) {
229                    Some(Ok(bytes)) => Poll::Ready(Some(Ok(bytes.freeze()))),
230                    Some(Err(err)) => Poll::Ready(Some(Err(map_decode_error(err)))),
231                    None => Poll::Ready(None),
232                }
233            },
234        }
235    }
236}
237
238impl Future for Pending {
239    type Output = Result<Inner, io::Error>;
240
241    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
242        match futures_core::ready!(Pin::new(&mut self.body).poll_peek(cx)) {
243            Some(Ok(_)) => {
244                // fallthrough
245            },
246            Some(Err(_e)) => {
247                // error was just a ref, so we need to really poll to move it
248                return Poll::Ready(Err(futures_core::ready!(
249                    Pin::new(&mut self.body).poll_next(cx)
250                )
251                .expect("just peeked Some")
252                .unwrap_err()));
253            },
254            None => return Poll::Ready(Ok(Inner::PlainText(BodyStream::empty()))),
255        };
256
257        let body = std::mem::replace(&mut self.body, BodyStream::empty().peekable());
258
259        match self.type_ {
260            DecoderType::Brotli => Poll::Ready(Ok(Inner::Brotli(FramedRead::with_capacity(
261                BrotliDecoder::new(StreamReader::new(body)),
262                BytesCodec::new(),
263                DECODER_BUFFER_SIZE,
264            )))),
265            DecoderType::Gzip => Poll::Ready(Ok(Inner::Gzip(FramedRead::with_capacity(
266                GzipDecoder::new(StreamReader::new(body)),
267                BytesCodec::new(),
268                DECODER_BUFFER_SIZE,
269            )))),
270            DecoderType::Deflate => Poll::Ready(Ok(Inner::Deflate(FramedRead::with_capacity(
271                ZlibDecoder::new(StreamReader::new(body)),
272                BytesCodec::new(),
273                DECODER_BUFFER_SIZE,
274            )))),
275            DecoderType::Zstd => Poll::Ready(Ok(Inner::Zstd(FramedRead::with_capacity(
276                ZstdDecoder::new(StreamReader::new(body)),
277                BytesCodec::new(),
278                DECODER_BUFFER_SIZE,
279            )))),
280        }
281    }
282}
283
284struct BodyStream {
285    body: BoxedBody,
286    is_secure_scheme: bool,
287    content_length: Option<ContentLength>,
288    total_read: u64,
289}
290
291impl BodyStream {
292    fn empty() -> Self {
293        BodyStream {
294            body: http_body_util::Empty::new()
295                .map_err(|_| unreachable!())
296                .boxed(),
297            is_secure_scheme: false,
298            content_length: None,
299            total_read: 0,
300        }
301    }
302
303    fn new(body: BoxedBody, is_secure_scheme: bool, content_length: Option<ContentLength>) -> Self {
304        BodyStream {
305            body,
306            is_secure_scheme,
307            content_length,
308            total_read: 0,
309        }
310    }
311}
312
313impl Stream for BodyStream {
314    type Item = Result<Bytes, io::Error>;
315
316    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
317        match futures_core::ready!(Pin::new(&mut self.body).poll_frame(cx)) {
318            Some(Ok(bytes)) => {
319                let Ok(bytes) = bytes.into_data() else {
320                    return Poll::Ready(None);
321                };
322                self.total_read += bytes.len() as u64;
323                Poll::Ready(Some(Ok(bytes)))
324            },
325            Some(Err(err)) => {
326                // To prevent truncation attacks rustls treats close connection without a close_notify as
327                // an error of type std::io::Error with ErrorKind::UnexpectedEof.
328                // https://docs.rs/rustls/latest/rustls/manual/_03_howto/index.html#unexpected-eof
329                //
330                // The error can be safely ignored if we known that all content was received or is explicitly
331                // set in preferences.
332                let all_content_read = self.content_length.is_some_and(|c| c.0 == self.total_read);
333                if self.is_secure_scheme && all_content_read {
334                    let source = err.source();
335                    let is_unexpected_eof = source
336                        .and_then(|e| e.downcast_ref::<io::Error>())
337                        .is_some_and(|e| e.kind() == io::ErrorKind::UnexpectedEof);
338                    if is_unexpected_eof {
339                        return Poll::Ready(None);
340                    }
341                }
342                Poll::Ready(Some(Err(io::Error::other(BodyStreamError(err.into())))))
343            },
344            None => Poll::Ready(None),
345        }
346    }
347}
net/decoder.rs

net/
decoder.rs
