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    /// Reads a Kafka UUID from its 16-byte network-order representation.
139    pub fn read_uuid(&mut self) -> Result<[u8; 16]> {
140        let bytes = self.read_exact(16)?;
141        let mut value = [0; 16];
142        value.copy_from_slice(bytes);
143        Ok(value)
144    }
145
146    pub fn read_f64(&mut self) -> Result<f64> {
147        let bytes = self.read_exact(8)?;
148        Ok(f64::from_bits(u64::from_be_bytes([
149            bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
150        ])))
151    }
152
153    pub fn read_string(&mut self) -> Result<String> {
154        let length = self.read_i16()?;
155        if length < 0 {
156            return Err(Error::NegativeLength {
157                kind: "string",
158                length: i32::from(length),
159            });
160        }
161        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow("string"))?;
162        self.read_utf8(length)
163    }
164
165    pub fn read_nullable_string(&mut self) -> Result<Option<String>> {
166        let length = self.read_i16()?;
167        if length == -1 {
168            return Ok(None);
169        }
170        if length < -1 {
171            return Err(Error::NegativeLength {
172                kind: "nullable string",
173                length: i32::from(length),
174            });
175        }
176        let length =
177            usize::try_from(length).map_err(|_| Error::LengthOverflow("nullable string"))?;
178        Ok(Some(self.read_utf8(length)?))
179    }
180
181    pub fn read_bytes(&mut self) -> Result<Vec<u8>> {
182        let length = self.read_i32()?;
183        if length < 0 {
184            return Err(Error::NegativeLength {
185                kind: "bytes",
186                length,
187            });
188        }
189        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow("bytes"))?;
190        Ok(self.read_exact(length)?.to_vec())
191    }
192
193    pub fn read_nullable_bytes(&mut self) -> Result<Option<Vec<u8>>> {
194        let length = self.read_i32()?;
195        if length == -1 {
196            return Ok(None);
197        }
198        if length < -1 {
199            return Err(Error::NegativeLength {
200                kind: "nullable bytes",
201                length,
202            });
203        }
204        let length =
205            usize::try_from(length).map_err(|_| Error::LengthOverflow("nullable bytes"))?;
206        Ok(Some(self.read_exact(length)?.to_vec()))
207    }
208
209    pub fn read_unsigned_varint(&mut self) -> Result<u32> {
210        let mut value = 0u32;
211        for shift in (0..=28).step_by(7) {
212            let byte = self.read_exact(1)?[0];
213            value |= u32::from(byte & 0x7f) << shift;
214            if byte & 0x80 == 0 {
215                return Ok(value);
216            }
217        }
218        Err(Error::VarintTooLong)
219    }
220
221    pub fn read_varint(&mut self) -> Result<i32> {
222        let value = self.read_unsigned_varint()?;
223        Ok(((value >> 1) as i32) ^ -((value & 1) as i32))
224    }
225
226    pub fn read_varlong(&mut self) -> Result<i64> {
227        let mut value = 0u64;
228        for shift in (0..=63).step_by(7) {
229            let byte = self.read_exact(1)?[0];
230            value |= u64::from(byte & 0x7f) << shift;
231            if byte & 0x80 == 0 {
232                return Ok(((value >> 1) as i64) ^ -((value & 1) as i64));
233            }
234        }
235        Err(Error::VarintTooLong)
236    }
237
238    pub fn read_varint_bytes(&mut self) -> Result<Vec<u8>> {
239        let length = self.read_varint()?;
240        if length < 0 {
241            return Err(Error::NegativeLength {
242                kind: "varint bytes",
243                length,
244            });
245        }
246        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow("varint bytes"))?;
247        Ok(self.read_exact(length)?.to_vec())
248    }
249
250    pub fn read_varint_nullable_bytes(&mut self) -> Result<Option<Vec<u8>>> {
251        let length = self.read_varint()?;
252        if length == -1 {
253            return Ok(None);
254        }
255        if length < -1 {
256            return Err(Error::NegativeLength {
257                kind: "varint nullable bytes",
258                length,
259            });
260        }
261        let length =
262            usize::try_from(length).map_err(|_| Error::LengthOverflow("varint nullable bytes"))?;
263        Ok(Some(self.read_exact(length)?.to_vec()))
264    }
265
266    pub fn read_compact_string(&mut self) -> Result<String> {
267        let encoded_length = self.read_unsigned_varint()?;
268        let length = encoded_length.checked_sub(1).ok_or(Error::NegativeLength {
269            kind: "compact string",
270            length: -1,
271        })?;
272        let length =
273            usize::try_from(length).map_err(|_| Error::LengthOverflow("compact string"))?;
274        self.read_utf8(length)
275    }
276
277    pub fn read_compact_nullable_string(&mut self) -> Result<Option<String>> {
278        let encoded_length = self.read_unsigned_varint()?;
279        if encoded_length == 0 {
280            return Ok(None);
281        }
282        let length = usize::try_from(encoded_length - 1)
283            .map_err(|_| Error::LengthOverflow("compact nullable string"))?;
284        Ok(Some(self.read_utf8(length)?))
285    }
286
287    pub fn read_compact_bytes(&mut self) -> Result<Vec<u8>> {
288        let encoded_length = self.read_unsigned_varint()?;
289        let length = encoded_length.checked_sub(1).ok_or(Error::NegativeLength {
290            kind: "compact bytes",
291            length: -1,
292        })?;
293        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow("compact bytes"))?;
294        Ok(self.read_exact(length)?.to_vec())
295    }
296
297    pub fn read_compact_nullable_bytes(&mut self) -> Result<Option<Vec<u8>>> {
298        let encoded_length = self.read_unsigned_varint()?;
299        if encoded_length == 0 {
300            return Ok(None);
301        }
302        let length = usize::try_from(encoded_length - 1)
303            .map_err(|_| Error::LengthOverflow("compact nullable bytes"))?;
304        Ok(Some(self.read_exact(length)?.to_vec()))
305    }
306
307    pub fn read_array<T>(
308        &mut self,
309        kind: &'static str,
310        mut read_item: impl FnMut(&mut Self) -> Result<T>,
311    ) -> Result<Option<Vec<T>>> {
312        let length = self.read_i32()?;
313        if length == -1 {
314            return Ok(None);
315        }
316        if length < -1 {
317            return Err(Error::NegativeLength { kind, length });
318        }
319        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow(kind))?;
320        self.ensure_collection_length(kind, length)?;
321        let mut values = Vec::with_capacity(length);
322        for _ in 0..length {
323            values.push(read_item(self)?);
324        }
325        Ok(Some(values))
326    }
327
328    pub fn read_compact_array<T>(
329        &mut self,
330        kind: &'static str,
331        mut read_item: impl FnMut(&mut Self) -> Result<T>,
332    ) -> Result<Option<Vec<T>>> {
333        let encoded_length = self.read_unsigned_varint()?;
334        if encoded_length == 0 {
335            return Ok(None);
336        }
337        let length =
338            usize::try_from(encoded_length - 1).map_err(|_| Error::LengthOverflow(kind))?;
339        self.ensure_collection_length(kind, length)?;
340        let mut values = Vec::with_capacity(length);
341        for _ in 0..length {
342            values.push(read_item(self)?);
343        }
344        Ok(Some(values))
345    }
346
347    pub fn read_tagged_fields(&mut self) -> Result<Vec<TaggedField>> {
348        let count = self.read_unsigned_varint()?;
349        let count = usize::try_from(count).map_err(|_| Error::LengthOverflow("tagged fields"))?;
350        self.ensure_collection_length("tagged fields", count)?;
351        let mut fields = Vec::with_capacity(count);
352        for _ in 0..count {
353            let tag = self.read_unsigned_varint()?;
354            let length = self.read_unsigned_varint()?;
355            let length =
356                usize::try_from(length).map_err(|_| Error::LengthOverflow("tagged field data"))?;
357            let data = self.read_exact(length)?.to_vec();
358            fields.push(TaggedField { tag, data });
359        }
360        Ok(fields)
361    }
362
363    pub fn read_exact(&mut self, length: usize) -> Result<&'a [u8]> {
364        if self.remaining() < length {
365            return Err(Error::UnexpectedEof {
366                needed: length,
367                remaining: self.remaining(),
368            });
369        }
370        let start = self.position;
371        self.position += length;
372        Ok(&self.input[start..self.position])
373    }
374
375    fn read_utf8(&mut self, length: usize) -> Result<String> {
376        let bytes = self.read_exact(length)?;
377        let value = core::str::from_utf8(bytes).map_err(|_| Error::InvalidUtf8)?;
378        Ok(value.to_owned())
379    }
380}
kafrust_protocol/codec/decode.rs

kafrust_protocol/codec/
decode.rs
