Skip to main content

kafrust_protocol/codec/
decode.rs

1use crate::error::{Error, Result};
2
3#[derive(Debug, Clone, PartialEq, Eq)]
4pub struct TaggedField {
5    pub tag: u32,
6    pub data: Vec<u8>,
7}
8
9#[derive(Debug, Clone)]
10pub struct Decoder<'a> {
11    input: &'a [u8],
12    position: usize,
13}
14
15impl<'a> Decoder<'a> {
16    pub fn new(input: &'a [u8]) -> Self {
17        Self { input, position: 0 }
18    }
19
20    pub fn remaining(&self) -> usize {
21        self.input.len().saturating_sub(self.position)
22    }
23
24    pub fn position(&self) -> usize {
25        self.position
26    }
27
28    pub fn is_empty(&self) -> bool {
29        self.remaining() == 0
30    }
31
32    pub fn read_i8(&mut self) -> Result<i8> {
33        Ok(self.read_exact(1)?[0] as i8)
34    }
35
36    pub fn read_bool(&mut self) -> Result<bool> {
37        match self.read_i8()? {
38            0 => Ok(false),
39            1 => Ok(true),
40            value => Err(Error::InvalidBool(value)),
41        }
42    }
43
44    pub fn read_i16(&mut self) -> Result<i16> {
45        let bytes = self.read_exact(2)?;
46        Ok(i16::from_be_bytes([bytes[0], bytes[1]]))
47    }
48
49    pub fn read_i32(&mut self) -> Result<i32> {
50        let bytes = self.read_exact(4)?;
51        Ok(i32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
52    }
53
54    pub fn read_i64(&mut self) -> Result<i64> {
55        let bytes = self.read_exact(8)?;
56        Ok(i64::from_be_bytes([
57            bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
58        ]))
59    }
60
61    pub fn read_string(&mut self) -> Result<String> {
62        let length = self.read_i16()?;
63        if length < 0 {
64            return Err(Error::NegativeLength {
65                kind: "string",
66                length: i32::from(length),
67            });
68        }
69        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow("string"))?;
70        self.read_utf8(length)
71    }
72
73    pub fn read_nullable_string(&mut self) -> Result<Option<String>> {
74        let length = self.read_i16()?;
75        if length == -1 {
76            return Ok(None);
77        }
78        if length < -1 {
79            return Err(Error::NegativeLength {
80                kind: "nullable string",
81                length: i32::from(length),
82            });
83        }
84        let length =
85            usize::try_from(length).map_err(|_| Error::LengthOverflow("nullable string"))?;
86        Ok(Some(self.read_utf8(length)?))
87    }
88
89    pub fn read_bytes(&mut self) -> Result<Vec<u8>> {
90        let length = self.read_i32()?;
91        if length < 0 {
92            return Err(Error::NegativeLength {
93                kind: "bytes",
94                length,
95            });
96        }
97        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow("bytes"))?;
98        Ok(self.read_exact(length)?.to_vec())
99    }
100
101    pub fn read_nullable_bytes(&mut self) -> Result<Option<Vec<u8>>> {
102        let length = self.read_i32()?;
103        if length == -1 {
104            return Ok(None);
105        }
106        if length < -1 {
107            return Err(Error::NegativeLength {
108                kind: "nullable bytes",
109                length,
110            });
111        }
112        let length =
113            usize::try_from(length).map_err(|_| Error::LengthOverflow("nullable bytes"))?;
114        Ok(Some(self.read_exact(length)?.to_vec()))
115    }
116
117    pub fn read_unsigned_varint(&mut self) -> Result<u32> {
118        let mut value = 0u32;
119        for shift in (0..=28).step_by(7) {
120            let byte = self.read_exact(1)?[0];
121            value |= u32::from(byte & 0x7f) << shift;
122            if byte & 0x80 == 0 {
123                return Ok(value);
124            }
125        }
126        Err(Error::VarintTooLong)
127    }
128
129    pub fn read_varint(&mut self) -> Result<i32> {
130        let value = self.read_unsigned_varint()?;
131        Ok(((value >> 1) as i32) ^ -((value & 1) as i32))
132    }
133
134    pub fn read_varlong(&mut self) -> Result<i64> {
135        let mut value = 0u64;
136        for shift in (0..=63).step_by(7) {
137            let byte = self.read_exact(1)?[0];
138            value |= u64::from(byte & 0x7f) << shift;
139            if byte & 0x80 == 0 {
140                return Ok(((value >> 1) as i64) ^ -((value & 1) as i64));
141            }
142        }
143        Err(Error::VarintTooLong)
144    }
145
146    pub fn read_varint_bytes(&mut self) -> Result<Vec<u8>> {
147        let length = self.read_varint()?;
148        if length < 0 {
149            return Err(Error::NegativeLength {
150                kind: "varint bytes",
151                length,
152            });
153        }
154        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow("varint bytes"))?;
155        Ok(self.read_exact(length)?.to_vec())
156    }
157
158    pub fn read_varint_nullable_bytes(&mut self) -> Result<Option<Vec<u8>>> {
159        let length = self.read_varint()?;
160        if length == -1 {
161            return Ok(None);
162        }
163        if length < -1 {
164            return Err(Error::NegativeLength {
165                kind: "varint nullable bytes",
166                length,
167            });
168        }
169        let length =
170            usize::try_from(length).map_err(|_| Error::LengthOverflow("varint nullable bytes"))?;
171        Ok(Some(self.read_exact(length)?.to_vec()))
172    }
173
174    pub fn read_compact_string(&mut self) -> Result<String> {
175        let encoded_length = self.read_unsigned_varint()?;
176        let length = encoded_length.checked_sub(1).ok_or(Error::NegativeLength {
177            kind: "compact string",
178            length: -1,
179        })?;
180        let length =
181            usize::try_from(length).map_err(|_| Error::LengthOverflow("compact string"))?;
182        self.read_utf8(length)
183    }
184
185    pub fn read_compact_nullable_string(&mut self) -> Result<Option<String>> {
186        let encoded_length = self.read_unsigned_varint()?;
187        if encoded_length == 0 {
188            return Ok(None);
189        }
190        let length = usize::try_from(encoded_length - 1)
191            .map_err(|_| Error::LengthOverflow("compact nullable string"))?;
192        Ok(Some(self.read_utf8(length)?))
193    }
194
195    pub fn read_compact_bytes(&mut self) -> Result<Vec<u8>> {
196        let encoded_length = self.read_unsigned_varint()?;
197        let length = encoded_length.checked_sub(1).ok_or(Error::NegativeLength {
198            kind: "compact bytes",
199            length: -1,
200        })?;
201        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow("compact bytes"))?;
202        Ok(self.read_exact(length)?.to_vec())
203    }
204
205    pub fn read_compact_nullable_bytes(&mut self) -> Result<Option<Vec<u8>>> {
206        let encoded_length = self.read_unsigned_varint()?;
207        if encoded_length == 0 {
208            return Ok(None);
209        }
210        let length = usize::try_from(encoded_length - 1)
211            .map_err(|_| Error::LengthOverflow("compact nullable bytes"))?;
212        Ok(Some(self.read_exact(length)?.to_vec()))
213    }
214
215    pub fn read_array<T>(
216        &mut self,
217        kind: &'static str,
218        mut read_item: impl FnMut(&mut Self) -> Result<T>,
219    ) -> Result<Option<Vec<T>>> {
220        let length = self.read_i32()?;
221        if length == -1 {
222            return Ok(None);
223        }
224        if length < -1 {
225            return Err(Error::NegativeLength { kind, length });
226        }
227        let length = usize::try_from(length).map_err(|_| Error::LengthOverflow(kind))?;
228        let mut values = Vec::with_capacity(length);
229        for _ in 0..length {
230            values.push(read_item(self)?);
231        }
232        Ok(Some(values))
233    }
234
235    pub fn read_compact_array<T>(
236        &mut self,
237        kind: &'static str,
238        mut read_item: impl FnMut(&mut Self) -> Result<T>,
239    ) -> Result<Option<Vec<T>>> {
240        let encoded_length = self.read_unsigned_varint()?;
241        if encoded_length == 0 {
242            return Ok(None);
243        }
244        let length =
245            usize::try_from(encoded_length - 1).map_err(|_| Error::LengthOverflow(kind))?;
246        let mut values = Vec::with_capacity(length);
247        for _ in 0..length {
248            values.push(read_item(self)?);
249        }
250        Ok(Some(values))
251    }
252
253    pub fn read_tagged_fields(&mut self) -> Result<Vec<TaggedField>> {
254        let count = self.read_unsigned_varint()?;
255        let count = usize::try_from(count).map_err(|_| Error::LengthOverflow("tagged fields"))?;
256        let mut fields = Vec::with_capacity(count);
257        for _ in 0..count {
258            let tag = self.read_unsigned_varint()?;
259            let length = self.read_unsigned_varint()?;
260            let length =
261                usize::try_from(length).map_err(|_| Error::LengthOverflow("tagged field data"))?;
262            let data = self.read_exact(length)?.to_vec();
263            fields.push(TaggedField { tag, data });
264        }
265        Ok(fields)
266    }
267
268    pub fn read_exact(&mut self, length: usize) -> Result<&'a [u8]> {
269        if self.remaining() < length {
270            return Err(Error::UnexpectedEof {
271                needed: length,
272                remaining: self.remaining(),
273            });
274        }
275        let start = self.position;
276        self.position += length;
277        Ok(&self.input[start..self.position])
278    }
279
280    fn read_utf8(&mut self, length: usize) -> Result<String> {
281        let bytes = self.read_exact(length)?;
282        let value = core::str::from_utf8(bytes).map_err(|_| Error::InvalidUtf8)?;
283        Ok(value.to_owned())
284    }
285}