kafrust_protocol/codec/
decode.rs

1use crate::error::{Error, Result};
2
3/// Resource limits applied while decoding broker responses.
4#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5pub struct DecodeLimits {
6    max_array_elements: usize,
7    max_decompressed_record_bytes: usize,
8}
9
10impl DecodeLimits {
11    /// Default maximum number of elements in one decoded Kafka array.
12    pub const DEFAULT_MAX_ARRAY_ELEMENTS: usize = 1_000_000;
13    /// Default maximum uncompressed size of one Kafka record batch.
14    pub const DEFAULT_MAX_DECOMPRESSED_RECORD_BYTES: usize = 64 * 1024 * 1024;
15
16    /// Creates the default decoding limits.
17    pub const fn new() -> Self {
18        Self {
19            max_array_elements: Self::DEFAULT_MAX_ARRAY_ELEMENTS,
20            max_decompressed_record_bytes: Self::DEFAULT_MAX_DECOMPRESSED_RECORD_BYTES,
21        }
22    }
23
24    /// Sets the maximum number of elements in one decoded Kafka array.
25    pub const fn with_max_array_elements(mut self, max: usize) -> Self {
26        self.max_array_elements = max;
27        self
28    }
29
30    /// Sets the maximum uncompressed size of one Kafka record batch.
31    pub const fn with_max_decompressed_record_bytes(mut self, max: usize) -> Self {
32        self.max_decompressed_record_bytes = max;
33        self
34    }
35
36    /// Returns the maximum number of elements in one decoded Kafka array.
37    pub const fn max_array_elements(self) -> usize {
38        self.max_array_elements
39    }
40
41    /// Returns the maximum uncompressed size of one Kafka record batch.
42    pub const fn max_decompressed_record_bytes(self) -> usize {
43        self.max_decompressed_record_bytes
44    }
45}
46
47impl Default for DecodeLimits {
48    fn default() -> Self {
49        Self::new()
50    }
51}
52
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct TaggedField {
55    pub tag: u32,
56    pub data: Vec<u8>,
57}
58
59#[derive(Debug, Clone)]
60pub struct Decoder<'a> {
61    input: &'a [u8],
62    position: usize,
63    limits: DecodeLimits,
64}
65
66impl<'a> Decoder<'a> {
67    pub fn new(input: &'a [u8]) -> Self {
68        Self::with_limits(input, DecodeLimits::default())
69    }
70
71    /// Creates a decoder with explicit resource limits.
72    pub fn with_limits(input: &'a [u8], limits: DecodeLimits) -> Self {
73        Self {
74            input,
75            position: 0,
76            limits,
77        }
78    }
79
80    /// Returns the resource limits inherited by nested decoders.
81    pub const fn limits(&self) -> DecodeLimits {
82        self.limits
83    }
84
85    /// Rejects a collection length before allocating storage for it.
86    pub fn ensure_collection_length(&self, kind: &'static str, length: usize) -> Result<()> {
87        if length > self.limits.max_array_elements {
88            return Err(Error::LimitExceeded {
89                kind,
90                actual: length,
91                max: self.limits.max_array_elements,
92            });
93        }
94        Ok(())
95    }
96
97    pub fn remaining(&self) -> usize {
98        self.input.len().saturating_sub(self.position)
99    }
100
101    pub fn position(&self) -> usize {
102        self.position
103    }
104
105    pub fn is_empty(&self) -> bool {
106        self.remaining() == 0
107    }
108
109    pub fn read_i8(&mut self) -> Result<i8> {
110        Ok(self.read_exact(1)?[0] as i8)
111    }
112
113    pub fn read_bool(&mut self) -> Result<bool> {
114        match self.read_i8()? {
115            0 => Ok(false),
116            1 => Ok(true),
117            value => Err(Error::InvalidBool(value)),
118        }
119    }
120
121    pub fn read_i16(&mut self) -> Result<i16> {
122        let bytes = self.read_exact(2)?;
123        Ok(i16::from_be_bytes([bytes[0], bytes[1]]))
124    }
125
126    pub fn read_i32(&mut self) -> Result<i32> {
127        let bytes = self.read_exact(4)?;
128        Ok(i32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
129    }
130
131    pub fn read_i64(&mut self) -> Result<i64> {
132        let bytes = self.read_exact(8)?;
133        Ok(i64::from_be_bytes([
134            bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
135        ]))
136    }
137
138    pub fn read_f64(&mut self) -> Result<f64> {
139        let bytes = self.read_exact(8)?;
140        Ok(f64::from_bits(u64::from_be_bytes([
141            bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
142        ])))
143    }
144
145    pub fn read_string(&mut self) -> Result<String> {
146        let length = self.read_i16()?;
147        if length < 0 {
148            return Err(Error::NegativeLength {
149                kind: "string",
150                length: i32::from(length),
151            });
152        }
153        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow("string"))?;
154        self.read_utf8(length)
155    }
156
157    pub fn read_nullable_string(&mut self) -> Result<Option<String>> {
158        let length = self.read_i16()?;
159        if length == -1 {
160            return Ok(None);
161        }
162        if length < -1 {
163            return Err(Error::NegativeLength {
164                kind: "nullable string",
165                length: i32::from(length),
166            });
167        }
168        let length =
169            usize::try_from(length).map_err(|_| Error::LengthOverflow("nullable string"))?;
170        Ok(Some(self.read_utf8(length)?))
171    }
172
173    pub fn read_bytes(&mut self) -> Result<Vec<u8>> {
174        let length = self.read_i32()?;
175        if length < 0 {
176            return Err(Error::NegativeLength {
177                kind: "bytes",
178                length,
179            });
180        }
181        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow("bytes"))?;
182        Ok(self.read_exact(length)?.to_vec())
183    }
184
185    pub fn read_nullable_bytes(&mut self) -> Result<Option<Vec<u8>>> {
186        let length = self.read_i32()?;
187        if length == -1 {
188            return Ok(None);
189        }
190        if length < -1 {
191            return Err(Error::NegativeLength {
192                kind: "nullable bytes",
193                length,
194            });
195        }
196        let length =
197            usize::try_from(length).map_err(|_| Error::LengthOverflow("nullable bytes"))?;
198        Ok(Some(self.read_exact(length)?.to_vec()))
199    }
200
201    pub fn read_unsigned_varint(&mut self) -> Result<u32> {
202        let mut value = 0u32;
203        for shift in (0..=28).step_by(7) {
204            let byte = self.read_exact(1)?[0];
205            value |= u32::from(byte & 0x7f) << shift;
206            if byte & 0x80 == 0 {
207                return Ok(value);
208            }
209        }
210        Err(Error::VarintTooLong)
211    }
212
213    pub fn read_varint(&mut self) -> Result<i32> {
214        let value = self.read_unsigned_varint()?;
215        Ok(((value >> 1) as i32) ^ -((value & 1) as i32))
216    }
217
218    pub fn read_varlong(&mut self) -> Result<i64> {
219        let mut value = 0u64;
220        for shift in (0..=63).step_by(7) {
221            let byte = self.read_exact(1)?[0];
222            value |= u64::from(byte & 0x7f) << shift;
223            if byte & 0x80 == 0 {
224                return Ok(((value >> 1) as i64) ^ -((value & 1) as i64));
225            }
226        }
227        Err(Error::VarintTooLong)
228    }
229
230    pub fn read_varint_bytes(&mut self) -> Result<Vec<u8>> {
231        let length = self.read_varint()?;
232        if length < 0 {
233            return Err(Error::NegativeLength {
234                kind: "varint bytes",
235                length,
236            });
237        }
238        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow("varint bytes"))?;
239        Ok(self.read_exact(length)?.to_vec())
240    }
241
242    pub fn read_varint_nullable_bytes(&mut self) -> Result<Option<Vec<u8>>> {
243        let length = self.read_varint()?;
244        if length == -1 {
245            return Ok(None);
246        }
247        if length < -1 {
248            return Err(Error::NegativeLength {
249                kind: "varint nullable bytes",
250                length,
251            });
252        }
253        let length =
254            usize::try_from(length).map_err(|_| Error::LengthOverflow("varint nullable bytes"))?;
255        Ok(Some(self.read_exact(length)?.to_vec()))
256    }
257
258    pub fn read_compact_string(&mut self) -> Result<String> {
259        let encoded_length = self.read_unsigned_varint()?;
260        let length = encoded_length.checked_sub(1).ok_or(Error::NegativeLength {
261            kind: "compact string",
262            length: -1,
263        })?;
264        let length =
265            usize::try_from(length).map_err(|_| Error::LengthOverflow("compact string"))?;
266        self.read_utf8(length)
267    }
268
269    pub fn read_compact_nullable_string(&mut self) -> Result<Option<String>> {
270        let encoded_length = self.read_unsigned_varint()?;
271        if encoded_length == 0 {
272            return Ok(None);
273        }
274        let length = usize::try_from(encoded_length - 1)
275            .map_err(|_| Error::LengthOverflow("compact nullable string"))?;
276        Ok(Some(self.read_utf8(length)?))
277    }
278
279    pub fn read_compact_bytes(&mut self) -> Result<Vec<u8>> {
280        let encoded_length = self.read_unsigned_varint()?;
281        let length = encoded_length.checked_sub(1).ok_or(Error::NegativeLength {
282            kind: "compact bytes",
283            length: -1,
284        })?;
285        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow("compact bytes"))?;
286        Ok(self.read_exact(length)?.to_vec())
287    }
288
289    pub fn read_compact_nullable_bytes(&mut self) -> Result<Option<Vec<u8>>> {
290        let encoded_length = self.read_unsigned_varint()?;
291        if encoded_length == 0 {
292            return Ok(None);
293        }
294        let length = usize::try_from(encoded_length - 1)
295            .map_err(|_| Error::LengthOverflow("compact nullable bytes"))?;
296        Ok(Some(self.read_exact(length)?.to_vec()))
297    }
298
299    pub fn read_array<T>(
300        &mut self,
301        kind: &'static str,
302        mut read_item: impl FnMut(&mut Self) -> Result<T>,
303    ) -> Result<Option<Vec<T>>> {
304        let length = self.read_i32()?;
305        if length == -1 {
306            return Ok(None);
307        }
308        if length < -1 {
309            return Err(Error::NegativeLength { kind, length });
310        }
311        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow(kind))?;
312        self.ensure_collection_length(kind, length)?;
313        let mut values = Vec::with_capacity(length);
314        for _ in 0..length {
315            values.push(read_item(self)?);
316        }
317        Ok(Some(values))
318    }
319
320    pub fn read_compact_array<T>(
321        &mut self,
322        kind: &'static str,
323        mut read_item: impl FnMut(&mut Self) -> Result<T>,
324    ) -> Result<Option<Vec<T>>> {
325        let encoded_length = self.read_unsigned_varint()?;
326        if encoded_length == 0 {
327            return Ok(None);
328        }
329        let length =
330            usize::try_from(encoded_length - 1).map_err(|_| Error::LengthOverflow(kind))?;
331        self.ensure_collection_length(kind, length)?;
332        let mut values = Vec::with_capacity(length);
333        for _ in 0..length {
334            values.push(read_item(self)?);
335        }
336        Ok(Some(values))
337    }
338
339    pub fn read_tagged_fields(&mut self) -> Result<Vec<TaggedField>> {
340        let count = self.read_unsigned_varint()?;
341        let count = usize::try_from(count).map_err(|_| Error::LengthOverflow("tagged fields"))?;
342        self.ensure_collection_length("tagged fields", count)?;
343        let mut fields = Vec::with_capacity(count);
344        for _ in 0..count {
345            let tag = self.read_unsigned_varint()?;
346            let length = self.read_unsigned_varint()?;
347            let length =
348                usize::try_from(length).map_err(|_| Error::LengthOverflow("tagged field data"))?;
349            let data = self.read_exact(length)?.to_vec();
350            fields.push(TaggedField { tag, data });
351        }
352        Ok(fields)
353    }
354
355    pub fn read_exact(&mut self, length: usize) -> Result<&'a [u8]> {
356        if self.remaining() < length {
357            return Err(Error::UnexpectedEof {
358                needed: length,
359                remaining: self.remaining(),
360            });
361        }
362        let start = self.position;
363        self.position += length;
364        Ok(&self.input[start..self.position])
365    }
366
367    fn read_utf8(&mut self, length: usize) -> Result<String> {
368        let bytes = self.read_exact(length)?;
369        let value = core::str::from_utf8(bytes).map_err(|_| Error::InvalidUtf8)?;
370        Ok(value.to_owned())
371    }
372}
kafrust_protocol/codec/decode.rs

kafrust_protocol/codec/
decode.rs
