quinn_proto/
varint.rs

1use std::{convert::TryInto, fmt};
2
3use bytes::{Buf, BufMut};
4use thiserror::Error;
5
6use crate::coding::{self, Codec, UnexpectedEnd};
7
8#[cfg(feature = "arbitrary")]
9use arbitrary::Arbitrary;
10
11/// An integer less than 2^62
12///
13/// Values of this type are suitable for encoding as QUIC variable-length integer.
14// It would be neat if we could express to Rust that the top two bits are available for use as enum
15// discriminants
16#[derive(Default, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
17pub struct VarInt(pub(crate) u64);
18
19impl VarInt {
20    /// The largest representable value
21    pub const MAX: Self = Self((1 << 62) - 1);
22    /// The largest encoded value length
23    pub const MAX_SIZE: usize = 8;
24
25    /// Construct a `VarInt` infallibly
26    pub const fn from_u32(x: u32) -> Self {
27        Self(x as u64)
28    }
29
30    /// Succeeds iff `x` < 2^62
31    pub fn from_u64(x: u64) -> Result<Self, VarIntBoundsExceeded> {
32        if x < 2u64.pow(62) {
33            Ok(Self(x))
34        } else {
35            Err(VarIntBoundsExceeded)
36        }
37    }
38
39    /// Create a VarInt without ensuring it's in range
40    ///
41    /// # Safety
42    ///
43    /// `x` must be less than 2^62.
44    pub const unsafe fn from_u64_unchecked(x: u64) -> Self {
45        Self(x)
46    }
47
48    /// Extract the integer value
49    pub const fn into_inner(self) -> u64 {
50        self.0
51    }
52
53    /// Compute the number of bytes needed to encode this value
54    pub(crate) const fn size(self) -> usize {
55        let x = self.0;
56        if x < 2u64.pow(6) {
57            1
58        } else if x < 2u64.pow(14) {
59            2
60        } else if x < 2u64.pow(30) {
61            4
62        } else if x < 2u64.pow(62) {
63            8
64        } else {
65            panic!("malformed VarInt");
66        }
67    }
68}
69
70impl From<VarInt> for u64 {
71    fn from(x: VarInt) -> Self {
72        x.0
73    }
74}
75
76impl From<u8> for VarInt {
77    fn from(x: u8) -> Self {
78        Self(x.into())
79    }
80}
81
82impl From<u16> for VarInt {
83    fn from(x: u16) -> Self {
84        Self(x.into())
85    }
86}
87
88impl From<u32> for VarInt {
89    fn from(x: u32) -> Self {
90        Self(x.into())
91    }
92}
93
94impl std::convert::TryFrom<u64> for VarInt {
95    type Error = VarIntBoundsExceeded;
96    /// Succeeds iff `x` < 2^62
97    fn try_from(x: u64) -> Result<Self, VarIntBoundsExceeded> {
98        Self::from_u64(x)
99    }
100}
101
102impl std::convert::TryFrom<u128> for VarInt {
103    type Error = VarIntBoundsExceeded;
104    /// Succeeds iff `x` < 2^62
105    fn try_from(x: u128) -> Result<Self, VarIntBoundsExceeded> {
106        Self::from_u64(x.try_into().map_err(|_| VarIntBoundsExceeded)?)
107    }
108}
109
110impl std::convert::TryFrom<usize> for VarInt {
111    type Error = VarIntBoundsExceeded;
112    /// Succeeds iff `x` < 2^62
113    fn try_from(x: usize) -> Result<Self, VarIntBoundsExceeded> {
114        Self::try_from(x as u64)
115    }
116}
117
118impl fmt::Debug for VarInt {
119    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
120        self.0.fmt(f)
121    }
122}
123
124impl fmt::Display for VarInt {
125    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
126        self.0.fmt(f)
127    }
128}
129
130#[cfg(feature = "arbitrary")]
131impl<'arbitrary> Arbitrary<'arbitrary> for VarInt {
132    fn arbitrary(u: &mut arbitrary::Unstructured<'arbitrary>) -> arbitrary::Result<Self> {
133        Ok(Self(u.int_in_range(0..=Self::MAX.0)?))
134    }
135}
136
137/// Error returned when constructing a `VarInt` from a value >= 2^62
138#[derive(Debug, Copy, Clone, Eq, PartialEq, Error)]
139#[error("value too large for varint encoding")]
140pub struct VarIntBoundsExceeded;
141
142impl Codec for VarInt {
143    fn decode<B: Buf>(r: &mut B) -> coding::Result<Self> {
144        if !r.has_remaining() {
145            return Err(UnexpectedEnd);
146        }
147        let mut buf = [0; 8];
148        buf[0] = r.get_u8();
149        let tag = buf[0] >> 6;
150        buf[0] &= 0b0011_1111;
151        let x = match tag {
152            0b00 => u64::from(buf[0]),
153            0b01 => {
154                if r.remaining() < 1 {
155                    return Err(UnexpectedEnd);
156                }
157                r.copy_to_slice(&mut buf[1..2]);
158                u64::from(u16::from_be_bytes(buf[..2].try_into().unwrap()))
159            }
160            0b10 => {
161                if r.remaining() < 3 {
162                    return Err(UnexpectedEnd);
163                }
164                r.copy_to_slice(&mut buf[1..4]);
165                u64::from(u32::from_be_bytes(buf[..4].try_into().unwrap()))
166            }
167            0b11 => {
168                if r.remaining() < 7 {
169                    return Err(UnexpectedEnd);
170                }
171                r.copy_to_slice(&mut buf[1..8]);
172                u64::from_be_bytes(buf)
173            }
174            _ => unreachable!(),
175        };
176        Ok(Self(x))
177    }
178
179    fn encode<B: BufMut>(&self, w: &mut B) {
180        let x = self.0;
181        if x < 2u64.pow(6) {
182            w.put_u8(x as u8);
183        } else if x < 2u64.pow(14) {
184            w.put_u16((0b01 << 14) | x as u16);
185        } else if x < 2u64.pow(30) {
186            w.put_u32((0b10 << 30) | x as u32);
187        } else if x < 2u64.pow(62) {
188            w.put_u64((0b11 << 62) | x);
189        } else {
190            unreachable!("malformed VarInt")
191        }
192    }
193}