1pub struct Reader<'a> {
11 pub data: &'a [u8],
12 pub pos: usize,
13}
14
15impl<'a> Reader<'a> {
16 pub fn new(data: &'a [u8]) -> Self {
17 Self { data, pos: 0 }
18 }
19
20 pub fn remaining(&self) -> &'a [u8] {
21 &self.data[self.pos..]
22 }
23
24 fn need(&self, n: usize) -> Result<(), DecodeErr> {
25 if self.pos + n > self.data.len() {
26 Err(DecodeErr::Eof)
27 } else {
28 Ok(())
29 }
30 }
31
32 pub fn read_u8(&mut self) -> Result<u8, DecodeErr> {
33 self.need(1)?;
34 let v = self.data[self.pos];
35 self.pos += 1;
36 Ok(v)
37 }
38
39 pub fn read_u32_le(&mut self) -> Result<u32, DecodeErr> {
40 self.need(4)?;
41 let v = u32::from_le_bytes(self.data[self.pos..self.pos + 4].try_into().unwrap());
42 self.pos += 4;
43 Ok(v)
44 }
45
46 pub fn read_raw(&mut self, n: usize) -> Result<&'a [u8], DecodeErr> {
47 self.need(n)?;
48 let v = &self.data[self.pos..self.pos + n];
49 self.pos += n;
50 Ok(v)
51 }
52
53 pub fn read_compact_u32(&mut self) -> Result<u32, DecodeErr> {
55 let first = self.read_u8()? as u32;
56 match first & 0b11 {
57 0b00 => Ok(first >> 2),
58 0b01 => {
59 let second = self.read_u8()? as u32;
60 Ok(((first | (second << 8)) >> 2) & 0x3FFF)
61 }
62 0b10 => {
63 self.need(3)?;
64 let b1 = self.data[self.pos] as u32;
65 let b2 = self.data[self.pos + 1] as u32;
66 let b3 = self.data[self.pos + 2] as u32;
67 self.pos += 3;
68 let val = first | (b1 << 8) | (b2 << 16) | (b3 << 24);
69 Ok(val >> 2)
70 }
71 0b11 => {
72 let byte_count = (first >> 2) + 4;
73 if byte_count > 4 {
74 return Err(DecodeErr::CompactTooLarge);
75 }
76 let mut val = 0u32;
77 for i in 0..byte_count as usize {
78 val |= (self.read_u8()? as u32) << (i * 8);
79 }
80 Ok(val)
81 }
82 _ => unreachable!(),
83 }
84 }
85
86 pub fn read_string(&mut self) -> Result<String, DecodeErr> {
88 let len = self.read_compact_u32()? as usize;
89 let bytes = self.read_raw(len)?;
90 String::from_utf8(bytes.to_vec()).map_err(|_| DecodeErr::InvalidUtf8)
91 }
92
93 pub fn read_var_bytes(&mut self) -> Result<Vec<u8>, DecodeErr> {
95 let len = self.read_compact_u32()? as usize;
96 Ok(self.read_raw(len)?.to_vec())
97 }
98
99 pub fn read_fixed_bytes(&mut self, n: usize) -> Result<Vec<u8>, DecodeErr> {
101 Ok(self.read_raw(n)?.to_vec())
102 }
103
104 pub fn read_option<T>(
106 &mut self,
107 f: impl FnOnce(&mut Self) -> Result<T, DecodeErr>,
108 ) -> Result<Option<T>, DecodeErr> {
109 match self.read_u8()? {
110 0 => Ok(None),
111 1 => f(self).map(Some),
112 _ => Err(DecodeErr::InvalidOption),
113 }
114 }
115
116 pub fn skip_rest(&mut self) {
118 self.pos = self.data.len();
119 }
120}
121
122#[derive(Debug)]
123pub enum DecodeErr {
124 Eof,
125 CompactTooLarge,
126 InvalidUtf8,
127 InvalidOption,
128 InvalidTag(u8),
129 BadMessage(&'static str),
130}
131
132impl std::fmt::Display for DecodeErr {
133 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134 match self {
135 Self::Eof => write!(f, "unexpected end of input"),
136 Self::CompactTooLarge => write!(f, "compact integer exceeds u32"),
137 Self::InvalidUtf8 => write!(f, "invalid UTF-8 in string"),
138 Self::InvalidOption => write!(f, "invalid option discriminant"),
139 Self::InvalidTag(t) => write!(f, "invalid tag: {t}"),
140 Self::BadMessage(msg) => write!(f, "{msg}"),
141 }
142 }
143}
144
145impl std::error::Error for DecodeErr {}
146
147pub fn encode_compact_u32(buf: &mut Vec<u8>, val: u32) {
153 if val < 0x40 {
154 buf.push((val as u8) << 2);
155 } else if val < 0x4000 {
156 let v = (val << 2) | 0b01;
157 buf.push(v as u8);
158 buf.push((v >> 8) as u8);
159 } else if val < 0x4000_0000 {
160 let v = (val << 2) | 0b10;
161 buf.push(v as u8);
162 buf.push((v >> 8) as u8);
163 buf.push((v >> 16) as u8);
164 buf.push((v >> 24) as u8);
165 } else {
166 buf.push(0b11); buf.push(val as u8);
168 buf.push((val >> 8) as u8);
169 buf.push((val >> 16) as u8);
170 buf.push((val >> 24) as u8);
171 }
172}
173
174pub fn encode_string(buf: &mut Vec<u8>, s: &str) {
176 encode_compact_u32(buf, s.len() as u32);
177 buf.extend_from_slice(s.as_bytes());
178}
179
180pub fn encode_tag(buf: &mut Vec<u8>, tag: u8) {
182 buf.push(tag);
183}
184
185pub fn encode_result_ok_void(buf: &mut Vec<u8>) {
187 buf.push(0x00);
188}
189
190pub fn encode_result_ok(buf: &mut Vec<u8>) {
192 buf.push(0x00);
193}
194
195pub fn encode_result_err(buf: &mut Vec<u8>) {
197 buf.push(0x01);
198}
199
200pub fn encode_var_bytes(buf: &mut Vec<u8>, data: &[u8]) {
202 encode_compact_u32(buf, data.len() as u32);
203 buf.extend_from_slice(data);
204}
205
206pub fn encode_option_none(buf: &mut Vec<u8>) {
208 buf.push(0x00);
209}
210
211pub fn encode_option_some(buf: &mut Vec<u8>) {
213 buf.push(0x01);
214}
215
216pub fn encode_vector_len(buf: &mut Vec<u8>, count: u32) {
218 encode_compact_u32(buf, count);
219}
220
221#[cfg(test)]
222mod tests {
223 use super::*;
224
225 #[test]
226 fn compact_round_trip() {
227 for val in [0u32, 1, 63, 64, 16383, 16384, 1_073_741_823, u32::MAX] {
228 let mut buf = Vec::new();
229 encode_compact_u32(&mut buf, val);
230 let mut r = Reader::new(&buf);
231 let decoded = r.read_compact_u32().unwrap();
232 assert_eq!(val, decoded, "compact round-trip failed for {val}");
233 assert_eq!(r.pos, buf.len());
234 }
235 }
236
237 #[test]
238 fn string_round_trip() {
239 for s in [
240 "",
241 "hello",
242 "dot://mytestapp.dot",
243 "a".repeat(1000).as_str(),
244 ] {
245 let mut buf = Vec::new();
246 encode_string(&mut buf, s);
247 let mut r = Reader::new(&buf);
248 let decoded = r.read_string().unwrap();
249 assert_eq!(s, decoded);
250 }
251 }
252
253 #[test]
254 fn var_bytes_round_trip() {
255 let data = vec![0xde, 0xad, 0xbe, 0xef];
256 let mut buf = Vec::new();
257 encode_var_bytes(&mut buf, &data);
258 let mut r = Reader::new(&buf);
259 let decoded = r.read_var_bytes().unwrap();
260 assert_eq!(data, decoded);
261 }
262
263 #[test]
264 fn var_bytes_empty() {
265 let mut buf = Vec::new();
266 encode_var_bytes(&mut buf, &[]);
267 assert_eq!(buf, vec![0x00]); let mut r = Reader::new(&buf);
269 let decoded = r.read_var_bytes().unwrap();
270 assert!(decoded.is_empty());
271 }
272
273 #[test]
274 fn reader_eof_on_empty() {
275 let mut r = Reader::new(&[]);
276 assert!(r.read_u8().is_err());
277 assert!(r.read_string().is_err());
278 assert!(r.read_compact_u32().is_err());
279 }
280
281 #[test]
282 fn reader_truncated_string() {
283 let mut buf = Vec::new();
285 encode_compact_u32(&mut buf, 10);
286 buf.extend_from_slice(b"abc");
287 let mut r = Reader::new(&buf);
288 assert!(r.read_string().is_err());
289 }
290}