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
5 use crate::{Error, MlsDecode, MlsEncode, MlsSize};
6 use alloc::vec::Vec;
7
8 #[derive(Debug, Clone, Copy, PartialEq)]
9 pub struct VarInt(pub u32);
10
11 impl VarInt {
12 pub const MAX: VarInt = VarInt((1 << 30) - 1);
13 }
14
15 impl From<VarInt> for u32 {
16 #[inline]
from(n: VarInt) -> u3217 fn from(n: VarInt) -> u32 {
18 n.0
19 }
20 }
21
22 impl TryFrom<u32> for VarInt {
23 type Error = Error;
24
try_from(n: u32) -> Result<Self, Error>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
32 impl TryFrom<usize> for VarInt {
33 type Error = Error;
34
try_from(n: usize) -> Result<Self, Error>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
42 impl MlsSize for VarInt {
43 #[inline]
mls_encoded_len(&self) -> usize44 fn mls_encoded_len(&self) -> usize {
45 count_bytes_to_encode_int(*self) as usize
46 }
47 }
48
49 impl MlsEncode for VarInt {
mls_encode(&self, writer: &mut Vec<u8>) -> Result<(), Error>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
70 impl MlsDecode for VarInt {
mls_decode(reader: &mut &[u8]) -> Result<Self, Error>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)]
96 enum LengthEncoding {
97 One = 1,
98 Two = 2,
99 Four = 4,
100 }
101
count_bytes_to_encode_int(n: VarInt) -> LengthEncoding102 fn 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)]
113 mod 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]
zero_is_convertible_to_varint()122 fn zero_is_convertible_to_varint() {
123 assert_matches!(VarInt::try_from(0u32).map(u32::from), Ok(0));
124 }
125
126 #[test]
successor_of_max_varint_is_not_convertible_to_varint()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]
zero_serializes_as_single_null_byte()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]
zero_roundtrips()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]
varint_max_roundtrips()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
decoding_matches_rfc(encoded: u32, decoded: u32)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]
decoding_0x25_matches_rfc_result()177 fn decoding_0x25_matches_rfc_result() {
178 decoding_matches_rfc(0x25, 37);
179 }
180
181 #[test]
decoding_0x7bbd_matches_rfc_result()182 fn decoding_0x7bbd_matches_rfc_result() {
183 decoding_matches_rfc(0x7bbd, 15293);
184 }
185
186 #[test]
decoding_0x9d7f3e7d_matches_rfc_result()187 fn decoding_0x9d7f3e7d_matches_rfc_result() {
188 decoding_matches_rfc(0x9d7f3e7d, 494878333);
189 }
190 }
191