mls_rs_codec/
varint.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// Copyright by contributors to this project.
3// SPDX-License-Identifier: (Apache-2.0 OR MIT)
4
5use crate::{Error, MlsDecode, MlsEncode, MlsSize};
6use alloc::vec::Vec;
7
8#[derive(Debug, Clone, Copy, PartialEq)]
9pub struct VarInt(pub u32);
10
11impl VarInt {
12    pub const MAX: VarInt = VarInt((1 << 30) - 1);
13}
14
15impl From<VarInt> for u32 {
16    #[inline]
17    fn from(n: VarInt) -> u32 {
18        n.0
19    }
20}
21
22impl TryFrom<u32> for VarInt {
23    type Error = Error;
24
25    fn try_from(n: u32) -> Result<Self, Error> {
26        (n <= u32::from(VarInt::MAX))
27            .then_some(VarInt(n))
28            .ok_or(Error::VarIntOutOfRange)
29    }
30}
31
32impl TryFrom<usize> for VarInt {
33    type Error = Error;
34
35    fn try_from(n: usize) -> Result<Self, Error> {
36        u32::try_from(n)
37            .map_err(|_| Error::VarIntOutOfRange)?
38            .try_into()
39    }
40}
41
42impl MlsSize for VarInt {
43    #[inline]
44    fn mls_encoded_len(&self) -> usize {
45        count_bytes_to_encode_int(*self) as usize
46    }
47}
48
49impl MlsEncode for VarInt {
50    fn mls_encode(&self, writer: &mut Vec<u8>) -> Result<(), Error> {
51        let mut bytes = self.0.to_be_bytes();
52
53        let bytes = match count_bytes_to_encode_int(*self) {
54            LengthEncoding::One => &bytes[3..],
55            LengthEncoding::Two => {
56                bytes[2] |= 0x40;
57                &bytes[2..]
58            }
59            LengthEncoding::Four => {
60                bytes[0] |= 0x80;
61                &bytes
62            }
63        };
64
65        writer.extend_from_slice(bytes);
66        Ok(())
67    }
68}
69
70impl MlsDecode for VarInt {
71    fn mls_decode(reader: &mut &[u8]) -> Result<Self, Error> {
72        let first = u8::mls_decode(reader)?;
73
74        let prefix = first >> 6;
75
76        let count = (prefix < 3)
77            .then_some(1 << prefix)
78            .ok_or(Error::InvalidVarIntPrefix(prefix))?;
79
80        let n = (1..count).try_fold(u32::from(first & 0x3f), |n, _| {
81            u8::mls_decode(reader).map(|b| (n << 8) | u32::from(b))
82        })?;
83
84        let n = VarInt(n);
85
86        if count_bytes_to_encode_int(n) as usize == count {
87            Ok(n)
88        } else {
89            Err(Error::VarIntMinimumLengthEncoding)
90        }
91    }
92}
93
94/// Number of bytes to encode a variable-size integer.
95#[derive(Debug)]
96enum LengthEncoding {
97    One = 1,
98    Two = 2,
99    Four = 4,
100}
101
102fn count_bytes_to_encode_int(n: VarInt) -> LengthEncoding {
103    let used_bits = 32 - n.0.leading_zeros();
104    match used_bits {
105        0..=6 => LengthEncoding::One,
106        7..=14 => LengthEncoding::Two,
107        15..=30 => LengthEncoding::Four,
108        _ => panic!("Such a large VarInt cannot be instantiated"),
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::VarInt;
115    use crate::{Error, MlsDecode, MlsEncode};
116    use assert_matches::assert_matches;
117
118    #[cfg(target_arch = "wasm32")]
119    use wasm_bindgen_test::wasm_bindgen_test as test;
120
121    #[test]
122    fn zero_is_convertible_to_varint() {
123        assert_matches!(VarInt::try_from(0u32).map(u32::from), Ok(0));
124    }
125
126    #[test]
127    fn successor_of_max_varint_is_not_convertible_to_varint() {
128        let n = u32::from(VarInt::MAX) + 1;
129        assert_matches!(VarInt::try_from(n), Err(Error::VarIntOutOfRange));
130    }
131
132    #[test]
133    fn zero_serializes_as_single_null_byte() {
134        assert_eq!(
135            VarInt::try_from(0u32).unwrap().mls_encode_to_vec().unwrap(),
136            [0]
137        );
138    }
139
140    #[test]
141    fn zero_roundtrips() {
142        let n = VarInt::try_from(0u32).unwrap();
143
144        let serialized = n.mls_encode_to_vec().unwrap();
145        let restored = VarInt::mls_decode(&mut &*serialized).unwrap();
146
147        assert_eq!(restored, n);
148    }
149
150    #[test]
151    fn varint_max_roundtrips() {
152        let n = VarInt::MAX;
153
154        let serialized = n.mls_encode_to_vec().unwrap();
155        let restored = VarInt::mls_decode(&mut &*serialized).unwrap();
156
157        assert_eq!(restored, n);
158    }
159
160    fn decoding_matches_rfc(encoded: u32, decoded: u32) {
161        let bytes = encoded.to_be_bytes();
162
163        let start = bytes
164            .iter()
165            .position(|&b| b != 0)
166            .unwrap_or(bytes.len() - 1);
167
168        let bytes = &bytes[start..];
169
170        assert_eq!(
171            VarInt::mls_decode(&mut &*bytes).unwrap(),
172            VarInt::try_from(decoded).unwrap()
173        );
174    }
175
176    #[test]
177    fn decoding_0x25_matches_rfc_result() {
178        decoding_matches_rfc(0x25, 37);
179    }
180
181    #[test]
182    fn decoding_0x7bbd_matches_rfc_result() {
183        decoding_matches_rfc(0x7bbd, 15293);
184    }
185
186    #[test]
187    fn decoding_0x9d7f3e7d_matches_rfc_result() {
188        decoding_matches_rfc(0x9d7f3e7d, 494878333);
189    }
190}
mls_rs_codec/varint.rs

mls_rs_codec/
varint.rs
