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 String(String),
12 Array(Vec<MsgValue>),
13 Map(BTreeMap<String, MsgValue>),
14}
15
16impl MsgValue {
17 pub fn get(&self, key: &str) -> Option<&MsgValue> {
18 match self {
19 MsgValue::Map(map) => map.get(key),
20 _ => None,
21 }
22 }
23
24 pub fn as_str(&self) -> Option<&str> {
25 match self {
26 MsgValue::String(value) => Some(value.as_str()),
27 _ => None,
28 }
29 }
30
31 pub fn as_array(&self) -> Option<&[MsgValue]> {
32 match self {
33 MsgValue::Array(values) => Some(values.as_slice()),
34 _ => None,
35 }
36 }
37
38 pub fn as_map(&self) -> Option<&BTreeMap<String, MsgValue>> {
39 match self {
40 MsgValue::Map(map) => Some(map),
41 _ => None,
42 }
43 }
44
45 pub fn as_u64(&self) -> Option<u64> {
46 match self {
47 MsgValue::UInt(value) => Some(*value),
48 MsgValue::Int(value) if *value >= 0 => Some(*value as u64),
49 _ => None,
50 }
51 }
52}
53
54pub fn encode(value: &MsgValue) -> Vec<u8> {
55 let mut buffer = Vec::new();
56 encode_value(value, &mut buffer);
57 buffer
58}
59
60pub fn decode(bytes: &[u8]) -> Result<MsgValue> {
61 let (value, consumed) = decode_value(bytes, 0)?;
62 if consumed != bytes.len() {
63 return Err(AppError::Msgpack(
64 "trailing bytes in msgpack payload".to_string(),
65 ));
66 }
67 Ok(value)
68}
69
70fn encode_value(value: &MsgValue, buffer: &mut Vec<u8>) {
71 match value {
72 MsgValue::Nil => buffer.push(0xc0),
73 MsgValue::Bool(false) => buffer.push(0xc2),
74 MsgValue::Bool(true) => buffer.push(0xc3),
75 MsgValue::Int(value) => encode_int(*value, buffer),
76 MsgValue::UInt(value) => encode_uint(*value, buffer),
77 MsgValue::String(value) => encode_str(value, buffer),
78 MsgValue::Array(values) => encode_array(values, buffer),
79 MsgValue::Map(map) => encode_map(map, buffer),
80 }
81}
82
83fn encode_uint(value: u64, buffer: &mut Vec<u8>) {
84 match value {
85 0..=127 => buffer.push(value as u8),
86 128..=255 => {
87 buffer.push(0xcc);
88 buffer.push(value as u8);
89 }
90 256..=65_535 => {
91 buffer.push(0xcd);
92 buffer.extend_from_slice(&(value as u16).to_be_bytes());
93 }
94 65_536..=4_294_967_295 => {
95 buffer.push(0xce);
96 buffer.extend_from_slice(&(value as u32).to_be_bytes());
97 }
98 _ => {
99 buffer.push(0xcf);
100 buffer.extend_from_slice(&value.to_be_bytes());
101 }
102 }
103}
104
105fn encode_int(value: i64, buffer: &mut Vec<u8>) {
106 if (-32..=127).contains(&value) {
107 if value >= 0 {
108 buffer.push(value as u8);
109 } else {
110 buffer.push((value as i8) as u8);
111 }
112 return;
113 }
114
115 if (-128..=-33).contains(&value) {
116 buffer.push(0xd0);
117 buffer.push(value as i8 as u8);
118 } else if (-32_768..=-129).contains(&value) {
119 buffer.push(0xd1);
120 buffer.extend_from_slice(&(value as i16).to_be_bytes());
121 } else if (-2_147_483_648..=-32_769).contains(&value) {
122 buffer.push(0xd2);
123 buffer.extend_from_slice(&(value as i32).to_be_bytes());
124 } else {
125 buffer.push(0xd3);
126 buffer.extend_from_slice(&value.to_be_bytes());
127 }
128}
129
130fn encode_str(value: &str, buffer: &mut Vec<u8>) {
131 let bytes = value.as_bytes();
132 let len = bytes.len();
133 if len <= 31 {
134 buffer.push(0xa0 | (len as u8));
135 } else if len <= 255 {
136 buffer.push(0xd9);
137 buffer.push(len as u8);
138 } else if len <= 65_535 {
139 buffer.push(0xda);
140 buffer.extend_from_slice(&(len as u16).to_be_bytes());
141 } else {
142 buffer.push(0xdb);
143 buffer.extend_from_slice(&(len as u32).to_be_bytes());
144 }
145 buffer.extend_from_slice(bytes);
146}
147
148fn encode_array(values: &[MsgValue], buffer: &mut Vec<u8>) {
149 let len = values.len();
150 if len <= 15 {
151 buffer.push(0x90 | (len as u8));
152 } else if len <= 65_535 {
153 buffer.push(0xdc);
154 buffer.extend_from_slice(&(len as u16).to_be_bytes());
155 } else {
156 buffer.push(0xdd);
157 buffer.extend_from_slice(&(len as u32).to_be_bytes());
158 }
159 for value in values {
160 encode_value(value, buffer);
161 }
162}
163
164fn encode_map(map: &BTreeMap<String, MsgValue>, buffer: &mut Vec<u8>) {
165 let len = map.len();
166 if len <= 15 {
167 buffer.push(0x80 | (len as u8));
168 } else if len <= 65_535 {
169 buffer.push(0xde);
170 buffer.extend_from_slice(&(len as u16).to_be_bytes());
171 } else {
172 buffer.push(0xdf);
173 buffer.extend_from_slice(&(len as u32).to_be_bytes());
174 }
175 for (key, value) in map {
176 encode_str(key, buffer);
177 encode_value(value, buffer);
178 }
179}
180
181fn decode_value(bytes: &[u8], index: usize) -> Result<(MsgValue, usize)> {
182 let prefix = *bytes
183 .get(index)
184 .ok_or_else(|| AppError::Msgpack("unexpected end of msgpack payload".to_string()))?;
185
186 match prefix {
187 0xc0 => Ok((MsgValue::Nil, index + 1)),
188 0xc2 => Ok((MsgValue::Bool(false), index + 1)),
189 0xc3 => Ok((MsgValue::Bool(true), index + 1)),
190 0xcc => {
191 let value = read_u8(bytes, index + 1)? as u64;
192 Ok((MsgValue::UInt(value), index + 2))
193 }
194 0xcd => {
195 let value = read_u16(bytes, index + 1)? as u64;
196 Ok((MsgValue::UInt(value), index + 3))
197 }
198 0xce => {
199 let value = read_u32(bytes, index + 1)? as u64;
200 Ok((MsgValue::UInt(value), index + 5))
201 }
202 0xcf => {
203 let value = read_u64(bytes, index + 1)?;
204 Ok((MsgValue::UInt(value), index + 9))
205 }
206 0xd0 => {
207 let value = read_i8(bytes, index + 1)? as i64;
208 Ok((MsgValue::Int(value), index + 2))
209 }
210 0xd1 => {
211 let value = read_i16(bytes, index + 1)? as i64;
212 Ok((MsgValue::Int(value), index + 3))
213 }
214 0xd2 => {
215 let value = read_i32(bytes, index + 1)? as i64;
216 Ok((MsgValue::Int(value), index + 5))
217 }
218 0xd3 => {
219 let value = read_i64(bytes, index + 1)?;
220 Ok((MsgValue::Int(value), index + 9))
221 }
222 0xd9 => {
223 let len = read_u8(bytes, index + 1)? as usize;
224 decode_str(bytes, index + 2, len)
225 }
226 0xda => {
227 let len = read_u16(bytes, index + 1)? as usize;
228 decode_str(bytes, index + 3, len)
229 }
230 0xdb => {
231 let len = read_u32(bytes, index + 1)? as usize;
232 decode_str(bytes, index + 5, len)
233 }
234 0xdc => {
235 let len = read_u16(bytes, index + 1)? as usize;
236 decode_array(bytes, index + 3, len)
237 }
238 0xdd => {
239 let len = read_u32(bytes, index + 1)? as usize;
240 decode_array(bytes, index + 5, len)
241 }
242 0xde => {
243 let len = read_u16(bytes, index + 1)? as usize;
244 decode_map(bytes, index + 3, len)
245 }
246 0xdf => {
247 let len = read_u32(bytes, index + 1)? as usize;
248 decode_map(bytes, index + 5, len)
249 }
250 0x80..=0x8f => {
251 let len = (prefix & 0x0f) as usize;
252 decode_map(bytes, index + 1, len)
253 }
254 0x90..=0x9f => {
255 let len = (prefix & 0x0f) as usize;
256 decode_array(bytes, index + 1, len)
257 }
258 0xa0..=0xbf => {
259 let len = (prefix & 0x1f) as usize;
260 decode_str(bytes, index + 1, len)
261 }
262 0xe0..=0xff => Ok((MsgValue::Int((prefix as i8) as i64), index + 1)),
263 0x00..=0x7f => Ok((MsgValue::UInt(prefix as u64), index + 1)),
264 _ => Err(AppError::Msgpack(format!(
265 "unsupported msgpack prefix: 0x{prefix:02x}"
266 ))),
267 }
268}
269
270fn decode_str(bytes: &[u8], index: usize, len: usize) -> Result<(MsgValue, usize)> {
271 let end = index
272 .checked_add(len)
273 .ok_or_else(|| AppError::Msgpack("msgpack string length overflow".to_string()))?;
274 let slice = bytes
275 .get(index..end)
276 .ok_or_else(|| AppError::Msgpack("unexpected end of msgpack string".to_string()))?;
277 let value = String::from_utf8(slice.to_vec())
278 .map_err(|e| AppError::Msgpack(format!("invalid utf-8 string: {e}")))?;
279 Ok((MsgValue::String(value), end))
280}
281
282fn decode_array(bytes: &[u8], mut index: usize, len: usize) -> Result<(MsgValue, usize)> {
283 let mut values = Vec::with_capacity(len);
284 for _ in 0..len {
285 let (value, next) = decode_value(bytes, index)?;
286 values.push(value);
287 index = next;
288 }
289 Ok((MsgValue::Array(values), index))
290}
291
292fn decode_map(bytes: &[u8], mut index: usize, len: usize) -> Result<(MsgValue, usize)> {
293 let mut map = BTreeMap::new();
294 for _ in 0..len {
295 let (key_value, next_key) = decode_value(bytes, index)?;
296 index = next_key;
297 let key = key_value
298 .as_str()
299 .ok_or_else(|| AppError::Msgpack("msgpack map key was not a string".to_string()))?;
300 let (value, next_value) = decode_value(bytes, index)?;
301 index = next_value;
302 map.insert(key.to_string(), value);
303 }
304 Ok((MsgValue::Map(map), index))
305}
306
307fn read_u8(bytes: &[u8], index: usize) -> Result<u8> {
308 bytes
309 .get(index)
310 .copied()
311 .ok_or_else(|| AppError::Msgpack("unexpected end of msgpack payload".to_string()))
312}
313
314fn read_i8(bytes: &[u8], index: usize) -> Result<i8> {
315 Ok(read_u8(bytes, index)? as i8)
316}
317
318fn read_u16(bytes: &[u8], index: usize) -> Result<u16> {
319 let end = index + 2;
320 let slice = bytes
321 .get(index..end)
322 .ok_or_else(|| AppError::Msgpack("unexpected end of msgpack payload".to_string()))?;
323 Ok(u16::from_be_bytes([slice[0], slice[1]]))
324}
325
326fn read_i16(bytes: &[u8], index: usize) -> Result<i16> {
327 Ok(i16::from_be_bytes(read_u16(bytes, index)?.to_be_bytes()))
328}
329
330fn read_u32(bytes: &[u8], index: usize) -> Result<u32> {
331 let end = index + 4;
332 let slice = bytes
333 .get(index..end)
334 .ok_or_else(|| AppError::Msgpack("unexpected end of msgpack payload".to_string()))?;
335 Ok(u32::from_be_bytes([slice[0], slice[1], slice[2], slice[3]]))
336}
337
338fn read_i32(bytes: &[u8], index: usize) -> Result<i32> {
339 Ok(i32::from_be_bytes(read_u32(bytes, index)?.to_be_bytes()))
340}
341
342fn read_u64(bytes: &[u8], index: usize) -> Result<u64> {
343 let end = index + 8;
344 let slice = bytes
345 .get(index..end)
346 .ok_or_else(|| AppError::Msgpack("unexpected end of msgpack payload".to_string()))?;
347 Ok(u64::from_be_bytes([
348 slice[0], slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7],
349 ]))
350}
351
352fn read_i64(bytes: &[u8], index: usize) -> Result<i64> {
353 Ok(i64::from_be_bytes(read_u64(bytes, index)?.to_be_bytes()))
354}