ddns/core/parser/
varint.rs

1use std::{cmp::Ordering, convert::TryFrom, fmt};
2
3use bytes::BufMut;
4use nom::{
5    IResult,
6    error::{ErrorKind, make_error},
7    number::streaming::be_u8,
8};
9
10/// An integer less than 2^62
11///
12/// Values of this type are suitable for encoding as QUIC variable-length integer.
13/// It would be neat if we could express to Rust that the top two bits are available for use as enum
14/// discriminants
15///
16/// See [variable-length integers](https://www.rfc-editor.org/rfc/rfc9000.html#name-variable-length-integer-enc)
17/// of [QUIC](https://www.rfc-editor.org/rfc/rfc9000.html) for more details.
18#[derive(Default, Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
19pub struct VarInt(u64);
20
21/// The maximum value that can be represented by a QUIC variable-length integer.
22pub const VARINT_MAX: u64 = 0x3fff_ffff_ffff_ffff;
23
24/// The number of bytes that a QUIC variable-length integer can be encoded in.
25///
26/// [`VarInt`] doesn't need to be encoded on the minimum number of bytes necessary,
27/// with the sole exception of the Frame Type field.
28#[allow(dead_code)]
29pub enum EncodeBytes {
30    One = 1,
31    Two = 2,
32    Four = 4,
33    Eight = 8,
34}
35
36impl VarInt {
37    /// The largest representable value
38    pub const MAX: Self = Self(VARINT_MAX);
39    /// The largest encoded value length
40    pub const MAX_SIZE: usize = 8;
41
42    /// Construct a `VarInt` from a [`u32`].
43    pub const fn from_u32(x: u32) -> Self {
44        Self(x as u64)
45    }
46
47    /// Construct a `VarInt` from a [`u64`].
48    /// Succeeds if `x` < 2^62.
49    pub const fn from_u64(value: u64) -> Result<Self, err::Overflow> {
50        if value <= VARINT_MAX {
51            Ok(Self(value))
52        } else {
53            Err(err::Overflow { value: value as _ })
54        }
55    }
56
57    /// Create a VarInt from a [`u64`] without ensuring it's in range
58    ///
59    /// # Safety
60    ///
61    /// `x` must be less than 2^62.
62    pub unsafe fn from_u64_unchecked(x: u64) -> Self {
63        Self(x)
64    }
65
66    /// Construct a `VarInt` from a [`u128`].
67    /// Succeeds if `x` < 2^62.
68    pub fn from_u128(value: u128) -> Result<Self, err::Overflow> {
69        if value <= VARINT_MAX as u128 {
70            Ok(Self(value as _))
71        } else {
72            Err(err::Overflow { value })
73        }
74    }
75
76    /// Extract the integer value
77    pub const fn into_inner(self) -> u64 {
78        self.0
79    }
80
81    /// Compute the number of bytes needed to encode this value
82    pub fn encoding_size(self) -> usize {
83        let x = self.0;
84        if x < (1 << 6) {
85            1
86        } else if x < (1 << 14) {
87            2
88        } else if x < (1 << 30) {
89            4
90        } else if x < (1 << 62) {
91            8
92        } else {
93            unreachable!("malformed VarInt");
94        }
95    }
96}
97
98impl From<VarInt> for u64 {
99    fn from(x: VarInt) -> Self {
100        x.0
101    }
102}
103
104impl From<u8> for VarInt {
105    fn from(x: u8) -> Self {
106        Self(x.into())
107    }
108}
109
110impl From<u16> for VarInt {
111    fn from(x: u16) -> Self {
112        Self(x.into())
113    }
114}
115
116impl From<u32> for VarInt {
117    fn from(x: u32) -> Self {
118        Self(x.into())
119    }
120}
121
122impl TryFrom<u128> for VarInt {
123    type Error = err::Overflow;
124
125    fn try_from(x: u128) -> Result<Self, Self::Error> {
126        Self::from_u128(x)
127    }
128}
129
130impl TryFrom<u64> for VarInt {
131    type Error = err::Overflow;
132
133    /// Succeeds if `x` < 2^62
134    fn try_from(x: u64) -> Result<Self, Self::Error> {
135        Self::from_u64(x)
136    }
137}
138
139impl TryFrom<usize> for VarInt {
140    type Error = err::Overflow;
141
142    /// Succeeds if `x` < 2^62
143    fn try_from(x: usize) -> Result<Self, Self::Error> {
144        Self::try_from(x as u64)
145    }
146}
147
148impl PartialEq<u64> for VarInt {
149    fn eq(&self, other: &u64) -> bool {
150        self.0.eq(other)
151    }
152}
153
154impl PartialOrd<u64> for VarInt {
155    fn partial_cmp(&self, other: &u64) -> Option<Ordering> {
156        self.0.partial_cmp(other)
157    }
158}
159
160impl fmt::Display for VarInt {
161    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
162        self.0.fmt(f)
163    }
164}
165
166impl fmt::LowerHex for VarInt {
167    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
168        self.0.fmt(f)
169    }
170}
171
172impl fmt::UpperHex for VarInt {
173    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
174        self.0.fmt(f)
175    }
176}
177
178/// Error module for VarInt
179pub mod err {
180    use std::fmt;
181
182    /// Overflow error indicating that a value exceeds 2^62
183    #[derive(Debug, Copy, Clone, Eq, PartialEq)]
184    pub struct Overflow {
185        pub(super) value: u128,
186    }
187
188    impl fmt::Display for Overflow {
189        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
190            write!(f, "Value({}) too large for varint encoding", self.value)
191        }
192    }
193
194    impl std::error::Error for Overflow {}
195}
196
197pub fn be_varint(input: &[u8]) -> IResult<&[u8], VarInt> {
198    let (remain, first_byte) = be_u8(input)?;
199    let len = 2usize.pow((first_byte >> 6) as u32);
200    if remain.len() + 1 < len {
201        return Err(nom::Err::Incomplete(nom::Needed::new(
202            len - (remain.len() + 1),
203        )));
204    }
205
206    let mut buf = [0u8; 8];
207    buf[0] = first_byte & 0b0011_1111;
208    buf[1..len].copy_from_slice(&remain[..len - 1]);
209    let value = u64::from_be_bytes(buf) >> (8 * (8 - len));
210    Ok((&remain[len - 1..], VarInt(value)))
211}
212
213pub trait WriteVarInt {
214    fn put_varint(&mut self, value: VarInt);
215}
216
217impl<T: BufMut> WriteVarInt for T {
218    fn put_varint(&mut self, VarInt(x): VarInt) {
219        if x < 1u64 << 6 {
220            self.put_u8(x as u8);
221        } else if x < 1u64 << 14 {
222            self.put_u16((0b01 << 14) | x as u16);
223        } else if x < 1u64 << 30 {
224            self.put_u32((0b10 << 30) | x as u32);
225        } else if x < 1u64 << 62 {
226            self.put_u64((0b11 << 62) | x);
227        } else {
228            unreachable!("malformed VarInt")
229        };
230    }
231}
232
233#[allow(dead_code)]
234pub fn varint_from_u64(value: u64) -> IResult<&'static [u8], VarInt> {
235    match VarInt::from_u64(value) {
236        Ok(v) => Ok((&[] as &[u8], v)),
237        Err(_e) => Err(nom::Err::Error(make_error(
238            &[] as &[u8],
239            ErrorKind::TooLarge,
240        ))),
241    }
242}
243
244#[cfg(test)]
245mod tests {
246    use bytes::BytesMut;
247
248    use super::*;
249
250    fn roundtrip(v: u64) {
251        let v = VarInt::from_u64(v).unwrap();
252        let mut buf = BytesMut::new();
253        buf.put_varint(v);
254        assert_eq!(buf.len(), v.encoding_size());
255        let (remain, decoded) = be_varint(&buf).unwrap();
256        assert!(remain.is_empty());
257        assert_eq!(decoded, v);
258    }
259
260    #[test]
261    fn quic_varint_roundtrip() {
262        for v in [
263            0u64,
264            1,
265            63,
266            64,
267            16383,
268            16384,
269            (1 << 30) - 1,
270            1 << 30,
271            (1 << 62) - 1,
272        ] {
273            roundtrip(v);
274        }
275    }
276
277    #[test]
278    fn quic_varint_rejects_incomplete() {
279        let v = VarInt::from_u64(64).unwrap();
280        let mut buf = BytesMut::new();
281        buf.put_varint(v);
282        let truncated = &buf[..1];
283        match be_varint(truncated) {
284            Err(nom::Err::Incomplete(_)) => {}
285            other => panic!("expected Incomplete, got {other:?}"),
286        }
287    }
288
289    #[test]
290    fn quic_varint_overflow() {
291        assert!(VarInt::from_u64(VARINT_MAX + 1).is_err());
292    }
293}
ddns/core/parser/varint.rs

ddns/core/parser/
varint.rs
