codecs/decoding/framing/
varint_length_delimited.rs

1use bytes::{Buf, Bytes, BytesMut};
2use snafu::Snafu;
3use tokio_util::codec::Decoder;
4use vector_config::configurable_component;
5
6use super::{BoxedFramingError, FramingError, StreamDecodingError};
7
8/// Errors that can occur during varint length delimited framing.
9#[derive(Debug, Snafu)]
10pub enum VarintFramingError {
11    #[snafu(display("Varint too large"))]
12    VarintOverflow,
13
14    #[snafu(display("Frame too large: {length} bytes (max: {max})"))]
15    FrameTooLarge { length: usize, max: usize },
16
17    #[snafu(display("Trailing data at EOF"))]
18    TrailingData,
19}
20
21impl StreamDecodingError for VarintFramingError {
22    fn can_continue(&self) -> bool {
23        match self {
24            // Varint overflow and frame too large are not recoverable
25            Self::VarintOverflow | Self::FrameTooLarge { .. } => false,
26            // Trailing data at EOF is not recoverable
27            Self::TrailingData => false,
28        }
29    }
30}
31
32impl FramingError for VarintFramingError {
33    fn as_any(&self) -> &dyn std::any::Any {
34        self as &dyn std::any::Any
35    }
36}
37
38/// Config used to build a `VarintLengthDelimitedDecoder`.
39#[configurable_component]
40#[derive(Debug, Clone, Default)]
41pub struct VarintLengthDelimitedDecoderConfig {
42    /// Maximum frame length
43    #[serde(default = "default_max_frame_length")]
44    pub max_frame_length: usize,
45}
46
47const fn default_max_frame_length() -> usize {
48    8 * 1_024 * 1_024
49}
50
51impl VarintLengthDelimitedDecoderConfig {
52    /// Build the `VarintLengthDelimitedDecoder` from this configuration.
53    pub fn build(&self) -> VarintLengthDelimitedDecoder {
54        VarintLengthDelimitedDecoder::new(self.max_frame_length)
55    }
56}
57
58/// A codec for handling bytes sequences whose length is encoded as a varint prefix.
59/// This is compatible with protobuf's length-delimited encoding.
60#[derive(Debug, Clone)]
61pub struct VarintLengthDelimitedDecoder {
62    max_frame_length: usize,
63}
64
65impl VarintLengthDelimitedDecoder {
66    /// Creates a new `VarintLengthDelimitedDecoder`.
67    pub fn new(max_frame_length: usize) -> Self {
68        Self { max_frame_length }
69    }
70
71    /// Decode a varint from the start of the buffer without consuming it.
72    fn decode_varint(&self, buf: &BytesMut) -> Result<Option<(u64, usize)>, BoxedFramingError> {
73        let mut value: u64 = 0;
74        let mut shift: u8 = 0;
75
76        for (index, byte) in buf.iter().enumerate() {
77            let byte_value = (*byte & 0x7F) as u64;
78            value |= byte_value << shift;
79
80            if *byte & 0x80 == 0 {
81                // Last byte of varint
82                return Ok(Some((value, index + 1)));
83            }
84
85            shift += 7;
86            if shift >= 64 {
87                return Err(VarintFramingError::VarintOverflow.into());
88            }
89        }
90
91        // Incomplete varint
92        Ok(None)
93    }
94}
95
96impl Default for VarintLengthDelimitedDecoder {
97    fn default() -> Self {
98        Self::new(default_max_frame_length())
99    }
100}
101
102impl Decoder for VarintLengthDelimitedDecoder {
103    type Item = Bytes;
104    type Error = BoxedFramingError;
105
106    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
107        // First, peek at the varint length prefix.
108        let (length, prefix_length) = match self.decode_varint(src)? {
109            Some((length, prefix_length)) => {
110                (usize::try_from(length).unwrap_or(usize::MAX), prefix_length)
111            }
112            None => return Ok(None), // Incomplete varint
113        };
114
115        // Check if the length is reasonable
116        if length > self.max_frame_length {
117            return Err(VarintFramingError::FrameTooLarge {
118                length,
119                max: self.max_frame_length,
120            }
121            .into());
122        }
123
124        // Check if we have enough data for the complete frame
125        if src.len() - prefix_length < length {
126            return Ok(None); // Incomplete frame
127        }
128
129        // Consume the length prefix and extract the frame
130        src.advance(prefix_length);
131        let frame = src.split_to(length).freeze();
132        Ok(Some(frame))
133    }
134
135    fn decode_eof(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
136        if src.is_empty() {
137            Ok(None)
138        } else {
139            // Try to decode what we have, even if incomplete
140            match self.decode(src)? {
141                Some(frame) => Ok(Some(frame)),
142                None => {
143                    // If we have data but couldn't decode it, it's trailing data
144                    if !src.is_empty() {
145                        Err(VarintFramingError::TrailingData.into())
146                    } else {
147                        Ok(None)
148                    }
149                }
150            }
151        }
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    #[test]
160    fn decode_single_byte_varint() {
161        let mut input = BytesMut::from(&[0x03, b'f', b'o', b'o'][..]);
162        let mut decoder = VarintLengthDelimitedDecoder::default();
163
164        assert_eq!(
165            decoder.decode(&mut input).unwrap().unwrap(),
166            Bytes::from("foo")
167        );
168        assert_eq!(decoder.decode(&mut input).unwrap(), None);
169    }
170
171    #[test]
172    fn decode_multi_byte_varint() {
173        // 300 in varint encoding: 0xAC 0x02
174        let mut input = BytesMut::from(&[0xAC, 0x02][..]);
175        // Add 300 bytes of data
176        input.extend_from_slice(&vec![b'x'; 300]);
177        let mut decoder = VarintLengthDelimitedDecoder::default();
178
179        let result = decoder.decode(&mut input).unwrap().unwrap();
180        assert_eq!(result.len(), 300);
181        assert_eq!(decoder.decode(&mut input).unwrap(), None);
182    }
183
184    #[test]
185    fn decode_incomplete_varint() {
186        let mut input = BytesMut::from(&[0x80][..]); // Incomplete varint
187        let mut decoder = VarintLengthDelimitedDecoder::default();
188
189        assert_eq!(decoder.decode(&mut input).unwrap(), None);
190    }
191
192    #[test]
193    fn decode_incomplete_frame() {
194        let mut input = BytesMut::from(&[0x05, b'f', b'o'][..]); // Length 5, but only 2 bytes
195        let mut decoder = VarintLengthDelimitedDecoder::default();
196
197        assert_eq!(decoder.decode(&mut input).unwrap(), None);
198    }
199
200    #[test]
201    fn decode_frames_split_across_read_buffer() {
202        const FRAME_LENGTH: usize = 87;
203        const FRAME_COUNT: usize = 200;
204        const READ_BUFFER_SIZE: usize = 8 * 1024;
205
206        let mut encoded = Vec::new();
207
208        for index in 0..FRAME_COUNT {
209            let byte = b'A' + (index % 26) as u8;
210
211            // FRAME_LENGTH fits in a single-byte varint.
212            encoded.push(FRAME_LENGTH as u8);
213            encoded.extend_from_slice(&[byte; FRAME_LENGTH]);
214        }
215
216        let mut decoder = VarintLengthDelimitedDecoder::default();
217        let mut input = BytesMut::new();
218        let mut decoded = Vec::new();
219        let mut chunks = encoded.chunks(READ_BUFFER_SIZE);
220
221        input.extend_from_slice(chunks.next().unwrap());
222
223        while let Some(frame) = decoder.decode(&mut input).unwrap() {
224            decoded.push(frame);
225        }
226
227        // 93 complete 88-byte frames fit in the first read, followed by the
228        // length prefix and seven payload bytes of the next frame.
229        assert_eq!(decoded.len(), 93);
230        assert_eq!(input.len(), 8);
231        assert_eq!(input[0], FRAME_LENGTH as u8);
232
233        for chunk in chunks {
234            input.extend_from_slice(chunk);
235
236            while let Some(frame) = decoder.decode(&mut input).unwrap() {
237                decoded.push(frame);
238            }
239        }
240
241        assert_eq!(decoder.decode_eof(&mut input).unwrap(), None);
242        assert!(input.is_empty());
243        assert_eq!(decoded.len(), FRAME_COUNT);
244
245        for (index, frame) in decoded.iter().enumerate() {
246            let expected = b'A' + (index % 26) as u8;
247
248            assert_eq!(frame.len(), FRAME_LENGTH);
249            assert!(frame.iter().all(|byte| *byte == expected));
250        }
251    }
252
253    #[test]
254    fn decode_incomplete_multi_byte_varint_leaves_buffer_intact() {
255        let mut input = BytesMut::from(&[0xAC, 0x02][..]);
256        input.extend_from_slice(&[b'x'; 100]);
257        let mut decoder = VarintLengthDelimitedDecoder::default();
258
259        assert_eq!(decoder.decode(&mut input).unwrap(), None);
260        assert_eq!(input.len(), 102);
261        assert_eq!(&input[..2], &[0xAC, 0x02][..]);
262
263        input.extend_from_slice(&[b'x'; 200]);
264        assert_eq!(decoder.decode(&mut input).unwrap().unwrap().len(), 300);
265        assert!(input.is_empty());
266    }
267
268    #[test]
269    fn decode_frame_too_large() {
270        let mut input =
271            BytesMut::from(&[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x01][..]);
272        let mut decoder = VarintLengthDelimitedDecoder::new(1000);
273
274        assert!(decoder.decode(&mut input).is_err());
275    }
276
277    #[test]
278    fn decode_trailing_data_at_eof() {
279        let mut input = BytesMut::from(&[0x03, b'f', b'o', b'o', b'e', b'x', b't', b'r', b'a'][..]);
280        let mut decoder = VarintLengthDelimitedDecoder::default();
281
282        // First decode should succeed
283        assert_eq!(
284            decoder.decode(&mut input).unwrap().unwrap(),
285            Bytes::from("foo")
286        );
287
288        // Second decode should fail with trailing data
289        assert!(decoder.decode_eof(&mut input).is_err());
290    }
291}
codecs/decoding/framing/varint_length_delimited.rs

codecs/decoding/framing/
varint_length_delimited.rs
