http_encoding/
coder.rs

1use core::{
2    pin::Pin,
3    task::{Context, Poll, ready},
4};
5
6use std::io;
7
8use bytes::Bytes;
9use http::HeaderMap;
10use http_body_alt::{Body, Frame, SizeHint};
11use pin_project_lite::pin_project;
12
13use super::error::CoderError;
14
15pin_project! {
16    /// A coder type that can be used for either encode or decode which determined by De type.
17    #[derive(Default)]
18    pub struct Coder<S, C = FeaturedCode>{
19        #[pin]
20        body: S,
21        coder: C,
22        trailers: Option<HeaderMap>
23    }
24}
25
26impl<S, C> Coder<S, C>
27where
28    S: Body,
29    C: Code<S::Data>,
30    S::Data: AsRef<[u8]>,
31{
32    /// Construct a new coder.
33    #[inline]
34    pub const fn new(body: S, coder: C) -> Self {
35        Self {
36            body,
37            coder,
38            trailers: None,
39        }
40    }
41
42    #[inline]
43    pub fn into_inner(self) -> S {
44        self.body
45    }
46}
47
48impl<S, C> Body for Coder<S, C>
49where
50    S: Body,
51    CoderError: From<S::Error>,
52    C: Code<S::Data>,
53{
54    type Data = C::Item;
55    type Error = CoderError;
56
57    fn poll_frame(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
58        let mut this = self.project();
59
60        while let Some(res) = ready!(this.body.as_mut().poll_frame(cx)) {
61            match res? {
62                Frame::Data(item) => {
63                    if let Some(item) = this.coder.code(item)? {
64                        return Poll::Ready(Some(Ok(Frame::Data(item))));
65                    }
66                }
67                Frame::Trailers(trailers) => {
68                    this.trailers.replace(trailers);
69                    break;
70                }
71            }
72        }
73
74        match this.coder.code_eof()? {
75            Some(res) => Poll::Ready(Some(Ok(Frame::Data(res)))),
76            None => Poll::Ready(this.trailers.take().map(|trailsers| Ok(Frame::Trailers(trailsers)))),
77        }
78    }
79
80    fn is_end_stream(&self) -> bool {
81        // forward is_end_stream to coder as it determines the data length after (de)compress.
82        self.coder.is_end_stream(&self.body) && self.trailers.is_none()
83    }
84
85    #[inline]
86    fn size_hint(&self) -> SizeHint {
87        // forward size_hint to coder as it determines the data length after (de)compress.
88        self.coder.size_hint(&self.body)
89    }
90}
91
92pub trait Code<T>: Sized {
93    type Item;
94
95    fn code(&mut self, item: T) -> io::Result<Option<Self::Item>>;
96
97    fn code_eof(&mut self) -> io::Result<Option<Self::Item>>;
98
99    /// A helper method for overriding associated input body's end_stream state.
100    /// by default it returns value the same as [`Body::is_end_stream`]'s default value.
101    /// in other word the default prediction is (de)compress can not hint if body has ended.
102    #[allow(unused_variables)]
103    #[inline]
104    fn is_end_stream(&self, body: &impl Body) -> bool {
105        false
106    }
107
108    /// A helper method for overriding associated input body's size_hint.
109    /// by default it returns value the same as [`SizeHint`]'s default value.
110    /// in other word the default prediction is (de)compress can not hint an exact size.
111    #[allow(unused_variables)]
112    #[inline]
113    fn size_hint(&self, body: &impl Body) -> SizeHint {
114        SizeHint::default()
115    }
116}
117
118/// coder serve as pass through that just forward items.
119pub struct NoOpCode;
120
121impl<T> Code<T> for NoOpCode
122where
123    T: AsRef<[u8]> + 'static,
124{
125    type Item = Bytes;
126
127    fn code(&mut self, item: T) -> io::Result<Option<Self::Item>> {
128        Ok(Some(
129            try_downcast_to_bytes(item).unwrap_or_else(|item| Bytes::copy_from_slice(item.as_ref())),
130        ))
131    }
132
133    #[inline]
134    fn code_eof(&mut self) -> io::Result<Option<Self::Item>> {
135        Ok(None)
136    }
137
138    // similar reason to size_hint.
139    #[inline]
140    fn is_end_stream(&self, body: &impl Body) -> bool {
141        body.is_end_stream()
142    }
143
144    // noop coder can take advantage of not doing any de/encoding work and hint the output body
145    // size. this would help downstream to infer the size of body and avoid going through
146    // transfer-encoding: chunked when possible.
147    #[inline]
148    fn size_hint(&self, body: &impl Body) -> SizeHint {
149        body.size_hint()
150    }
151}
152
153pub enum FeaturedCode {
154    NoOp(NoOpCode),
155    #[cfg(feature = "br")]
156    DecodeBr(super::brotli::Decoder),
157    #[cfg(feature = "br")]
158    EncodeBr(super::brotli::Encoder),
159    #[cfg(feature = "gz")]
160    DecodeGz(super::gzip::Decoder),
161    #[cfg(feature = "gz")]
162    EncodeGz(super::gzip::Encoder),
163    #[cfg(feature = "de")]
164    DecodeDe(super::deflate::Decoder),
165    #[cfg(feature = "de")]
166    EncodeDe(super::deflate::Encoder),
167    #[cfg(feature = "zs")]
168    DecodeZs(super::zstandard::Decoder),
169    #[cfg(feature = "zs")]
170    EncodeZs(super::zstandard::Encoder),
171}
172
173impl Default for FeaturedCode {
174    fn default() -> Self {
175        Self::NoOp(NoOpCode)
176    }
177}
178
179impl<T> Code<T> for FeaturedCode
180where
181    T: AsRef<[u8]> + 'static,
182{
183    type Item = Bytes;
184
185    fn code(&mut self, item: T) -> io::Result<Option<Self::Item>> {
186        match self {
187            Self::NoOp(coder) => coder.code(item),
188            #[cfg(feature = "br")]
189            Self::DecodeBr(coder) => coder.code(item),
190            #[cfg(feature = "br")]
191            Self::EncodeBr(coder) => coder.code(item),
192            #[cfg(feature = "gz")]
193            Self::DecodeGz(coder) => coder.code(item),
194            #[cfg(feature = "gz")]
195            Self::EncodeGz(coder) => coder.code(item),
196            #[cfg(feature = "de")]
197            Self::DecodeDe(coder) => coder.code(item),
198            #[cfg(feature = "de")]
199            Self::EncodeDe(coder) => coder.code(item),
200            #[cfg(feature = "zs")]
201            Self::DecodeZs(coder) => coder.code(item),
202            #[cfg(feature = "zs")]
203            Self::EncodeZs(coder) => coder.code(item),
204        }
205    }
206
207    fn code_eof(&mut self) -> io::Result<Option<Self::Item>> {
208        match self {
209            Self::NoOp(coder) => <NoOpCode as Code<T>>::code_eof(coder),
210            #[cfg(feature = "br")]
211            Self::DecodeBr(coder) => <super::brotli::Decoder as Code<T>>::code_eof(coder),
212            #[cfg(feature = "br")]
213            Self::EncodeBr(coder) => <super::brotli::Encoder as Code<T>>::code_eof(coder),
214            #[cfg(feature = "gz")]
215            Self::DecodeGz(coder) => <super::gzip::Decoder as Code<T>>::code_eof(coder),
216            #[cfg(feature = "gz")]
217            Self::EncodeGz(coder) => <super::gzip::Encoder as Code<T>>::code_eof(coder),
218            #[cfg(feature = "de")]
219            Self::DecodeDe(coder) => <super::deflate::Decoder as Code<T>>::code_eof(coder),
220            #[cfg(feature = "de")]
221            Self::EncodeDe(coder) => <super::deflate::Encoder as Code<T>>::code_eof(coder),
222            #[cfg(feature = "zs")]
223            Self::DecodeZs(coder) => <super::zstandard::Decoder as Code<T>>::code_eof(coder),
224            #[cfg(feature = "zs")]
225            Self::EncodeZs(coder) => <super::zstandard::Encoder as Code<T>>::code_eof(coder),
226        }
227    }
228
229    fn is_end_stream(&self, body: &impl Body) -> bool {
230        match self {
231            Self::NoOp(coder) => <NoOpCode as Code<T>>::is_end_stream(coder, body),
232            #[cfg(feature = "br")]
233            Self::DecodeBr(coder) => <super::brotli::Decoder as Code<T>>::is_end_stream(coder, body),
234            #[cfg(feature = "br")]
235            Self::EncodeBr(coder) => <super::brotli::Encoder as Code<T>>::is_end_stream(coder, body),
236            #[cfg(feature = "gz")]
237            Self::DecodeGz(coder) => <super::gzip::Decoder as Code<T>>::is_end_stream(coder, body),
238            #[cfg(feature = "gz")]
239            Self::EncodeGz(coder) => <super::gzip::Encoder as Code<T>>::is_end_stream(coder, body),
240            #[cfg(feature = "de")]
241            Self::DecodeDe(coder) => <super::deflate::Decoder as Code<T>>::is_end_stream(coder, body),
242            #[cfg(feature = "de")]
243            Self::EncodeDe(coder) => <super::deflate::Encoder as Code<T>>::is_end_stream(coder, body),
244            #[cfg(feature = "zs")]
245            Self::DecodeZs(coder) => <super::zstandard::Decoder as Code<T>>::is_end_stream(coder, body),
246            #[cfg(feature = "zs")]
247            Self::EncodeZs(coder) => <super::zstandard::Encoder as Code<T>>::is_end_stream(coder, body),
248        }
249    }
250
251    fn size_hint(&self, body: &impl Body) -> SizeHint {
252        match self {
253            Self::NoOp(coder) => <NoOpCode as Code<T>>::size_hint(coder, body),
254            #[cfg(feature = "br")]
255            Self::DecodeBr(coder) => <super::brotli::Decoder as Code<T>>::size_hint(coder, body),
256            #[cfg(feature = "br")]
257            Self::EncodeBr(coder) => <super::brotli::Encoder as Code<T>>::size_hint(coder, body),
258            #[cfg(feature = "gz")]
259            Self::DecodeGz(coder) => <super::gzip::Decoder as Code<T>>::size_hint(coder, body),
260            #[cfg(feature = "gz")]
261            Self::EncodeGz(coder) => <super::gzip::Encoder as Code<T>>::size_hint(coder, body),
262            #[cfg(feature = "de")]
263            Self::DecodeDe(coder) => <super::deflate::Decoder as Code<T>>::size_hint(coder, body),
264            #[cfg(feature = "de")]
265            Self::EncodeDe(coder) => <super::deflate::Encoder as Code<T>>::size_hint(coder, body),
266            #[cfg(feature = "zs")]
267            Self::DecodeZs(coder) => <super::zstandard::Decoder as Code<T>>::size_hint(coder, body),
268            #[cfg(feature = "zs")]
269            Self::EncodeZs(coder) => <super::zstandard::Encoder as Code<T>>::size_hint(coder, body),
270        }
271    }
272}
273
274#[cfg(any(feature = "gz", feature = "de"))]
275macro_rules! code_impl {
276    ($coder: ident) => {
277        impl<T> crate::Code<T> for $coder<crate::writer::BytesMutWriter>
278        where
279            T: AsRef<[u8]>,
280        {
281            type Item = ::bytes::Bytes;
282
283            fn code(&mut self, item: T) -> ::std::io::Result<Option<Self::Item>> {
284                use ::std::io::Write;
285
286                self.write_all(item.as_ref())?;
287                let b = self.get_mut().take();
288                if !b.is_empty() { Ok(Some(b)) } else { Ok(None) }
289            }
290
291            fn code_eof(&mut self) -> ::std::io::Result<Option<Self::Item>> {
292                self.try_finish()?;
293                let b = self.get_mut().take();
294                if !b.is_empty() { Ok(Some(b)) } else { Ok(None) }
295            }
296        }
297    };
298}
299
300fn try_downcast_to_bytes<T: 'static>(item: T) -> Result<Bytes, T> {
301    use core::any::Any;
302
303    let item = &mut Some(item);
304    match (item as &mut dyn Any).downcast_mut::<Option<Bytes>>() {
305        Some(bytes) => Ok(bytes.take().unwrap()),
306        None => Err(item.take().unwrap()),
307    }
308}
309
310#[cfg(test)]
311mod test {
312    use super::*;
313
314    #[test]
315    fn downcast_bytes() {
316        let bytes = Bytes::new();
317        assert!(try_downcast_to_bytes(bytes).is_ok());
318        let bytes = Vec::<u8>::new();
319        assert!(try_downcast_to_bytes(bytes).is_err());
320    }
321}
http_encoding/coder.rs

http_encoding/
coder.rs
