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

kafrust_protocol/codec/
decode.rs
