1use super::values::{Constants, Value};
18use alloc::vec::Vec;
19
20#[derive(Debug, PartialEq)]
22pub enum EncoderError {
23    TooMuchNesting,
24    DuplicateMapKey,
25}
26
27pub fn write(value: Value, encoded_cbor: &mut Vec<u8>) -> Result<(), EncoderError> {
30    write_nested(value, encoded_cbor, None)
31}
32
33pub fn write_nested(
37    value: Value,
38    encoded_cbor: &mut Vec<u8>,
39    max_nest: Option<i8>,
40) -> Result<(), EncoderError> {
41    let mut writer = Writer::new(encoded_cbor);
42    writer.encode_cbor(value, max_nest)
43}
44
45struct Writer<'a> {
46    encoded_cbor: &'a mut Vec<u8>,
47}
48
49impl<'a> Writer<'a> {
50    pub fn new(encoded_cbor: &mut Vec<u8>) -> Writer {
51        Writer { encoded_cbor }
52    }
53
54    fn encode_cbor(
55        &mut self,
56        value: Value,
57        remaining_depth: Option<i8>,
58    ) -> Result<(), EncoderError> {
59        if remaining_depth.map_or(false, |d| d < 0) {
60            return Err(EncoderError::TooMuchNesting);
61        }
62        let type_label = value.type_label();
63        match value {
64            Value::Unsigned(unsigned) => self.start_item(type_label, unsigned),
65            Value::Negative(negative) => self.start_item(type_label, -(negative + 1) as u64),
66            Value::ByteString(byte_string) => {
67                self.start_item(type_label, byte_string.len() as u64);
68                self.encoded_cbor.extend(byte_string);
69            }
70            Value::TextString(text_string) => {
71                self.start_item(type_label, text_string.len() as u64);
72                self.encoded_cbor.extend(text_string.into_bytes());
73            }
74            Value::Array(array) => {
75                self.start_item(type_label, array.len() as u64);
76                for el in array {
77                    self.encode_cbor(el, remaining_depth.map(|d| d - 1))?;
78                }
79            }
80            Value::Map(mut map) => {
81                map.sort_by(|a, b| a.0.cmp(&b.0));
82                let map_len = map.len();
83                map.dedup_by(|a, b| a.0.eq(&b.0));
84                if map_len != map.len() {
85                    return Err(EncoderError::DuplicateMapKey);
86                }
87                self.start_item(type_label, map_len as u64);
88                for (k, v) in map {
89                    self.encode_cbor(k, remaining_depth.map(|d| d - 1))?;
90                    self.encode_cbor(v, remaining_depth.map(|d| d - 1))?;
91                }
92            }
93            Value::Tag(tag, inner_value) => {
94                self.start_item(type_label, tag);
95                self.encode_cbor(*inner_value, remaining_depth.map(|d| d - 1))?;
96            }
97            Value::Simple(simple_value) => self.start_item(type_label, simple_value as u64),
98        }
99        Ok(())
100    }
101
102    fn start_item(&mut self, type_label: u8, size: u64) {
103        let (mut first_byte, shift) = match size {
104            0..=23 => (size as u8, 0),
105            24..=0xFF => (Constants::ADDITIONAL_INFORMATION_1_BYTE, 1),
106            0x100..=0xFFFF => (Constants::ADDITIONAL_INFORMATION_2_BYTES, 2),
107            0x10000..=0xFFFF_FFFF => (Constants::ADDITIONAL_INFORMATION_4_BYTES, 4),
108            _ => (Constants::ADDITIONAL_INFORMATION_8_BYTES, 8),
109        };
110        first_byte |= type_label << Constants::MAJOR_TYPE_BIT_SHIFT;
111        self.encoded_cbor.push(first_byte);
112
113        for i in (0..shift).rev() {
114            self.encoded_cbor.push((size >> (i * 8)) as u8);
115        }
116    }
117}
118
119#[cfg(test)]
120mod test {
121    use super::*;
122    use crate::{
123        cbor_array, cbor_array_vec, cbor_bytes, cbor_false, cbor_int, cbor_map, cbor_null,
124        cbor_tagged, cbor_text, cbor_true, cbor_undefined,
125    };
126    use alloc::vec;
127
128    fn write_return(value: Value) -> Option<Vec<u8>> {
129        let mut encoded_cbor = Vec::new();
130        if write(value, &mut encoded_cbor).is_ok() {
131            Some(encoded_cbor)
132        } else {
133            None
134        }
135    }
136
137    #[test]
138    fn test_write_unsigned() {
139        let cases = vec![
140            (0, vec![0x00]),
141            (1, vec![0x01]),
142            (10, vec![0x0A]),
143            (23, vec![0x17]),
144            (24, vec![0x18, 0x18]),
145            (25, vec![0x18, 0x19]),
146            (100, vec![0x18, 0x64]),
147            (1000, vec![0x19, 0x03, 0xE8]),
148            (1000000, vec![0x1A, 0x00, 0x0F, 0x42, 0x40]),
149            (0xFFFFFFFF, vec![0x1A, 0xFF, 0xFF, 0xFF, 0xFF]),
150            (
151                0x100000000,
152                vec![0x1B, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00],
153            ),
154            (
155                core::i64::MAX,
156                vec![0x1B, 0x7F, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF],
157            ),
158        ];
159        for (unsigned, correct_cbor) in cases {
160            assert_eq!(write_return(cbor_int!(unsigned)), Some(correct_cbor));
161        }
162    }
163
164    #[test]
165    fn test_write_negative() {
166        let cases = vec![
167            (-1, vec![0x20]),
168            (-10, vec![0x29]),
169            (-23, vec![0x36]),
170            (-24, vec![0x37]),
171            (-25, vec![0x38, 0x18]),
172            (-100, vec![0x38, 0x63]),
173            (-1000, vec![0x39, 0x03, 0xE7]),
174            (-4294967296, vec![0x3A, 0xFF, 0xFF, 0xFF, 0xFF]),
175            (
176                -4294967297,
177                vec![0x3B, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00],
178            ),
179            (
180                core::i64::MIN,
181                vec![0x3B, 0x7F, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF],
182            ),
183        ];
184        for (negative, correct_cbor) in cases {
185            assert_eq!(write_return(cbor_int!(negative)), Some(correct_cbor));
186        }
187    }
188
189    #[test]
190    fn test_write_byte_string() {
191        let cases = vec![
192            (vec![], vec![0x40]),
193            (
194                vec![0x01, 0x02, 0x03, 0x04],
195                vec![0x44, 0x01, 0x02, 0x03, 0x04],
196            ),
197        ];
198        for (byte_string, correct_cbor) in cases {
199            assert_eq!(write_return(cbor_bytes!(byte_string)), Some(correct_cbor));
200        }
201    }
202
203    #[test]
204    fn test_write_text_string() {
205        let unicode_3byte = vec![0xE6, 0xB0, 0xB4];
206        let cases = vec![
207            ("", vec![0x60]),
208            ("a", vec![0x61, 0x61]),
209            ("IETF", vec![0x64, 0x49, 0x45, 0x54, 0x46]),
210            ("\"\\", vec![0x62, 0x22, 0x5C]),
211            ("ü", vec![0x62, 0xC3, 0xBC]),
212            (
213                core::str::from_utf8(&unicode_3byte).unwrap(),
214                vec![0x63, 0xE6, 0xB0, 0xB4],
215            ),
216            ("𐅑", vec![0x64, 0xF0, 0x90, 0x85, 0x91]),
217        ];
218        for (text_string, correct_cbor) in cases {
219            assert_eq!(write_return(cbor_text!(text_string)), Some(correct_cbor));
220        }
221    }
222
223    #[test]
224    fn test_write_array() {
225        let value_vec: Vec<_> = (1..26).collect();
226        let expected_cbor = vec![
227            0x98, 0x19, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
229            0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x18, 0x18, 0x19,
230        ];
231        assert_eq!(
232            write_return(cbor_array_vec!(value_vec)),
233            Some(expected_cbor)
234        );
235    }
236
237    #[test]
238    fn test_write_map() {
239        let value_map = cbor_map! {
240            0 => "a",
241            23 => "b",
242            24 => "c",
243            core::u8::MAX as i64 => "d",
244            256 => "e",
245            core::u16::MAX as i64 => "f",
246            65536 => "g",
247            core::u32::MAX as i64 => "h",
248            4294967296_i64 => "i",
249            core::i64::MAX => "j",
250            -1 => "k",
251            -24 => "l",
252            -25 => "m",
253            -256 => "n",
254            -257 => "o",
255            -65537 => "p",
256            -4294967296_i64 => "q",
257            -4294967297_i64 => "r",
258            core::i64::MIN => "s",
259            b"a" => 2,
260            b"bar" => 3,
261            b"foo" => 4,
262            "" => ".",
263            "e" => "E",
264            "aa" => "AA",
265        };
266        let expected_cbor = vec![
267            0xb8, 0x19, 0x00, 0x61, 0x61, 0x17, 0x61, 0x62, 0x18, 0x18, 0x61, 0x63, 0x18, 0xFF, 0x61, 0x64, 0x19, 0x01, 0x00, 0x61, 0x65, 0x19, 0xFF, 0xFF, 0x61, 0x66, 0x1A, 0x00, 0x01, 0x00, 0x00, 0x61, 0x67, 0x1A, 0xFF, 0xFF, 0xFF, 0xFF, 0x61, 0x68, 0x1B, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x61, 0x69, 0x1b, 0x7f, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x61, 0x6a, 0x20, 0x61, 0x6b, 0x37, 0x61, 0x6c, 0x38, 0x18, 0x61, 0x6d, 0x38, 0xFF, 0x61, 0x6e, 0x39, 0x01, 0x00, 0x61, 0x6f, 0x3A, 0x00, 0x01, 0x00, 0x00, 0x61, 0x70, 0x3A, 0xFF, 0xFF, 0xFF, 0xFF, 0x61, 0x71, 0x3B, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x61, 0x72, 0x3b, 0x7f, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x61, 0x73, 0x41, b'a', 0x02, 0x43, b'b', b'a', b'r', 0x03, 0x43, b'f', b'o', b'o', 0x04, 0x60, 0x61, 0x2e, 0x61, 0x65, 0x61, 0x45, 0x62, 0x61, 0x61, 0x62, 0x41, 0x41, ];
316        assert_eq!(write_return(value_map), Some(expected_cbor));
317    }
318
319    #[test]
320    fn test_write_map_sorted() {
321        let sorted_map = cbor_map! {
322            0 => "a",
323            1 => "b",
324            -1 => "c",
325            -2 => "d",
326            b"a" => "e",
327            b"b" => "f",
328            "" => "g",
329            "c" => "h",
330        };
331        let unsorted_map = cbor_map! {
332            1 => "b",
333            -2 => "d",
334            b"b" => "f",
335            "c" => "h",
336            "" => "g",
337            b"a" => "e",
338            -1 => "c",
339            0 => "a",
340        };
341        assert_eq!(write_return(sorted_map), write_return(unsorted_map));
342    }
343
344    #[test]
345    fn test_write_map_duplicates() {
346        let duplicate0 = cbor_map! {
347            0 => "a",
348            -1 => "c",
349            b"a" => "e",
350            "c" => "g",
351            0 => "b",
352        };
353        assert_eq!(write_return(duplicate0), None);
354        let duplicate1 = cbor_map! {
355            0 => "a",
356            -1 => "c",
357            b"a" => "e",
358            "c" => "g",
359            -1 => "d",
360        };
361        assert_eq!(write_return(duplicate1), None);
362        let duplicate2 = cbor_map! {
363            0 => "a",
364            -1 => "c",
365            b"a" => "e",
366            "c" => "g",
367            b"a" => "f",
368        };
369        assert_eq!(write_return(duplicate2), None);
370        let duplicate3 = cbor_map! {
371            0 => "a",
372            -1 => "c",
373            b"a" => "e",
374            "c" => "g",
375            "c" => "h",
376        };
377        assert_eq!(write_return(duplicate3), None);
378    }
379
380    #[test]
381    fn test_write_map_with_array() {
382        let value_map = cbor_map! {
383            "a" => 1,
384            "b" => cbor_array![2, 3],
385        };
386        let expected_cbor = vec![
387            0xa2, 0x61, 0x61, 0x01, 0x61, 0x62, 0x82, 0x02, 0x03,
392        ];
393        assert_eq!(write_return(value_map), Some(expected_cbor));
394    }
395
396    #[test]
397    fn test_write_nested_map() {
398        let value_map = cbor_map! {
399            "a" => 1,
400            "b" => cbor_map! {
401                "c" => 2,
402                "d" => 3,
403            },
404        };
405        let expected_cbor = vec![
406            0xa2, 0x61, 0x61, 0x01, 0x61, 0x62, 0xa2, 0x61, 0x63, 0x02, 0x61, 0x64, 0x03,
413        ];
414        assert_eq!(write_return(value_map), Some(expected_cbor));
415    }
416
417    #[test]
418    fn test_write_tagged() {
419        let cases = vec![
420            (cbor_tagged!(6, cbor_int!(0x42)), vec![0xc6, 0x18, 0x42]),
421            (cbor_tagged!(1, cbor_true!()), vec![0xc1, 0xf5]),
422            (
423                cbor_tagged!(
424                    1000,
425                    cbor_map! {
426                        "a" => 1,
427                        "b" => cbor_array![2, 3],
428                    }
429                ),
430                vec![
431                    0xd9, 0x03, 0xe8, 0xa2, 0x61, 0x61, 0x01, 0x61, 0x62, 0x82, 0x02, 0x03,
436                ],
437            ),
438        ];
439        for (value, correct_cbor) in cases {
440            assert_eq!(write_return(value), Some(correct_cbor));
441        }
442    }
443
444    #[test]
445    fn test_write_simple() {
446        let cases = vec![
447            (cbor_false!(), vec![0xF4]),
448            (cbor_true!(), vec![0xF5]),
449            (cbor_null!(), vec![0xF6]),
450            (cbor_undefined!(), vec![0xF7]),
451        ];
452        for (value, correct_cbor) in cases {
453            assert_eq!(write_return(value), Some(correct_cbor));
454        }
455    }
456
457    #[test]
458    fn test_write_single_levels() {
459        let simple_array: Value = cbor_array![2];
460        let simple_map: Value = cbor_map! {"b" => 3};
461        let positive_cases = vec![
462            (cbor_int!(1), 0),
463            (cbor_bytes!(vec![0x01, 0x02, 0x03, 0x04]), 0),
464            (cbor_text!("a"), 0),
465            (cbor_array![], 0),
466            (cbor_map! {}, 0),
467            (simple_array.clone(), 1),
468            (simple_map.clone(), 1),
469        ];
470        let negative_cases = vec![(simple_array.clone(), 0), (simple_map.clone(), 0)];
471        for (value, level) in positive_cases {
472            let mut buf = Vec::new();
473            let mut writer = Writer::new(&mut buf);
474            assert!(writer.encode_cbor(value, Some(level)).is_ok());
475        }
476        for (value, level) in negative_cases {
477            let mut buf = Vec::new();
478            let mut writer = Writer::new(&mut buf);
479            assert!(!writer.encode_cbor(value, Some(level)).is_ok());
480        }
481    }
482
483    #[test]
484    fn test_write_nested_map_levels() {
485        let cbor_map: Value = cbor_map! {
486            "a" => 1,
487            "b" => cbor_map! {
488                "c" => 2,
489                "d" => 3,
490            },
491        };
492
493        let mut buf = Vec::new();
494        let mut writer = Writer::new(&mut buf);
495        assert!(writer.encode_cbor(cbor_map.clone(), Some(2)).is_ok());
496        assert!(writer.encode_cbor(cbor_map.clone(), None).is_ok());
497        writer = Writer::new(&mut buf);
498        assert!(writer.encode_cbor(cbor_map, Some(1)).is_err());
499    }
500
501    #[test]
502    fn test_write_unbalanced_nested_containers() {
503        let cbor_array: Value = cbor_array![
504            1,
505            2,
506            3,
507            cbor_map! {
508                "a" => 1,
509                "b" => cbor_map! {
510                    "c" => 2,
511                    "d" => 3,
512                },
513            },
514        ];
515
516        let mut buf = Vec::new();
517        let mut writer = Writer::new(&mut buf);
518        assert!(writer.encode_cbor(cbor_array.clone(), Some(3)).is_ok());
519        writer = Writer::new(&mut buf);
520        assert!(writer.encode_cbor(cbor_array, Some(2)).is_err());
521    }
522
523    #[test]
524    fn test_write_overly_nested() {
525        let cbor_map: Value = cbor_map! {
526            "a" => 1,
527            "b" => cbor_map! {
528                "c" => 2,
529                "d" => 3,
530                "h" => cbor_map! {
531                    "e" => 4,
532                    "f" => 5,
533                    "g" => cbor_array![
534                        6,
535                        7,
536                        cbor_array![
537                            8
538                        ]
539                    ],
540                },
541            },
542        };
543
544        let mut buf = Vec::new();
545        let mut writer = Writer::new(&mut buf);
546        assert!(writer.encode_cbor(cbor_map.clone(), Some(5)).is_ok());
547        assert!(writer.encode_cbor(cbor_map.clone(), None).is_ok());
548        writer = Writer::new(&mut buf);
549        assert!(writer.encode_cbor(cbor_map, Some(4)).is_err());
550    }
551}