Skip to main content

simoom_quic/
coding.rs

1//! Variable-length integers (RFC 9000 ยง16) and a bounds-checked reader.
2
3/// Largest value a varint can hold (2^62 - 1).
4pub const VARINT_MAX: u64 = (1 << 62) - 1;
5
6/// Encoded size of `v` in bytes (1, 2, 4 or 8). `v` must be at most [`VARINT_MAX`].
7#[inline]
8pub fn varint_len(v: u64) -> usize {
9    if v < 1 << 6 {
10        1
11    } else if v < 1 << 14 {
12        2
13    } else if v < 1 << 30 {
14        4
15    } else {
16        8
17    }
18}
19
20/// Appends `v` as a varint. Panics in debug builds when `v` exceeds [`VARINT_MAX`].
21#[inline]
22pub fn put_varint(out: &mut Vec<u8>, v: u64) {
23    debug_assert!(v <= VARINT_MAX);
24    match varint_len(v) {
25        1 => out.push(v as u8),
26        2 => out.extend_from_slice(&(v as u16 | 0x4000).to_be_bytes()),
27        4 => out.extend_from_slice(&(v as u32 | 0x8000_0000).to_be_bytes()),
28        _ => out.extend_from_slice(&(v | 0xc000_0000_0000_0000).to_be_bytes()),
29    }
30}
31
32/// Writes `v` as a varint of exactly `len` bytes (used for patched length fields).
33pub fn put_varint_fixed(out: &mut [u8], v: u64, len: usize) {
34    let b = v.to_be_bytes();
35    out[..len].copy_from_slice(&b[8 - len..]);
36    out[0] |= match len {
37        1 => 0,
38        2 => 0x40,
39        4 => 0x80,
40        _ => 0xc0,
41    };
42}
43
44/// A cursor over a byte slice; every read is bounds-checked.
45#[derive(Clone, Copy)]
46pub struct Reader<'a> {
47    buf: &'a [u8],
48    pos: usize,
49}
50
51impl<'a> Reader<'a> {
52    pub fn new(buf: &'a [u8]) -> Self {
53        Reader { buf, pos: 0 }
54    }
55    #[inline]
56    pub fn remaining(&self) -> usize {
57        self.buf.len() - self.pos
58    }
59    #[inline]
60    pub fn is_empty(&self) -> bool {
61        self.pos == self.buf.len()
62    }
63    pub fn position(&self) -> usize {
64        self.pos
65    }
66    pub fn rest(&self) -> &'a [u8] {
67        &self.buf[self.pos..]
68    }
69    #[inline]
70    pub fn u8(&mut self) -> Option<u8> {
71        let b = *self.buf.get(self.pos)?;
72        self.pos += 1;
73        Some(b)
74    }
75    pub fn u16(&mut self) -> Option<u16> {
76        let s = self.bytes(2)?;
77        Some(u16::from_be_bytes([s[0], s[1]]))
78    }
79    pub fn u32(&mut self) -> Option<u32> {
80        let s = self.bytes(4)?;
81        Some(u32::from_be_bytes([s[0], s[1], s[2], s[3]]))
82    }
83    #[inline]
84    pub fn bytes(&mut self, n: usize) -> Option<&'a [u8]> {
85        if self.remaining() < n {
86            return None;
87        }
88        let s = &self.buf[self.pos..self.pos + n];
89        self.pos += n;
90        Some(s)
91    }
92    /// Reads a varint, returning the value and its encoded length.
93    #[inline]
94    pub fn varint_with_len(&mut self) -> Option<(u64, usize)> {
95        let first = *self.buf.get(self.pos)?;
96        let len = 1usize << (first >> 6);
97        let s = self.bytes(len)?;
98        let mut v = (first & 0x3f) as u64;
99        for &b in &s[1..] {
100            v = (v << 8) | b as u64;
101        }
102        Some((v, len))
103    }
104    #[inline]
105    pub fn varint(&mut self) -> Option<u64> {
106        self.varint_with_len().map(|(v, _)| v)
107    }
108    /// Reads a varint length followed by that many bytes.
109    pub fn vbytes(&mut self) -> Option<&'a [u8]> {
110        let n = self.varint()?;
111        if n > self.remaining() as u64 {
112            return None;
113        }
114        self.bytes(n as usize)
115    }
116}
117
118/// Decodes a single varint from the front of `buf`.
119pub fn get_varint(buf: &[u8]) -> Option<(u64, usize)> {
120    Reader::new(buf).varint_with_len()
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126
127    #[test]
128    fn rfc9000_appendix_a1_examples() {
129        let cases: [(&[u8], u64); 4] = [
130            (&[0xc2, 0x19, 0x7c, 0x5e, 0xff, 0x14, 0xe8, 0x8c], 151_288_809_941_952_652),
131            (&[0x9d, 0x7f, 0x3e, 0x7d], 494_878_333),
132            (&[0x7b, 0xbd], 15_293),
133            (&[0x25], 37),
134        ];
135        for (bytes, v) in cases {
136            assert_eq!(get_varint(bytes), Some((v, bytes.len())));
137            let mut out = Vec::new();
138            put_varint(&mut out, v);
139            assert_eq!(out, bytes);
140        }
141        // A two-byte encoding of 37 decodes too (non-minimal is legal for varints).
142        assert_eq!(get_varint(&[0x40, 0x25]), Some((37, 2)));
143        assert_eq!(get_varint(&[0x40]), None);
144        assert_eq!(get_varint(&[]), None);
145    }
146
147    #[test]
148    fn boundaries() {
149        for v in [0, 63, 64, 16383, 16384, (1 << 30) - 1, 1 << 30, VARINT_MAX] {
150            let mut out = Vec::new();
151            put_varint(&mut out, v);
152            assert_eq!(out.len(), varint_len(v));
153            assert_eq!(get_varint(&out), Some((v, out.len())));
154            let mut fixed = [0u8; 8];
155            put_varint_fixed(&mut fixed, v.min(16383), 2);
156            assert_eq!(get_varint(&fixed[..2]).unwrap().0, v.min(16383));
157        }
158    }
159}