1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

use super::{Decoder, Encoder};
use bytes::{Buf, BytesMut};

#[allow(missing_docs)]
pub trait SkipAheadHandler: Sized + std::fmt::Debug {
    fn continue_skipping(self, src: &[u8]) -> anyhow::Result<(usize, Option<Self>)>;
}

impl SkipAheadHandler for () {
    fn continue_skipping(self, _: &[u8]) -> anyhow::Result<(usize, Option<Self>)> {
        Ok((0, None))
    }
}

#[allow(missing_docs)]
pub trait DecoderWithSkipAhead: Decoder {
    type Handler: SkipAheadHandler;

    fn prepare_skip_ahead(&mut self, src: &mut BytesMut) -> Self::Handler;
}

#[derive(Debug)]
pub struct LimitCodec<C: DecoderWithSkipAhead> {
    inner: C,
    max_frame_size: usize,
    skip_ahead_state: Option<<C as DecoderWithSkipAhead>::Handler>,
    decoder_defunct: bool,
}

impl<C> LimitCodec<C>
where
    C: DecoderWithSkipAhead,
{
    #[allow(missing_docs)]
    pub fn new(inner: C, max_frame_size: usize) -> Self {
        Self {
            inner,
            max_frame_size,
            skip_ahead_state: None,
            decoder_defunct: false,
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub enum LimitError<E: std::error::Error + 'static> {
    #[error("frame size limit exceeded (detected at {0} bytes)")]
    LimitExceeded(usize),
    #[error("codec couldn't recover from invalid or too big frame")]
    Defunct,
    #[error(transparent)]
    Inner(#[from] E),
}

impl<C> Encoder for LimitCodec<C>
where
    C: Encoder + DecoderWithSkipAhead,
{
    type Error = LimitError<<C as Encoder>::Error>;
    type Item = <C as Encoder>::Item;

    fn encode(&mut self, src: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
        let mut tmp_dst = dst.split_off(dst.len());
        self.inner.encode(src, &mut tmp_dst)?;

        if tmp_dst.len() > self.max_frame_size {
            return Err(LimitError::LimitExceeded(tmp_dst.len()));
        }

        dst.unsplit(tmp_dst);
        Ok(())
    }
}

impl<C> Decoder for LimitCodec<C>
where
    C: DecoderWithSkipAhead,
{
    type Error = LimitError<<C as Decoder>::Error>;
    type Item = <C as Decoder>::Item;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        while let Some(sas) = self.skip_ahead_state.take() {
            match sas.continue_skipping(src) {
                Ok((amount, next)) => {
                    self.skip_ahead_state = next;
                    debug_assert!(amount <= src.len());
                    src.advance(amount);
                    debug_assert!(amount != 0 || self.skip_ahead_state.is_none());
                    if src.is_empty() {
                        return Ok(None);
                    }
                },
                Err(_err) => {
                    self.decoder_defunct = true;
                },
            }
        }

        if self.decoder_defunct {
            src.clear();
            return Err(LimitError::Defunct);
        }
        match self.inner.decode(src) {
            Ok(None) if src.len() > self.max_frame_size => {
                self.skip_ahead_state = Some(self.inner.prepare_skip_ahead(src));
                Err(LimitError::LimitExceeded(src.len()))
            },
            Ok(x) => Ok(x),
            Err(x) => Err(LimitError::Inner(x)),
        }
    }
}