1use std::collections::BTreeMap;
2
3use crate::error::{AppError, Result};
4
5#[derive(Debug, Clone, PartialEq)]
6pub enum MsgValue {
7 Nil,
8 Bool(bool),
9 Int(i64),
10 UInt(u64),
11 Float(f64),
12 String(String),
13 Array(Vec<MsgValue>),
14 Map(BTreeMap<String, MsgValue>),
15}
16
17impl MsgValue {
18 pub fn get(&self, key: &str) -> Option<&MsgValue> {
19 match self {
20 MsgValue::Map(map) => map.get(key),
21 _ => None,
22 }
23 }
24
25 pub fn as_str(&self) -> Option<&str> {
26 match self {
27 MsgValue::String(value) => Some(value.as_str()),
28 _ => None,
29 }
30 }
31
32 pub fn as_array(&self) -> Option<&[MsgValue]> {
33 match self {
34 MsgValue::Array(values) => Some(values.as_slice()),
35 _ => None,
36 }
37 }
38
39 pub fn as_map(&self) -> Option<&BTreeMap<String, MsgValue>> {
40 match self {
41 MsgValue::Map(map) => Some(map),
42 _ => None,
43 }
44 }
45
46 pub fn as_u64(&self) -> Option<u64> {
47 match self {
48 MsgValue::UInt(value) => Some(*value),
49 MsgValue::Int(value) if *value >= 0 => Some(*value as u64),
50 _ => None,
51 }
52 }
53}
54
55pub fn encode(value: &MsgValue) -> Vec<u8> {
56 let mut buffer = Vec::new();
57 encode_value(value, &mut buffer);
58 buffer
59}
60
61pub fn decode(bytes: &[u8]) -> Result<MsgValue> {
62 let (value, consumed) = decode_value(bytes, 0)?;
63 if consumed != bytes.len() {
64 return Err(AppError::Msgpack(
65 "trailing bytes in msgpack payload".to_string(),
66 ));
67 }
68 Ok(value)
69}
70
71fn encode_value(value: &MsgValue, buffer: &mut Vec<u8>) {
72 match value {
73 MsgValue::Nil => buffer.push(0xc0),
74 MsgValue::Bool(false) => buffer.push(0xc2),
75 MsgValue::Bool(true) => buffer.push(0xc3),
76 MsgValue::Int(value) => encode_int(*value, buffer),
77 MsgValue::UInt(value) => encode_uint(*value, buffer),
78 MsgValue::Float(value) => {
79 buffer.push(0xcb);
80 buffer.extend_from_slice(&value.to_be_bytes());
81 }
82 MsgValue::String(value) => encode_str(value, buffer),
83 MsgValue::Array(values) => encode_array(values, buffer),
84 MsgValue::Map(map) => encode_map(map, buffer),
85 }
86}
87
88fn encode_uint(value: u64, buffer: &mut Vec<u8>) {
89 match value {
90 0..=127 => buffer.push(value as u8),
91 128..=255 => {
92 buffer.push(0xcc);
93 buffer.push(value as u8);
94 }
95 256..=65_535 => {
96 buffer.push(0xcd);
97 buffer.extend_from_slice(&(value as u16).to_be_bytes());
98 }
99 65_536..=4_294_967_295 => {
100 buffer.push(0xce);
101 buffer.extend_from_slice(&(value as u32).to_be_bytes());
102 }
103 _ => {
104 buffer.push(0xcf);
105 buffer.extend_from_slice(&value.to_be_bytes());
106 }
107 }
108}
109
110fn encode_int(value: i64, buffer: &mut Vec<u8>) {
111 if (-32..=127).contains(&value) {
112 if value >= 0 {
113 buffer.push(value as u8);
114 } else {
115 buffer.push((value as i8) as u8);
116 }
117 return;
118 }
119
120 if (-128..=-33).contains(&value) {
121 buffer.push(0xd0);
122 buffer.push(value as i8 as u8);
123 } else if (-32_768..=-129).contains(&value) {
124 buffer.push(0xd1);
125 buffer.extend_from_slice(&(value as i16).to_be_bytes());
126 } else if (-2_147_483_648..=-32_769).contains(&value) {
127 buffer.push(0xd2);
128 buffer.extend_from_slice(&(value as i32).to_be_bytes());
129 } else {
130 buffer.push(0xd3);
131 buffer.extend_from_slice(&value.to_be_bytes());
132 }
133}
134
135fn encode_str(value: &str, buffer: &mut Vec<u8>) {
136 let bytes = value.as_bytes();
137 let len = bytes.len();
138 if len <= 31 {
139 buffer.push(0xa0 | (len as u8));
140 } else if len <= 255 {
141 buffer.push(0xd9);
142 buffer.push(len as u8);
143 } else if len <= 65_535 {
144 buffer.push(0xda);
145 buffer.extend_from_slice(&(len as u16).to_be_bytes());
146 } else {
147 buffer.push(0xdb);
148 buffer.extend_from_slice(&(len as u32).to_be_bytes());
149 }
150 buffer.extend_from_slice(bytes);
151}
152
153fn encode_array(values: &[MsgValue], buffer: &mut Vec<u8>) {
154 let len = values.len();
155 if len <= 15 {
156 buffer.push(0x90 | (len as u8));
157 } else if len <= 65_535 {
158 buffer.push(0xdc);
159 buffer.extend_from_slice(&(len as u16).to_be_bytes());
160 } else {
161 buffer.push(0xdd);
162 buffer.extend_from_slice(&(len as u32).to_be_bytes());
163 }
164 for value in values {
165 encode_value(value, buffer);
166 }
167}
168
169fn encode_map(map: &BTreeMap<String, MsgValue>, buffer: &mut Vec<u8>) {
170 let len = map.len();
171 if len <= 15 {
172 buffer.push(0x80 | (len as u8));
173 } else if len <= 65_535 {
174 buffer.push(0xde);
175 buffer.extend_from_slice(&(len as u16).to_be_bytes());
176 } else {
177 buffer.push(0xdf);
178 buffer.extend_from_slice(&(len as u32).to_be_bytes());
179 }
180 for (key, value) in map {
181 encode_str(key, buffer);
182 encode_value(value, buffer);
183 }
184}
185
186fn decode_value(bytes: &[u8], index: usize) -> Result<(MsgValue, usize)> {
187 let prefix = *bytes
188 .get(index)
189 .ok_or_else(|| AppError::Msgpack("unexpected end of msgpack payload".to_string()))?;
190
191 match prefix {
192 0xc0 => Ok((MsgValue::Nil, index + 1)),
193 0xc2 => Ok((MsgValue::Bool(false), index + 1)),
194 0xc3 => Ok((MsgValue::Bool(true), index + 1)),
195 0xc4 => {
196 let len = read_u8(bytes, index + 1)? as usize;
197 let slice = bytes.get(index + 2..index + 2 + len).ok_or_else(|| {
198 AppError::Msgpack("unexpected end of msgpack binary data".to_string())
199 })?;
200 let value = String::from_utf8_lossy(slice).into_owned();
201 Ok((MsgValue::String(value), index + 2 + len))
202 }
203 0xc5 => {
204 let len = read_u16(bytes, index + 1)? as usize;
205 let slice = bytes.get(index + 3..index + 3 + len).ok_or_else(|| {
206 AppError::Msgpack("unexpected end of msgpack binary data".to_string())
207 })?;
208 let value = String::from_utf8_lossy(slice).into_owned();
209 Ok((MsgValue::String(value), index + 3 + len))
210 }
211 0xc6 => {
212 let len = read_u32(bytes, index + 1)? as usize;
213 let slice = bytes.get(index + 5..index + 5 + len).ok_or_else(|| {
214 AppError::Msgpack("unexpected end of msgpack binary data".to_string())
215 })?;
216 let value = String::from_utf8_lossy(slice).into_owned();
217 Ok((MsgValue::String(value), index + 5 + len))
218 }
219 0xca => {
220 let end = index + 5;
221 let slice = bytes.get(index + 1..end).ok_or_else(|| {
222 AppError::Msgpack("unexpected end of msgpack payload".to_string())
223 })?;
224 let value = f32::from_be_bytes([slice[0], slice[1], slice[2], slice[3]]);
225 Ok((MsgValue::Float(value as f64), end))
226 }
227 0xcb => {
228 let end = index + 9;
229 let slice = bytes.get(index + 1..end).ok_or_else(|| {
230 AppError::Msgpack("unexpected end of msgpack payload".to_string())
231 })?;
232 let value = f64::from_be_bytes([
233 slice[0], slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7],
234 ]);
235 Ok((MsgValue::Float(value), end))
236 }
237 0xcc => {
238 let value = read_u8(bytes, index + 1)? as u64;
239 Ok((MsgValue::UInt(value), index + 2))
240 }
241 0xcd => {
242 let value = read_u16(bytes, index + 1)? as u64;
243 Ok((MsgValue::UInt(value), index + 3))
244 }
245 0xce => {
246 let value = read_u32(bytes, index + 1)? as u64;
247 Ok((MsgValue::UInt(value), index + 5))
248 }
249 0xcf => {
250 let value = read_u64(bytes, index + 1)?;
251 Ok((MsgValue::UInt(value), index + 9))
252 }
253 0xd0 => {
254 let value = read_i8(bytes, index + 1)? as i64;
255 Ok((MsgValue::Int(value), index + 2))
256 }
257 0xd1 => {
258 let value = read_i16(bytes, index + 1)? as i64;
259 Ok((MsgValue::Int(value), index + 3))
260 }
261 0xd2 => {
262 let value = read_i32(bytes, index + 1)? as i64;
263 Ok((MsgValue::Int(value), index + 5))
264 }
265 0xd3 => {
266 let value = read_i64(bytes, index + 1)?;
267 Ok((MsgValue::Int(value), index + 9))
268 }
269 0xd9 => {
270 let len = read_u8(bytes, index + 1)? as usize;
271 decode_str(bytes, index + 2, len)
272 }
273 0xda => {
274 let len = read_u16(bytes, index + 1)? as usize;
275 decode_str(bytes, index + 3, len)
276 }
277 0xdb => {
278 let len = read_u32(bytes, index + 1)? as usize;
279 decode_str(bytes, index + 5, len)
280 }
281 0xdc => {
282 let len = read_u16(bytes, index + 1)? as usize;
283 decode_array(bytes, index + 3, len)
284 }
285 0xdd => {
286 let len = read_u32(bytes, index + 1)? as usize;
287 decode_array(bytes, index + 5, len)
288 }
289 0xde => {
290 let len = read_u16(bytes, index + 1)? as usize;
291 decode_map(bytes, index + 3, len)
292 }
293 0xdf => {
294 let len = read_u32(bytes, index + 1)? as usize;
295 decode_map(bytes, index + 5, len)
296 }
297 0x80..=0x8f => {
298 let len = (prefix & 0x0f) as usize;
299 decode_map(bytes, index + 1, len)
300 }
301 0x90..=0x9f => {
302 let len = (prefix & 0x0f) as usize;
303 decode_array(bytes, index + 1, len)
304 }
305 0xa0..=0xbf => {
306 let len = (prefix & 0x1f) as usize;
307 decode_str(bytes, index + 1, len)
308 }
309 0xe0..=0xff => Ok((MsgValue::Int((prefix as i8) as i64), index + 1)),
310 0x00..=0x7f => Ok((MsgValue::UInt(prefix as u64), index + 1)),
311 _ => Err(AppError::Msgpack(format!(
312 "unsupported msgpack prefix: 0x{prefix:02x}"
313 ))),
314 }
315}
316
317fn decode_str(bytes: &[u8], index: usize, len: usize) -> Result<(MsgValue, usize)> {
318 let end = index
319 .checked_add(len)
320 .ok_or_else(|| AppError::Msgpack("msgpack string length overflow".to_string()))?;
321 let slice = bytes
322 .get(index..end)
323 .ok_or_else(|| AppError::Msgpack("unexpected end of msgpack string".to_string()))?;
324 let value = String::from_utf8(slice.to_vec())
325 .map_err(|e| AppError::Msgpack(format!("invalid utf-8 string: {e}")))?;
326 Ok((MsgValue::String(value), end))
327}
328
329fn decode_array(bytes: &[u8], mut index: usize, len: usize) -> Result<(MsgValue, usize)> {
330 let mut values = Vec::with_capacity(len);
331 for _ in 0..len {
332 let (value, next) = decode_value(bytes, index)?;
333 values.push(value);
334 index = next;
335 }
336 Ok((MsgValue::Array(values), index))
337}
338
339fn decode_map(bytes: &[u8], mut index: usize, len: usize) -> Result<(MsgValue, usize)> {
340 let mut map = BTreeMap::new();
341 for _ in 0..len {
342 let (key_value, next_key) = decode_value(bytes, index)?;
343 index = next_key;
344 let key = key_value
345 .as_str()
346 .ok_or_else(|| AppError::Msgpack("msgpack map key was not a string".to_string()))?;
347 let (value, next_value) = decode_value(bytes, index)?;
348 index = next_value;
349 map.insert(key.to_string(), value);
350 }
351 Ok((MsgValue::Map(map), index))
352}
353
354fn read_u8(bytes: &[u8], index: usize) -> Result<u8> {
355 bytes
356 .get(index)
357 .copied()
358 .ok_or_else(|| AppError::Msgpack("unexpected end of msgpack payload".to_string()))
359}
360
361fn read_i8(bytes: &[u8], index: usize) -> Result<i8> {
362 Ok(read_u8(bytes, index)? as i8)
363}
364
365fn read_u16(bytes: &[u8], index: usize) -> Result<u16> {
366 let end = index + 2;
367 let slice = bytes
368 .get(index..end)
369 .ok_or_else(|| AppError::Msgpack("unexpected end of msgpack payload".to_string()))?;
370 Ok(u16::from_be_bytes([slice[0], slice[1]]))
371}
372
373fn read_i16(bytes: &[u8], index: usize) -> Result<i16> {
374 Ok(i16::from_be_bytes(read_u16(bytes, index)?.to_be_bytes()))
375}
376
377fn read_u32(bytes: &[u8], index: usize) -> Result<u32> {
378 let end = index + 4;
379 let slice = bytes
380 .get(index..end)
381 .ok_or_else(|| AppError::Msgpack("unexpected end of msgpack payload".to_string()))?;
382 Ok(u32::from_be_bytes([slice[0], slice[1], slice[2], slice[3]]))
383}
384
385fn read_i32(bytes: &[u8], index: usize) -> Result<i32> {
386 Ok(i32::from_be_bytes(read_u32(bytes, index)?.to_be_bytes()))
387}
388
389fn read_u64(bytes: &[u8], index: usize) -> Result<u64> {
390 let end = index + 8;
391 let slice = bytes
392 .get(index..end)
393 .ok_or_else(|| AppError::Msgpack("unexpected end of msgpack payload".to_string()))?;
394 Ok(u64::from_be_bytes([
395 slice[0], slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7],
396 ]))
397}
398
399fn read_i64(bytes: &[u8], index: usize) -> Result<i64> {
400 Ok(i64::from_be_bytes(read_u64(bytes, index)?.to_be_bytes()))
401}
402
403#[cfg(test)]
404mod tests {
405 use super::*;
406
407 #[test]
408 fn roundtrip_nil() {
409 let v = MsgValue::Nil;
410 let bytes = encode(&v);
411 let decoded = decode(&bytes).unwrap();
412 assert_eq!(v, decoded);
413 }
414
415 #[test]
416 fn roundtrip_bool() {
417 for b in [true, false] {
418 let v = MsgValue::Bool(b);
419 let bytes = encode(&v);
420 let decoded = decode(&bytes).unwrap();
421 assert_eq!(v, decoded);
422 }
423 }
424
425 #[test]
426 fn roundtrip_positive_fixint() {
427 for i in [0u64, 1, 127] {
428 let v = MsgValue::UInt(i);
429 let bytes = encode(&v);
430 let decoded = decode(&bytes).unwrap();
431 assert_eq!(v, decoded);
432 }
433 }
434
435 #[test]
436 fn roundtrip_u8() {
437 let v = MsgValue::UInt(200);
438 let bytes = encode(&v);
439 assert_eq!(bytes[0], 0xcc);
440 assert_eq!(bytes[1], 200);
441 let decoded = decode(&bytes).unwrap();
442 assert_eq!(v, decoded);
443 }
444
445 #[test]
446 fn roundtrip_u16() {
447 let v = MsgValue::UInt(1000);
448 let bytes = encode(&v);
449 assert_eq!(bytes[0], 0xcd);
450 let decoded = decode(&bytes).unwrap();
451 assert_eq!(v, decoded);
452 }
453
454 #[test]
455 fn roundtrip_u32() {
456 let v = MsgValue::UInt(100_000);
457 let bytes = encode(&v);
458 assert_eq!(bytes[0], 0xce);
459 let decoded = decode(&bytes).unwrap();
460 assert_eq!(v, decoded);
461 }
462
463 #[test]
464 fn roundtrip_negative_fixint() {
465 for i in [-1i64, -32] {
466 let v = MsgValue::Int(i);
467 let bytes = encode(&v);
468 let decoded = decode(&bytes).unwrap();
469 assert_eq!(v, decoded);
470 }
471 }
472
473 #[test]
474 fn roundtrip_i8() {
475 let v = MsgValue::Int(-100);
476 let bytes = encode(&v);
477 assert_eq!(bytes[0], 0xd0);
478 let decoded = decode(&bytes).unwrap();
479 assert_eq!(v, decoded);
480 }
481
482 #[test]
483 fn roundtrip_fixstr() {
484 let v = MsgValue::String("hello".to_string());
485 let bytes = encode(&v);
486 assert_eq!(bytes[0], 0xa5);
487 let decoded = decode(&bytes).unwrap();
488 assert_eq!(v, decoded);
489 }
490
491 #[test]
492 fn roundtrip_str8() {
493 let s = "a".repeat(32);
494 let v = MsgValue::String(s);
495 let bytes = encode(&v);
496 assert_eq!(bytes[0], 0xd9);
497 assert_eq!(bytes[1], 32);
498 let decoded = decode(&bytes).unwrap();
499 assert_eq!(v, decoded);
500 }
501
502 #[test]
503 fn roundtrip_fixarray() {
504 let v = MsgValue::Array(vec![
505 MsgValue::UInt(1),
506 MsgValue::String("two".to_string()),
507 MsgValue::Bool(true),
508 ]);
509 let bytes = encode(&v);
510 assert_eq!(bytes[0], 0x93);
511 let decoded = decode(&bytes).unwrap();
512 assert_eq!(v, decoded);
513 }
514
515 #[test]
516 fn roundtrip_fixmap() {
517 let mut map = BTreeMap::new();
518 map.insert("key".to_string(), MsgValue::String("value".to_string()));
519 let v = MsgValue::Map(map);
520 let bytes = encode(&v);
521 assert_eq!(bytes[0], 0x81);
522 let decoded = decode(&bytes).unwrap();
523 assert_eq!(v, decoded);
524 }
525
526 #[test]
527 fn decode_bin8() {
528 let bytes = vec![0xc4, 0x05, 0x68, 0x65, 0x6c, 0x6c, 0x6f];
530 let decoded = decode(&bytes).unwrap();
531 assert_eq!(decoded, MsgValue::String("hello".to_string()));
532 }
533
534 #[test]
535 fn decode_bin16() {
536 let mut bytes = vec![0xc5, 0x00, 0x05, 0x77, 0x6f, 0x72, 0x6c, 0x64];
538 let decoded = decode(&bytes).unwrap();
539 assert_eq!(decoded, MsgValue::String("world".to_string()));
540
541 let long = "x".repeat(300);
543 bytes = vec![0xc5, 0x01, 0x2c];
544 bytes.extend_from_slice(long.as_bytes());
545 let decoded = decode(&bytes).unwrap();
546 assert_eq!(decoded, MsgValue::String(long));
547 }
548
549 #[test]
550 fn decode_bin32() {
551 let long = "z".repeat(70000);
552 let mut bytes = vec![0xc6, 0x00, 0x01, 0x11, 0x70];
553 bytes.extend_from_slice(long.as_bytes());
554 let decoded = decode(&bytes).unwrap();
555 assert_eq!(decoded, MsgValue::String(long));
556 }
557
558 #[test]
559 fn decode_float32() {
560 let mut bytes = vec![0xca];
561 bytes.extend_from_slice(&3.14f32.to_be_bytes());
562 let decoded = decode(&bytes).unwrap();
563 match decoded {
564 MsgValue::Float(f) => assert!((f - 3.14).abs() < 0.001),
565 _ => panic!("expected Float"),
566 }
567 }
568
569 #[test]
570 fn decode_float64() {
571 let mut bytes = vec![0xcb];
572 bytes.extend_from_slice(&std::f64::consts::PI.to_be_bytes());
573 let decoded = decode(&bytes).unwrap();
574 match decoded {
575 MsgValue::Float(f) => assert!((f - std::f64::consts::PI).abs() < 1e-15),
576 _ => panic!("expected Float"),
577 }
578 }
579
580 #[test]
581 fn decode_msf_auth_response() {
582 let mut bytes = Vec::new();
585 bytes.push(0x82);
587 bytes.extend_from_slice(&[0xc4, 0x06, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74]);
589 bytes.extend_from_slice(&[0xc4, 0x07, 0x73, 0x75, 0x63, 0x63, 0x65, 0x73, 0x73]);
591 bytes.extend_from_slice(&[0xc4, 0x05, 0x74, 0x6f, 0x6b, 0x65, 0x6e]);
593 bytes.extend_from_slice(&[0xc4, 0x08, 0x61, 0x62, 0x63, 0x64, 0x31, 0x32, 0x33, 0x34]);
595
596 let decoded = decode(&bytes).unwrap();
597 let result = decoded.get("result").and_then(|v| v.as_str()).unwrap();
598 let token = decoded.get("token").and_then(|v| v.as_str()).unwrap();
599 assert_eq!(result, "success");
600 assert_eq!(token, "abcd1234");
601 }
602
603 #[test]
604 fn decode_msf_error_response() {
605 let mut bytes = Vec::new();
608 bytes.push(0x82); bytes.extend_from_slice(&[0xc4, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72]); bytes.push(0xc3); bytes.extend_from_slice(&[
612 0xc4, 0x0c, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x6e, 0x67,
613 ]); bytes.extend_from_slice(&[
615 0xc4, 0x16, 0x49, 0x6e, 0x76, 0x61, 0x6c, 0x69, 0x64, 0x20, 0x4d, 0x65, 0x73, 0x73,
616 0x61, 0x67, 0x65, 0x20, 0x46, 0x6f, 0x72, 0x6d, 0x61, 0x74,
617 ]); let decoded = decode(&bytes).unwrap();
620 let error = decoded.get("error").unwrap();
621 assert_eq!(error, &MsgValue::Bool(true));
622 let error_string = decoded
623 .get("error_string")
624 .and_then(|v| v.as_str())
625 .unwrap();
626 assert_eq!(error_string, "Invalid Message Format");
627 }
628
629 #[test]
630 fn decode_trailing_bytes_fails() {
631 let mut bytes = vec![0xa5, 0x68, 0x65, 0x6c, 0x6c, 0x6f]; bytes.push(0x00); let result = decode(&bytes);
634 assert!(result.is_err());
635 assert!(result.unwrap_err().to_string().contains("trailing bytes"));
636 }
637
638 #[test]
639 fn encode_array_format_auth_login() {
640 let v = MsgValue::Array(vec![
643 MsgValue::String("auth.login".to_string()),
644 MsgValue::String("Mithun".to_string()),
645 MsgValue::String("Mithun@2806".to_string()),
646 ]);
647 let bytes = encode(&v);
648 assert_eq!(bytes[0], 0x93);
650 let decoded = decode(&bytes).unwrap();
651 assert_eq!(v, decoded);
652 }
653}