Skip to main content

po_wire/
varint.rs

1//! QUIC-style Variable-Length Integer encoding.
2//!
3//! Uses the RFC 9000 §16 format: the two most significant bits of the first byte
4//! encode the length of the integer (1, 2, 4, or 8 bytes), leaving 6, 14, 30, or
5//! 62 bits for the value respectively.
6//!
7//! | Prefix (2 bits) | Length | Usable Bits | Max Value              |
8//! |-----------------|--------|-------------|------------------------|
9//! | 00              | 1 byte | 6           | 63                     |
10//! | 01              | 2 bytes| 14          | 16,383                 |
11//! | 10              | 4 bytes| 30          | 1,073,741,823          |
12//! | 11              | 8 bytes| 62          | 4,611,686,018,427,387,903 |
13
14use crate::error::WireError;
15
16/// Maximum value encodable in a QUIC VarInt (2^62 - 1).
17pub const VARINT_MAX: u64 = (1 << 62) - 1;
18
19/// Encode a `u64` value as a QUIC-style VarInt into `buf`.
20///
21/// Returns the number of bytes written.
22///
23/// # Errors
24/// - `WireError::BufferTooSmall` if `buf` is not large enough.
25/// - `WireError::InvalidVarInt` if `value` exceeds `VARINT_MAX`.
26#[inline]
27pub fn encode(value: u64, buf: &mut [u8]) -> Result<usize, WireError> {
28    if value > VARINT_MAX {
29        return Err(WireError::InvalidVarInt);
30    }
31
32    if value <= 63 {
33        // 1 byte: prefix 00
34        if buf.is_empty() {
35            return Err(WireError::BufferTooSmall { needed: 1, available: 0 });
36        }
37        buf[0] = value as u8;
38        Ok(1)
39    } else if value <= 16_383 {
40        // 2 bytes: prefix 01
41        if buf.len() < 2 {
42            return Err(WireError::BufferTooSmall { needed: 2, available: buf.len() });
43        }
44        let v = (value as u16) | 0x4000;
45        buf[..2].copy_from_slice(&v.to_be_bytes());
46        Ok(2)
47    } else if value <= 1_073_741_823 {
48        // 4 bytes: prefix 10
49        if buf.len() < 4 {
50            return Err(WireError::BufferTooSmall { needed: 4, available: buf.len() });
51        }
52        let v = (value as u32) | 0x8000_0000;
53        buf[..4].copy_from_slice(&v.to_be_bytes());
54        Ok(4)
55    } else {
56        // 8 bytes: prefix 11
57        if buf.len() < 8 {
58            return Err(WireError::BufferTooSmall { needed: 8, available: buf.len() });
59        }
60        let v = value | 0xC000_0000_0000_0000;
61        buf[..8].copy_from_slice(&v.to_be_bytes());
62        Ok(8)
63    }
64}
65
66/// Decode a QUIC-style VarInt from `buf`.
67///
68/// Returns `(value, bytes_consumed)`.
69///
70/// # Errors
71/// - `WireError::Incomplete` if `buf` is too short.
72#[inline]
73pub fn decode(buf: &[u8]) -> Result<(u64, usize), WireError> {
74    if buf.is_empty() {
75        return Err(WireError::Incomplete { needed_min: 1, available: 0 });
76    }
77
78    let prefix = buf[0] >> 6;
79    let len = 1usize << prefix; // 1, 2, 4, or 8
80
81    if buf.len() < len {
82        return Err(WireError::Incomplete { needed_min: len, available: buf.len() });
83    }
84
85    let value = match len {
86        1 => (buf[0] & 0x3F) as u64,
87        2 => {
88            let mut bytes = [0u8; 2];
89            bytes.copy_from_slice(&buf[..2]);
90            bytes[0] &= 0x3F;
91            u16::from_be_bytes(bytes) as u64
92        }
93        4 => {
94            let mut bytes = [0u8; 4];
95            bytes.copy_from_slice(&buf[..4]);
96            bytes[0] &= 0x3F;
97            u32::from_be_bytes(bytes) as u64
98        }
99        8 => {
100            let mut bytes = [0u8; 8];
101            bytes.copy_from_slice(&buf[..8]);
102            bytes[0] &= 0x3F;
103            u64::from_be_bytes(bytes)
104        }
105        _ => unreachable!(),
106    };
107
108    Ok((value, len))
109}
110
111/// Returns the number of bytes needed to encode `value` as a VarInt.
112#[inline]
113pub const fn encoded_len(value: u64) -> usize {
114    if value <= 63 {
115        1
116    } else if value <= 16_383 {
117        2
118    } else if value <= 1_073_741_823 {
119        4
120    } else {
121        8
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn roundtrip_1byte() {
131        let mut buf = [0u8; 8];
132        for v in [0, 1, 42, 63] {
133            let n = encode(v, &mut buf).unwrap();
134            assert_eq!(n, 1, "value {v} should encode in 1 byte");
135            let (decoded, consumed) = decode(&buf[..n]).unwrap();
136            assert_eq!(decoded, v);
137            assert_eq!(consumed, 1);
138        }
139    }
140
141    #[test]
142    fn roundtrip_2bytes() {
143        let mut buf = [0u8; 8];
144        for v in [64, 100, 15293, 16383] {
145            let n = encode(v, &mut buf).unwrap();
146            assert_eq!(n, 2, "value {v} should encode in 2 bytes");
147            let (decoded, consumed) = decode(&buf[..n]).unwrap();
148            assert_eq!(decoded, v);
149            assert_eq!(consumed, 2);
150        }
151    }
152
153    #[test]
154    fn roundtrip_4bytes() {
155        let mut buf = [0u8; 8];
156        for v in [16384, 100_000, 1_073_741_823] {
157            let n = encode(v, &mut buf).unwrap();
158            assert_eq!(n, 4, "value {v} should encode in 4 bytes");
159            let (decoded, consumed) = decode(&buf[..n]).unwrap();
160            assert_eq!(decoded, v);
161            assert_eq!(consumed, 4);
162        }
163    }
164
165    #[test]
166    fn roundtrip_8bytes() {
167        let mut buf = [0u8; 8];
168        for v in [1_073_741_824, u64::MAX >> 2, VARINT_MAX] {
169            let n = encode(v, &mut buf).unwrap();
170            assert_eq!(n, 8, "value {v} should encode in 8 bytes");
171            let (decoded, consumed) = decode(&buf[..n]).unwrap();
172            assert_eq!(decoded, v);
173            assert_eq!(consumed, 8);
174        }
175    }
176
177    #[test]
178    fn encoded_len_matches() {
179        assert_eq!(encoded_len(0), 1);
180        assert_eq!(encoded_len(63), 1);
181        assert_eq!(encoded_len(64), 2);
182        assert_eq!(encoded_len(16383), 2);
183        assert_eq!(encoded_len(16384), 4);
184        assert_eq!(encoded_len(1_073_741_823), 4);
185        assert_eq!(encoded_len(1_073_741_824), 8);
186    }
187
188    #[test]
189    fn overflow_rejected() {
190        let mut buf = [0u8; 8];
191        assert_eq!(encode(VARINT_MAX + 1, &mut buf), Err(WireError::InvalidVarInt));
192    }
193
194    #[test]
195    fn buffer_too_small() {
196        let mut buf = [0u8; 1];
197        assert!(matches!(
198            encode(15293, &mut buf),
199            Err(WireError::BufferTooSmall { needed: 2, .. })
200        ));
201    }
202
203    #[test]
204    fn incomplete_decode() {
205        let buf = [0x40]; // prefix 01 means 2 bytes, but only 1 available
206        assert!(matches!(
207            decode(&buf),
208            Err(WireError::Incomplete { needed_min: 2, available: 1 })
209        ));
210    }
211
212    #[test]
213    fn empty_decode() {
214        assert!(matches!(
215            decode(&[]),
216            Err(WireError::Incomplete { needed_min: 1, available: 0 })
217        ));
218    }
219
220    #[test]
221    fn boundary_values() {
222        let mut buf = [0u8; 8];
223        // Test exact boundary transitions
224        let boundaries = [63, 64, 16383, 16384, 1_073_741_823, 1_073_741_824];
225        for v in boundaries {
226            let n = encode(v, &mut buf).unwrap();
227            let (decoded, _) = decode(&buf[..n]).unwrap();
228            assert_eq!(decoded, v, "boundary value {v} failed roundtrip");
229        }
230    }
231}