rama_http/layer/decompression/
body.rs

1#![allow(unused_imports)]
2
3use crate::HeaderMap;
4use crate::body::{Body, Frame, SizeHint, StreamingBody};
5use crate::layer::util::compression::{
6    AsyncReadBody, BodyIntoStream, CompressionLevel, DecorateAsyncRead, WrapBody,
7    compressed_body_poll_frame, impl_decorate_async_read,
8};
9use rama_core::error::BoxError;
10
11use async_compression::tokio::bufread::BrotliDecoder;
12use async_compression::tokio::bufread::GzipDecoder;
13use async_compression::tokio::bufread::ZlibDecoder;
14use async_compression::tokio::bufread::ZstdDecoder;
15use pin_project_lite::pin_project;
16use rama_core::bytes::{Buf, Bytes};
17use rama_core::futures::ready;
18use rama_core::stream::io::StreamReader;
19use std::task::Context;
20use std::{io, marker::PhantomData, pin::Pin, task::Poll};
21
22pin_project! {
23    /// Response body of [`RequestDecompression`] and [`Decompression`].
24    ///
25    /// [`RequestDecompression`]: super::RequestDecompression
26    /// [`Decompression`]: super::Decompression
27    pub struct DecompressionBody<B>
28    where
29        B: StreamingBody
30    {
31        #[pin]
32        pub(crate) inner: BodyInner<B>,
33    }
34}
35
36impl<B> Default for DecompressionBody<B>
37where
38    B: StreamingBody + Default,
39{
40    fn default() -> Self {
41        Self {
42            inner: BodyInner::Identity {
43                inner: B::default(),
44            },
45        }
46    }
47}
48
49impl<B> DecompressionBody<B>
50where
51    B: StreamingBody,
52{
53    pub(crate) fn new(inner: BodyInner<B>) -> Self {
54        Self { inner }
55    }
56}
57
58type GzipBody<B> = WrapBody<GzipDecoder<B>>;
59type DeflateBody<B> = WrapBody<ZlibDecoder<B>>;
60type BrotliBody<B> = WrapBody<BrotliDecoder<B>>;
61type ZstdBody<B> = WrapBody<ZstdDecoder<B>>;
62
63pin_project! {
64    #[project = BodyInnerProj]
65    pub(crate) enum BodyInner<B>
66    where
67        B: StreamingBody,
68    {
69        Gzip {
70            #[pin]
71            inner: GzipBody<B>,
72        },
73        Deflate {
74            #[pin]
75            inner: DeflateBody<B>,
76        },
77        Brotli {
78            #[pin]
79            inner: BrotliBody<B>,
80        },
81        Zstd {
82            #[pin]
83            inner: ZstdBody<B>,
84        },
85        Identity {
86            #[pin]
87            inner: B,
88        },
89    }
90}
91
92impl<B: StreamingBody> BodyInner<B> {
93    pub(crate) fn gzip(inner: WrapBody<GzipDecoder<B>>) -> Self {
94        Self::Gzip { inner }
95    }
96
97    pub(crate) fn deflate(inner: WrapBody<ZlibDecoder<B>>) -> Self {
98        Self::Deflate { inner }
99    }
100
101    pub(crate) fn brotli(inner: WrapBody<BrotliDecoder<B>>) -> Self {
102        Self::Brotli { inner }
103    }
104
105    pub(crate) fn zstd(inner: WrapBody<ZstdDecoder<B>>) -> Self {
106        Self::Zstd { inner }
107    }
108
109    pub(crate) fn identity(inner: B) -> Self {
110        Self::Identity { inner }
111    }
112}
113
114impl<B> StreamingBody for DecompressionBody<B>
115where
116    B: StreamingBody<Error: Into<BoxError>>,
117{
118    type Data = Bytes;
119    type Error = BoxError;
120
121    fn poll_frame(
122        self: Pin<&mut Self>,
123        cx: &mut Context<'_>,
124    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
125        compressed_body_poll_frame!(self, cx)
126    }
127
128    fn size_hint(&self) -> SizeHint {
129        match self.inner {
130            BodyInner::Identity { ref inner } => inner.size_hint(),
131            BodyInner::Brotli { .. }
132            | BodyInner::Gzip { .. }
133            | BodyInner::Deflate { .. }
134            | BodyInner::Zstd { .. } => SizeHint::default(),
135        }
136    }
137}
138
139impl_decorate_async_read!(GzipDecoder: |input, _quality| {
140    GzipDecoder::new(input)
141});
142
143impl_decorate_async_read!(ZlibDecoder: |input, _quality| {
144    ZlibDecoder::new(input)
145});
146
147impl_decorate_async_read!(BrotliDecoder: |input, _quality| {
148    BrotliDecoder::new(input)
149});
150
151impl_decorate_async_read!(ZstdDecoder: |input, _quality| {
152    let mut decoder = ZstdDecoder::new(input);
153    decoder.multiple_members(true);
154    decoder
155});
rama_http/layer/decompression/body.rs

rama_http/layer/decompression/
body.rs
