async_codec_util/decoder/
decode_exact.rs

1use std::error::Error;
2use std::fmt::{self, Debug, Display, Formatter};
3
4use async_codec::{AsyncDecode, PollDec, DecodeError};
5use async_codec::PollDec::{Done, Progress, Pending, Errored};
6use atm_io_utils::limited_reader::LimitedReader;
7use futures_core::task::Context;
8use futures_io::AsyncRead;
9
10/// The error of a `DecodeExact`.
11///
12/// The inner decoder is limited in the number of bytes it can read. It errors with `UnexpectedEof`
13/// when it tries to read more than the allowed amount of bytes.
14#[derive(Debug)]
15pub enum DecodeExactError<E, I> {
16    /// The inner decoder finished decoding too early, after the contained number of bytes.
17    /// This error also contains the decoded item.
18    Early(I, usize),
19    /// The inner decoder errored.
20    Inner(E),
21}
22
23impl<E: Display, I> Display for DecodeExactError<E, I> {
24    fn fmt(&self, f: &mut Formatter) -> Result<(), fmt::Error> {
25        match *self {
26            DecodeExactError::Early(_, read) => {
27                write!(f, "Decoder finished early after reading {} bytes", read)
28            }
29            DecodeExactError::Inner(ref err) => write!(f, "Decode exact inner error: {}", err),
30        }
31    }
32}
33
34impl<E: Error, I: Debug> Error for DecodeExactError<E, I> {
35    fn description(&self) -> &str {
36        match *self {
37            DecodeExactError::Early(_, _) => "decoder finished early",
38            DecodeExactError::Inner(ref err) => err.description(),
39        }
40    }
41
42    fn cause(&self) -> Option<&Error> {
43        match *self {
44            DecodeExactError::Early(_, _) => None,
45            DecodeExactError::Inner(ref err) => Some(err),
46        }
47    }
48}
49
50impl<E, I> From<E> for DecodeExactError<E, I> {
51    fn from(err: E) -> DecodeExactError<E, I> {
52        DecodeExactError::Inner(err)
53    }
54}
55
56/// Wraps a decoder and emits an error if it does not stop encoding at exactly the specified byte
57/// count.
58pub struct DecodeExact<D> {
59    target: usize,
60    read: usize,
61    dec: D,
62}
63
64impl<D> DecodeExact<D> {
65    /// Create a new `DecodeExact`, delegating to the given `dec` and checking that it decodes
66    /// exactly `target` bytes.
67    pub fn new(dec: D, target: usize) -> DecodeExact<D> {
68        DecodeExact {
69            target,
70            read: 0,
71            dec,
72        }
73    }
74}
75
76impl<D> AsyncDecode for DecodeExact<D>
77    where D: AsyncDecode
78{
79    type Item = D::Item;
80    type Error = DecodeExactError<D::Error, D::Item>;
81
82    fn poll_decode<R: AsyncRead>(mut self,
83                                 cx: &mut Context,
84                                 reader: &mut R)
85                                 -> PollDec<Self::Item, Self, Self::Error> {
86        match self.dec
87                  .poll_decode(cx, &mut LimitedReader::new(reader, self.target - self.read)) {
88            Done(item, read) => {
89                self.read += read;
90                debug_assert!(self.read <= self.target);
91
92                if self.read < self.target {
93                    Errored(DecodeError::DataError(DecodeExactError::Early(item, self.read)))
94                } else {
95                    Done(item, read)
96                }
97            }
98            Progress(inner, read) => {
99                self.read += read;
100                debug_assert!(self.read <= self.target);
101
102                self.dec = inner;
103                Progress(self, read)
104            }
105            Pending(inner) => {
106                self.dec = inner;
107                Pending(self)
108            }
109            Errored(err) => {
110                match err {
111                    DecodeError::DataError(inner_err) => {
112                        Errored(DecodeError::DataError(inner_err.into()))
113                    }
114                    DecodeError::ReaderError(inner_err) => Errored(inner_err.into()),
115                }
116            }
117        }
118    }
119}