Skip to main content

dynamic/
msgpack.rs

1pub trait MsgPack {
2    fn encode(&self, buf: &mut Vec<u8>);
3}
4
5use std::collections::BTreeMap;
6
7use byteorder::{BigEndian, WriteBytesExt};
8use smol_str::SmolStr;
9impl MsgPack for &str {
10    fn encode(&self, buf: &mut Vec<u8>) {
11        let length = self.len();
12        if length < 0x20 {
13            buf.push(0xa0 | length as u8);
14        } else if length < 0x100 {
15            buf.push(0xd9);
16            buf.push(length as u8);
17        } else if length < 0x10000 {
18            buf.push(0xda);
19            buf.write_u16::<BigEndian>(length as u16).unwrap();
20        } else {
21            buf.push(0xdb);
22            buf.write_u32::<BigEndian>(length as u32).unwrap();
23        }
24        buf.extend_from_slice(self.as_bytes());
25    }
26}
27
28impl MsgPack for i64 {
29    fn encode(&self, buf: &mut Vec<u8>) {
30        let value = *self;
31        if value >= 0 && value < 128 {
32            buf.push(value as u8);
33        } else if value < 0 && value > -32 {
34            let raw = (value as i8) as u8;
35            buf.push(raw);
36        } else {
37            if value >= -0x80 && value <= 0x7f {
38                buf.push(0xd0);
39                buf.write_i8(value as i8).unwrap();
40            } else if value >= -0x8000 && value <= 0x7fff {
41                buf.push(0xd1);
42                buf.write_i16::<BigEndian>(value as i16).unwrap();
43            } else if value >= -0x8000_0000 && value <= 0x7fff_ffff {
44                buf.push(0xd2);
45                buf.write_i32::<BigEndian>(value as i32).unwrap();
46            } else {
47                buf.push(0xd3);
48                buf.write_i64::<BigEndian>(value).unwrap();
49            }
50        }
51    }
52}
53
54use super::Dynamic;
55
56fn encode_bytes(raw: &[u8], buf: &mut Vec<u8>) {
57    let length = raw.len();
58    if length < 0x100 {
59        buf.push(0xc4);
60        buf.push(length as u8);
61    } else if length < 0x10000 {
62        buf.push(0xc5);
63        buf.write_u16::<BigEndian>(length as u16).unwrap();
64    } else {
65        buf.push(0xc6);
66        buf.write_u32::<BigEndian>(length as u32).unwrap();
67    }
68    buf.extend_from_slice(&raw);
69}
70
71fn encode_vec(raw: &[u8], len: usize, tag: u8, buf: &mut Vec<u8>) {
72    if len < 0x100 {
73        buf.push(0xc7);
74        buf.push(len as u8);
75        buf.push(tag);
76    } else if len < 0x10000 {
77        buf.push(0xc8);
78        buf.write_u16::<BigEndian>(len as u16).unwrap();
79        buf.push(tag);
80    } else {
81        buf.push(0xc9);
82        buf.write_u32::<BigEndian>(len as u32).unwrap();
83        buf.push(tag);
84    }
85    buf.extend_from_slice(raw);
86}
87
88impl MsgPack for Dynamic {
89    fn encode(&self, buf: &mut Vec<u8>) {
90        match self {
91            Dynamic::Iter { idx: _, keys: _, value: _ } => {}
92            Dynamic::Null => buf.push(0xc0),
93            Dynamic::Bool(b) => buf.push(if *b { 0xc3 } else { 0xc2 }),
94            Dynamic::I8(i) => (*i as i64).encode(buf),
95            Dynamic::I16(i) => (*i as i64).encode(buf),
96            Dynamic::I32(i) => (*i as i64).encode(buf),
97            Dynamic::I64(i) => (*i).encode(buf),
98            Dynamic::U8(i) => (*i as i64).encode(buf),
99            Dynamic::U16(i) => (*i as i64).encode(buf),
100            Dynamic::U32(i) => (*i as i64).encode(buf),
101            Dynamic::U64(i) => (*i as i64).encode(buf),
102            Dynamic::F32(f) => {
103                buf.push(0xca);
104                let int_value = f32::to_bits(*f);
105                buf.write_u32::<BigEndian>(int_value).unwrap();
106            }
107            Dynamic::F64(f) => {
108                buf.push(0xcb);
109                let int_value = f64::to_bits(*f);
110                buf.write_u64::<BigEndian>(int_value).unwrap();
111            }
112            Dynamic::String(s) => s.as_str().encode(buf),
113            Dynamic::StringBuf(s) => s.as_str().encode(buf),
114            Dynamic::Bytes(raw) => {
115                encode_bytes(raw.as_slice(), buf);
116            }
117            Dynamic::VecI8(vec) => {
118                let len = self.len();
119                encode_vec(vec.as_slice(), len, 1, buf);
120            }
121            Dynamic::VecU16(vec) => {
122                let len = self.len();
123                encode_vec(vec.as_slice(), len, 2, buf);
124            }
125            Dynamic::VecI16(vec) => {
126                let len = self.len();
127                encode_vec(vec.as_slice(), len, 3, buf);
128            }
129            Dynamic::VecU32(vec) => {
130                let len = self.len();
131                encode_vec(vec.as_slice(), len, 4, buf);
132            }
133            Dynamic::VecI32(vec) => {
134                let len = self.len();
135                encode_vec(vec.as_slice(), len, 5, buf);
136            }
137            Dynamic::VecF32(vec) => {
138                let len = self.len();
139                encode_vec(vec.as_slice(), len, 6, buf);
140            }
141            Dynamic::VecU64(vec) => {
142                let len = self.len();
143                encode_vec(bytemuck::cast_slice(vec.as_slice()), len, 7, buf);
144            }
145            Dynamic::VecI64(vec) => {
146                let len = self.len();
147                encode_vec(bytemuck::cast_slice(vec.as_slice()), len, 8, buf);
148            }
149            Dynamic::VecF64(vec) => {
150                let len = self.len();
151                encode_vec(bytemuck::cast_slice(vec.as_slice()), len, 9, buf);
152            }
153            Dynamic::List(raw) => {
154                let length = raw.read().unwrap().len();
155                if length < 0x10 {
156                    buf.push(0x90 | length as u8);
157                } else if length < 0x10000 {
158                    buf.push(0xdc);
159                    buf.write_u16::<BigEndian>(length as u16).unwrap();
160                } else {
161                    buf.push(0xdd);
162                    buf.write_u32::<BigEndian>(length as u32).unwrap();
163                }
164                for item in raw.read().unwrap().iter() {
165                    item.encode(buf);
166                }
167            }
168            Dynamic::Map(raw) => {
169                let length = raw.read().unwrap().len();
170                if length < 16 {
171                    buf.push(0x80 | length as u8);
172                } else if length <= 0x10000 {
173                    buf.push(0xde);
174                    buf.write_u16::<BigEndian>(length as u16).unwrap();
175                } else {
176                    buf.push(0xdf);
177                    buf.write_u32::<BigEndian>(length as u32).unwrap();
178                }
179                for (k, v) in raw.read().unwrap().iter() {
180                    k.as_str().encode(buf);
181                    v.encode(buf);
182                }
183            }
184            Dynamic::Struct { .. } => {
185                let keys = self.keys();
186                let length = keys.len();
187                if length < 16 {
188                    buf.push(0x80 | length as u8);
189                } else if length <= 0x10000 {
190                    buf.push(0xde);
191                    buf.write_u16::<BigEndian>(length as u16).unwrap();
192                } else {
193                    buf.push(0xdf);
194                    buf.write_u32::<BigEndian>(length as u32).unwrap();
195                }
196                for key in keys {
197                    key.as_str().encode(buf);
198                    self.get_dynamic(key.as_str()).unwrap_or(Dynamic::Null).encode(buf);
199                }
200            }
201            Dynamic::Custom(value) => {
202                buf.push(0x81);
203                "@custom".encode(buf);
204                value.custom_type_name().encode(buf);
205            }
206        }
207    }
208}
209use anyhow::{Result, anyhow};
210
211pub trait MsgUnpack: Sized {
212    fn decode(buf: &[u8]) -> Result<(Self, usize)>;
213    fn decode_array(buf: &[u8], length: usize) -> Result<(Vec<Self>, usize)> {
214        let mut cursor = 0usize;
215        let mut result = Vec::with_capacity(length);
216        for _ in 0..length {
217            let (value, size) = Self::decode(&buf[cursor..])?;
218            result.push(value);
219            cursor += size;
220        }
221        Ok((result, cursor))
222    }
223}
224
225#[inline]
226pub(crate) fn read_8(raw: &[u8]) -> u8 {
227    raw[0]
228}
229
230#[inline]
231pub(crate) fn read_16(raw: &[u8]) -> u16 {
232    raw[1] as u16 | (raw[0] as u16) << 8
233}
234
235#[inline]
236pub(crate) fn read_32(raw: &[u8]) -> u32 {
237    raw[3] as u32 | (raw[2] as u32) << 8 | (raw[1] as u32) << 16 | (raw[0] as u32) << 24
238}
239
240#[inline]
241pub(crate) fn read_64(raw: &[u8]) -> u64 {
242    raw[7] as u64 | (raw[6] as u64) << 8 | (raw[5] as u64) << 16 | (raw[4] as u64) << 24 | (raw[3] as u64) << 32 | (raw[2] as u64) << 40 | (raw[1] as u64) << 48 | (raw[0] as u64) << 56
243}
244
245fn vec_to_map(kvs: Vec<Dynamic>) -> Result<Dynamic> {
246    let mut map: BTreeMap<SmolStr, Dynamic> = BTreeMap::new();
247    let mut key: Option<Dynamic> = None;
248    for kv in kvs {
249        if let Some(k) = key.take() {
250            map.insert(SmolStr::from(k.to_string()), kv);
251        } else {
252            key = Some(kv);
253        }
254    }
255    Ok(Dynamic::Map(Arc::new(RwLock::new(map))))
256}
257
258const TAG_LEN: [usize; 10] = [0, 1, 2, 2, 4, 4, 4, 8, 8, 8];
259
260use bytemuck::pod_collect_to_vec;
261
262fn to_vec(buf: &[u8], tag: u8) -> Result<Dynamic> {
263    match tag {
264        1 => Ok(Dynamic::from(bytemuck::cast_slice::<_, i8>(buf))),
265        2 => {
266            let aligned = pod_collect_to_vec::<_, u16>(buf);
267            Ok(Dynamic::from(aligned.as_slice()))
268        }
269        3 => {
270            let aligned = pod_collect_to_vec::<_, i16>(buf);
271            Ok(Dynamic::from(aligned.as_slice()))
272        }
273        4 => {
274            let aligned = pod_collect_to_vec::<_, u32>(buf);
275            Ok(Dynamic::from(aligned.as_slice()))
276        }
277        5 => {
278            let aligned = pod_collect_to_vec::<_, i32>(buf);
279            Ok(Dynamic::from(aligned.as_slice()))
280        }
281        6 => {
282            let aligned = pod_collect_to_vec::<_, f32>(buf);
283            Ok(Dynamic::from(aligned.as_slice()))
284        }
285        7 => {
286            let aligned = pod_collect_to_vec::<_, u64>(buf);
287            Ok(Dynamic::from(aligned.as_slice()))
288        }
289        8 => {
290            let aligned = pod_collect_to_vec::<_, i64>(buf);
291            Ok(Dynamic::from(aligned.as_slice()))
292        }
293        9 => {
294            let aligned = pod_collect_to_vec::<_, f64>(buf);
295            Ok(Dynamic::from(aligned.as_slice()))
296        }
297        _ => Err(anyhow!("unknow tag {}", tag)),
298    }
299}
300
301use std::sync::{Arc, RwLock};
302
303impl MsgUnpack for Dynamic {
304    fn decode(buf: &[u8]) -> Result<(Self, usize)> {
305        assert_err!(buf.len() < 1, anyhow!("no data"));
306        let first_byte = buf[0];
307        assert_ok!(first_byte <= 0x7f, (Dynamic::from(first_byte as i64), 1));
308        assert_ok!(first_byte >= 0xe0, (Dynamic::from(first_byte as i64 - 256), 1));
309        if first_byte >= 0x80 && first_byte <= 0x8f {
310            let len = (first_byte & 0x0f) as usize;
311            let (value, size) = Self::decode_array(&buf[1..], len * 2)?;
312            return vec_to_map(value).map(|r| (r, 1 + size));
313        }
314        if first_byte >= 0x90 && first_byte <= 0x9f {
315            let len = (first_byte & 0x0f) as usize;
316            let (value, size) = Self::decode_array(&buf[1..], len)?;
317            return Ok((Dynamic::List(Arc::new(RwLock::new(value))), 1 + size));
318        }
319
320        if first_byte >= 0xa0 && first_byte <= 0xbf {
321            let len = (first_byte & 0x1f) as usize;
322            assert_err!(buf.len() < 1 + len, anyhow!("no data"));
323            return Ok((Dynamic::from_utf8(&buf[1..1 + len])?, 1 + len));
324        }
325
326        assert_ok!(first_byte == 0xc0, (Dynamic::Null, 1));
327        assert_err!(first_byte == 0xc1, anyhow!("0xc1 never used"));
328        assert_ok!(first_byte == 0xc2, (false.into(), 1));
329        assert_ok!(first_byte == 0xc3, (true.into(), 1));
330
331        if first_byte == 0xc4 {
332            assert_err!(buf.len() < 2, anyhow!("no data"));
333            let len = read_8(&buf[1..]) as usize;
334            assert_err!(buf.len() < 2 + len, anyhow!("no data"));
335            return Ok((Dynamic::from(&buf[2..2 + len]), 2 + len));
336        }
337
338        if first_byte == 0xc5 {
339            assert_err!(buf.len() < 3, anyhow!("no data"));
340            let len = read_16(&buf[1..]) as usize;
341            assert_err!(buf.len() < 2 + len, anyhow!("no data"));
342            return Ok((Dynamic::from(&buf[3..3 + len]), 3 + len));
343        }
344
345        if first_byte == 0xc6 {
346            assert_err!(buf.len() < 5, anyhow!("no data"));
347            let len = read_32(&buf[1..]) as usize;
348            assert_err!(buf.len() < 5 + len, anyhow!("no data"));
349            return Ok((Dynamic::from(&buf[5..5 + len]), 5 + len));
350        }
351
352        if first_byte == 0xc7 {
353            assert_err!(buf.len() < 7, anyhow!("no data"));
354            let len = read_8(&buf[1..]) as usize;
355            let tag = read_8(&buf[2..]);
356            let byte_len = len * TAG_LEN[tag as usize];
357            return Ok((to_vec(&buf[3..3 + byte_len], tag)?, 3 + byte_len));
358        }
359
360        if first_byte == 0xc8 {
361            assert_err!(buf.len() < 8, anyhow!("no data"));
362            let len = read_16(&buf[1..]) as usize;
363            let tag = read_8(&buf[3..]);
364            let byte_len = len * TAG_LEN[tag as usize];
365            return Ok((to_vec(&buf[4..4 + byte_len], tag)?, 4 + byte_len));
366        }
367
368        if first_byte == 0xc9 {
369            assert_err!(buf.len() < 10, anyhow!("no data"));
370            let len = read_32(&buf[1..]) as usize;
371            let tag = read_8(&buf[5..]);
372            let byte_len = len * TAG_LEN[tag as usize];
373            return Ok((to_vec(&buf[6..6 + byte_len], tag)?, 6 + byte_len));
374        }
375
376        if first_byte == 0xca {
377            assert_err!(buf.len() < 5, anyhow!("no data"));
378            let raw_value = read_32(&buf[1..]) as u32;
379            let value = f32::from_bits(raw_value);
380            return Ok((Dynamic::from(value as f64), 5));
381        }
382
383        if first_byte == 0xcb {
384            assert_err!(buf.len() < 9, anyhow!("no data"));
385            let raw_value = read_64(&buf[1..]);
386            let value = f64::from_bits(raw_value);
387            return Ok((Dynamic::from(value), 9));
388        }
389
390        if first_byte == 0xd0 {
391            assert_err!(buf.len() < 2, anyhow!("no data"));
392            let raw_value = read_8(&buf[1..]);
393            let value = raw_value as i8 as i64;
394            return Ok((Dynamic::from(value), 2));
395        }
396
397        if first_byte == 0xcd {
398            assert_err!(buf.len() < 3, anyhow!("no data"));
399            let value = read_16(&buf[1..]);
400            return Ok((Dynamic::from(value as i64), 3));
401        }
402
403        if first_byte == 0xce {
404            assert_err!(buf.len() < 5, anyhow!("no data"));
405            let value = read_32(&buf[1..]);
406            return Ok((Dynamic::from(value as i64), 5));
407        }
408
409        if first_byte == 0xcf {
410            assert_err!(buf.len() < 9, anyhow!("no data"));
411            let value = read_64(&buf[1..]);
412            return Ok((Dynamic::from(value as i64), 9));
413        }
414
415        if first_byte == 0xd3 {
416            assert_err!(buf.len() < 9, anyhow!("no data"));
417            let raw_value = read_64(&buf[1..]);
418            let value = raw_value as i64;
419            return Ok((Dynamic::from(value), 9));
420        }
421
422        if first_byte == 0xd1 {
423            assert_err!(buf.len() < 3, anyhow!("no data"));
424            let raw_value = read_16(&buf[1..]);
425            let value = u16::cast_signed(raw_value) as i64;
426            return Ok((Dynamic::from(value), 3));
427        }
428
429        if first_byte == 0xd2 {
430            assert_err!(buf.len() < 5, anyhow!("no data"));
431            let raw_value = read_32(&buf[1..]);
432            let value = u32::cast_signed(raw_value) as i64;
433            return Ok((Dynamic::from(value), 5));
434        }
435
436        if first_byte == 0xd3 {
437            assert_err!(buf.len() < 9, anyhow!("no data"));
438            let raw_value = read_64(&buf[1..]);
439            let value = u64::cast_signed(raw_value);
440            return Ok((Dynamic::from(value), 9));
441        }
442
443        if first_byte == 0xd4 {
444            assert_err!(buf.len() < 3, anyhow!("no data"));
445            let _type_id = u8::cast_signed(buf[1]);
446            //let _value = raw[2..3].to_vec();
447            return Ok((Dynamic::Null, 3));
448        }
449
450        if first_byte == 0xd5 {
451            assert_err!(buf.len() < 4, anyhow!("no data"));
452            let _type_id = u8::cast_signed(buf[1]);
453            return Ok((Dynamic::Null, 4));
454        }
455
456        if first_byte == 0xd6 {
457            assert_err!(buf.len() < 6, anyhow!("no data"));
458            let _type_id = u8::cast_signed(buf[1]);
459            return Ok((Dynamic::Null, 6));
460        }
461
462        if first_byte == 0xd7 {
463            assert_err!(buf.len() < 10, anyhow!("no data"));
464            let _type_id = buf[1] as i8;
465            return Ok((Dynamic::Null, 10));
466        }
467
468        if first_byte == 0xd8 {
469            assert_err!(buf.len() < 18, anyhow!("no data"));
470            let _type_id = buf[1] as i8;
471            return Ok((Dynamic::Null, 18));
472        }
473
474        if first_byte == 0xd9 {
475            assert_err!(buf.len() < 2, anyhow!("no data"));
476            let len = read_8(&buf[1..]) as usize;
477            assert_err!(buf.len() < 2 + len, anyhow!("no data"));
478            return Ok((Dynamic::from_utf8(&buf[2..2 + len])?, 2 + len));
479        }
480
481        if first_byte == 0xda {
482            assert_err!(buf.len() < 3, anyhow!("no data"));
483            let len = read_16(&buf[1..]) as usize;
484            assert_err!(buf.len() < 3 + len, anyhow!("no data"));
485            return Ok((Dynamic::from_utf8(&buf[3..3 + len])?, 3 + len));
486        }
487
488        if first_byte == 0xdb {
489            assert_err!(buf.len() < 5, anyhow!("no data"));
490            let len = read_32(&buf[1..]) as usize;
491            assert_err!(buf.len() < 5 + len, anyhow!("no data"));
492            return Ok((Dynamic::from_utf8(&buf[5..5 + len])?, 5 + len));
493        }
494
495        if first_byte == 0xdc {
496            assert_err!(buf.len() < 3, anyhow!("no data"));
497            let len = read_16(&buf[1..]) as usize;
498            let (value, size) = Self::decode_array(&buf[3..], len)?;
499            return Ok((Dynamic::List(Arc::new(RwLock::new(value))), 3 + size));
500        }
501
502        if first_byte == 0xdd {
503            assert_err!(buf.len() < 5, anyhow!("no data"));
504            let len = read_32(&buf[1..]) as usize;
505            let (value, size) = Self::decode_array(&buf[5..], len)?;
506            return Ok((Dynamic::List(Arc::new(RwLock::new(value))), 5 + size));
507        }
508
509        if first_byte == 0xde {
510            assert_err!(buf.len() < 3, anyhow!("no data"));
511            let len = read_16(&buf[1..]) as usize;
512            let (value, size) = Self::decode_array(&buf[3..], len * 2)?;
513            return vec_to_map(value).map(|r| (r, 3 + size));
514        }
515
516        if first_byte == 0xdf {
517            assert_err!(buf.len() < 5, anyhow!("no data"));
518            let len = read_32(&buf[1..]) as usize;
519            let (value, size) = Self::decode_array(&buf[5..], len * 2)?;
520            return vec_to_map(value).map(|r| (r, 5 + size));
521        }
522        Err(anyhow!("error code {}", first_byte))
523    }
524}