1use std::convert::{TryFrom, TryInto};
6use std::fmt;
7
8use thiserror::Error;
9
10use super::{Decode, DecodeError, Encode, EncodeError};
11
12#[derive(Debug, Copy, Clone, Eq, PartialEq, Error)]
14#[error("value out of range")]
15pub struct BoundsExceeded;
16
17#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
23pub struct VarInt(u64);
24
25impl VarInt {
26 pub const MAX: Self = Self((1 << 62) - 1);
28
29 pub const ZERO: Self = Self(0);
31
32 pub const fn from_u32(x: u32) -> Self {
35 Self(x as u64)
36 }
37
38 pub const fn from_u64(x: u64) -> Option<Self> {
40 if x <= Self::MAX.0 { Some(Self(x)) } else { None }
41 }
42
43 pub const fn from_u128(x: u128) -> Option<Self> {
45 if x <= Self::MAX.0 as u128 {
46 Some(Self(x as u64))
47 } else {
48 None
49 }
50 }
51
52 pub const fn into_inner(self) -> u64 {
54 self.0
55 }
56
57 pub const fn from_zigzag(signed: i64) -> Result<Self, BoundsExceeded> {
63 const RANGE: i64 = 1 << 61;
64 if signed < -RANGE || signed >= RANGE {
65 return Err(BoundsExceeded);
66 }
67 Ok(Self(((signed << 1) ^ (signed >> 63)) as u64))
68 }
69
70 pub const fn to_zigzag(self) -> i64 {
72 let v = self.0;
73 ((v >> 1) as i64) ^ -((v & 1) as i64)
74 }
75}
76
77impl From<VarInt> for u64 {
78 fn from(x: VarInt) -> Self {
79 x.0
80 }
81}
82
83impl From<VarInt> for usize {
84 fn from(x: VarInt) -> Self {
85 x.0 as usize
86 }
87}
88
89impl From<VarInt> for u128 {
90 fn from(x: VarInt) -> Self {
91 x.0 as u128
92 }
93}
94
95impl From<u8> for VarInt {
96 fn from(x: u8) -> Self {
97 Self(x.into())
98 }
99}
100
101impl From<u16> for VarInt {
102 fn from(x: u16) -> Self {
103 Self(x.into())
104 }
105}
106
107impl From<u32> for VarInt {
108 fn from(x: u32) -> Self {
109 Self(x.into())
110 }
111}
112
113impl TryFrom<u64> for VarInt {
114 type Error = BoundsExceeded;
115
116 fn try_from(x: u64) -> Result<Self, BoundsExceeded> {
118 let x = Self(x);
119 if x <= Self::MAX { Ok(x) } else { Err(BoundsExceeded) }
120 }
121}
122
123impl TryFrom<u128> for VarInt {
124 type Error = BoundsExceeded;
125
126 fn try_from(x: u128) -> Result<Self, BoundsExceeded> {
128 if x <= Self::MAX.into() {
129 Ok(Self(x as u64))
130 } else {
131 Err(BoundsExceeded)
132 }
133 }
134}
135
136impl TryFrom<usize> for VarInt {
137 type Error = BoundsExceeded;
138
139 fn try_from(x: usize) -> Result<Self, BoundsExceeded> {
141 Self::try_from(x as u64)
142 }
143}
144
145impl TryFrom<VarInt> for u32 {
146 type Error = BoundsExceeded;
147
148 fn try_from(x: VarInt) -> Result<Self, BoundsExceeded> {
150 if x.0 <= u32::MAX.into() {
151 Ok(x.0 as u32)
152 } else {
153 Err(BoundsExceeded)
154 }
155 }
156}
157
158impl TryFrom<VarInt> for u16 {
159 type Error = BoundsExceeded;
160
161 fn try_from(x: VarInt) -> Result<Self, BoundsExceeded> {
163 if x.0 <= u16::MAX.into() {
164 Ok(x.0 as u16)
165 } else {
166 Err(BoundsExceeded)
167 }
168 }
169}
170
171impl TryFrom<VarInt> for u8 {
172 type Error = BoundsExceeded;
173
174 fn try_from(x: VarInt) -> Result<Self, BoundsExceeded> {
176 if x.0 <= u8::MAX.into() {
177 Ok(x.0 as u8)
178 } else {
179 Err(BoundsExceeded)
180 }
181 }
182}
183
184impl fmt::Display for VarInt {
185 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
186 self.0.fmt(f)
187 }
188}
189
190impl VarInt {
191 pub fn decode_quic<R: bytes::Buf>(r: &mut R) -> Result<Self, DecodeError> {
193 if !r.has_remaining() {
194 return Err(DecodeError::Short);
195 }
196
197 let b = r.get_u8();
198 let tag = b >> 6;
199
200 let mut buf = [0u8; 8];
201 buf[0] = b & 0b0011_1111;
202
203 let x = match tag {
204 0b00 => u64::from(buf[0]),
205 0b01 => {
206 if !r.has_remaining() {
207 return Err(DecodeError::Short);
208 }
209 r.copy_to_slice(buf[1..2].as_mut());
210 u64::from(u16::from_be_bytes(buf[..2].try_into().unwrap()))
211 }
212 0b10 => {
213 if r.remaining() < 3 {
214 return Err(DecodeError::Short);
215 }
216 r.copy_to_slice(buf[1..4].as_mut());
217 u64::from(u32::from_be_bytes(buf[..4].try_into().unwrap()))
218 }
219 0b11 => {
220 if r.remaining() < 7 {
221 return Err(DecodeError::Short);
222 }
223 r.copy_to_slice(buf[1..8].as_mut());
224 u64::from_be_bytes(buf)
225 }
226 _ => unreachable!(),
227 };
228
229 Ok(Self(x))
230 }
231
232 pub fn encode_quic<W: bytes::BufMut>(&self, w: &mut W) -> Result<(), EncodeError> {
234 let remaining = w.remaining_mut();
235 if self.0 < (1u64 << 6) {
236 if remaining < 1 {
237 return Err(EncodeError::Short);
238 }
239 w.put_u8(self.0 as u8);
240 } else if self.0 < (1u64 << 14) {
241 if remaining < 2 {
242 return Err(EncodeError::Short);
243 }
244 w.put_u16((0b01 << 14) | self.0 as u16);
245 } else if self.0 < (1u64 << 30) {
246 if remaining < 4 {
247 return Err(EncodeError::Short);
248 }
249 w.put_u32((0b10 << 30) | self.0 as u32);
250 } else if self.0 < (1u64 << 62) {
251 if remaining < 8 {
252 return Err(EncodeError::Short);
253 }
254 w.put_u64((0b11 << 62) | self.0);
255 } else {
256 return Err(BoundsExceeded.into());
257 }
258 Ok(())
259 }
260
261 fn decode_leading_ones<R: bytes::Buf>(r: &mut R, version: ietf::Version) -> Result<Self, DecodeError> {
274 if !r.has_remaining() {
275 return Err(DecodeError::Short);
276 }
277
278 let b = r.get_u8();
279 let ones = b.leading_ones() as usize;
280
281 match ones {
282 0 => {
283 Ok(Self(u64::from(b)))
285 }
286 1 => {
287 if !r.has_remaining() {
289 return Err(DecodeError::Short);
290 }
291 let hi = u64::from(b & 0x3F);
292 let lo = u64::from(r.get_u8());
293 Ok(Self((hi << 8) | lo))
294 }
295 2 => {
296 if r.remaining() < 2 {
298 return Err(DecodeError::Short);
299 }
300 let hi = u64::from(b & 0x1F);
301 let mut buf = [0u8; 2];
302 r.copy_to_slice(&mut buf);
303 Ok(Self((hi << 16) | u64::from(u16::from_be_bytes(buf))))
304 }
305 3 => {
306 if r.remaining() < 3 {
308 return Err(DecodeError::Short);
309 }
310 let hi = u64::from(b & 0x0F);
311 let mut buf = [0u8; 3];
312 r.copy_to_slice(&mut buf);
313 Ok(Self(
314 (hi << 24) | u64::from(buf[0]) << 16 | u64::from(buf[1]) << 8 | u64::from(buf[2]),
315 ))
316 }
317 4 => {
318 if r.remaining() < 4 {
320 return Err(DecodeError::Short);
321 }
322 let hi = u64::from(b & 0x07);
323 let mut buf = [0u8; 4];
324 r.copy_to_slice(&mut buf);
325 Ok(Self((hi << 32) | u64::from(u32::from_be_bytes(buf))))
326 }
327 5 => {
328 if r.remaining() < 5 {
330 return Err(DecodeError::Short);
331 }
332 let hi = u64::from(b & 0x03);
333 let mut buf = [0u8; 5];
334 r.copy_to_slice(&mut buf);
335 let lo = u64::from(buf[0]) << 32
336 | u64::from(buf[1]) << 24
337 | u64::from(buf[2]) << 16
338 | u64::from(buf[3]) << 8
339 | u64::from(buf[4]);
340 Ok(Self((hi << 40) | lo))
341 }
342 6 => {
343 if matches!(version, ietf::Version::Draft17) {
345 return Err(DecodeError::InvalidValue);
346 }
347 if r.remaining() < 6 {
348 return Err(DecodeError::Short);
349 }
350 let hi = u64::from(b & 0x01);
351 let mut buf = [0u8; 8];
352 r.copy_to_slice(&mut buf[2..]);
353 Ok(Self((hi << 48) | u64::from_be_bytes(buf)))
354 }
355 7 => {
356 if r.remaining() < 7 {
358 return Err(DecodeError::Short);
359 }
360 let mut buf = [0u8; 8];
361 buf[0] = 0;
362 r.copy_to_slice(&mut buf[1..]);
363 Ok(Self(u64::from_be_bytes(buf)))
364 }
365 8 => {
366 if r.remaining() < 8 {
368 return Err(DecodeError::Short);
369 }
370 let mut buf = [0u8; 8];
371 r.copy_to_slice(&mut buf);
372 Ok(Self(u64::from_be_bytes(buf)))
373 }
374 _ => unreachable!(),
375 }
376 }
377
378 fn encode_leading_ones<W: bytes::BufMut>(&self, w: &mut W, _version: ietf::Version) -> Result<(), EncodeError> {
384 let x = self.0;
385 let remaining = w.remaining_mut();
386
387 if x < (1 << 7) {
388 if remaining < 1 {
390 return Err(EncodeError::Short);
391 }
392 w.put_u8(x as u8);
393 } else if x < (1 << 14) {
394 if remaining < 2 {
396 return Err(EncodeError::Short);
397 }
398 w.put_u8(0x80 | (x >> 8) as u8);
399 w.put_u8(x as u8);
400 } else if x < (1 << 21) {
401 if remaining < 3 {
403 return Err(EncodeError::Short);
404 }
405 w.put_u8(0xC0 | (x >> 16) as u8);
406 w.put_u16(x as u16);
407 } else if x < (1 << 28) {
408 if remaining < 4 {
410 return Err(EncodeError::Short);
411 }
412 w.put_u8(0xE0 | (x >> 24) as u8);
413 w.put_u8((x >> 16) as u8);
414 w.put_u16(x as u16);
415 } else if x < (1 << 35) {
416 if remaining < 5 {
418 return Err(EncodeError::Short);
419 }
420 w.put_u8(0xF0 | (x >> 32) as u8);
421 w.put_u32(x as u32);
422 } else if x < (1 << 42) {
423 if remaining < 6 {
425 return Err(EncodeError::Short);
426 }
427 w.put_u8(0xF8 | (x >> 40) as u8);
428 w.put_u8((x >> 32) as u8);
429 w.put_u32(x as u32);
430 } else if x < (1 << 56) {
431 if remaining < 8 {
433 return Err(EncodeError::Short);
434 }
435 w.put_u8(0xFE);
436 w.put_u8((x >> 48) as u8);
438 w.put_u16((x >> 32) as u16);
439 w.put_u32(x as u32);
440 } else {
441 if remaining < 9 {
443 return Err(EncodeError::Short);
444 }
445 w.put_u8(0xFF);
446 w.put_u64(x);
447 }
448
449 Ok(())
450 }
451}
452
453use crate::{Version, ietf, lite};
454
455impl Encode<lite::Version> for VarInt {
457 fn encode<W: bytes::BufMut>(&self, w: &mut W, _: lite::Version) -> Result<(), EncodeError> {
458 self.encode_quic(w)
459 }
460}
461
462impl Decode<lite::Version> for VarInt {
463 fn decode<R: bytes::Buf>(r: &mut R, _: lite::Version) -> Result<Self, DecodeError> {
464 Self::decode_quic(r)
465 }
466}
467
468impl Encode<ietf::Version> for VarInt {
470 fn encode<W: bytes::BufMut>(&self, w: &mut W, version: ietf::Version) -> Result<(), EncodeError> {
471 match version {
472 ietf::Version::Draft14 | ietf::Version::Draft15 | ietf::Version::Draft16 => self.encode_quic(w),
473 _ => self.encode_leading_ones(w, version),
474 }
475 }
476}
477
478impl Decode<ietf::Version> for VarInt {
479 fn decode<R: bytes::Buf>(r: &mut R, version: ietf::Version) -> Result<Self, DecodeError> {
480 match version {
481 ietf::Version::Draft14 | ietf::Version::Draft15 | ietf::Version::Draft16 => Self::decode_quic(r),
482 _ => Self::decode_leading_ones(r, version),
483 }
484 }
485}
486
487impl Encode<Version> for VarInt {
489 fn encode<W: bytes::BufMut>(&self, w: &mut W, version: Version) -> Result<(), EncodeError> {
490 match version {
491 Version::Lite(v) => self.encode(w, v),
492 Version::Ietf(v) => self.encode(w, v),
493 }
494 }
495}
496
497impl Decode<Version> for VarInt {
498 fn decode<R: bytes::Buf>(r: &mut R, version: Version) -> Result<Self, DecodeError> {
499 match version {
500 Version::Lite(v) => Self::decode(r, v),
501 Version::Ietf(v) => Self::decode(r, v),
502 }
503 }
504}
505
506impl<V: Copy> Encode<V> for u64
508where
509 VarInt: Encode<V>,
510{
511 fn encode<W: bytes::BufMut>(&self, w: &mut W, version: V) -> Result<(), EncodeError> {
512 VarInt::try_from(*self)?.encode(w, version)
513 }
514}
515
516impl<V: Copy> Decode<V> for u64
517where
518 VarInt: Decode<V>,
519{
520 fn decode<R: bytes::Buf>(r: &mut R, version: V) -> Result<Self, DecodeError> {
521 VarInt::decode(r, version).map(|v| v.into_inner())
522 }
523}
524
525impl<V: Copy> Encode<V> for usize
526where
527 VarInt: Encode<V>,
528{
529 fn encode<W: bytes::BufMut>(&self, w: &mut W, version: V) -> Result<(), EncodeError> {
530 VarInt::try_from(*self)?.encode(w, version)
531 }
532}
533
534impl<V: Copy> Decode<V> for usize
535where
536 VarInt: Decode<V>,
537{
538 fn decode<R: bytes::Buf>(r: &mut R, version: V) -> Result<Self, DecodeError> {
539 VarInt::decode(r, version).map(|v| v.into_inner() as usize)
540 }
541}
542
543impl<V: Copy> Encode<V> for u32
544where
545 VarInt: Encode<V>,
546{
547 fn encode<W: bytes::BufMut>(&self, w: &mut W, version: V) -> Result<(), EncodeError> {
548 VarInt::from(*self).encode(w, version)
549 }
550}
551
552impl<V: Copy> Decode<V> for u32
553where
554 VarInt: Decode<V>,
555{
556 fn decode<R: bytes::Buf>(r: &mut R, version: V) -> Result<Self, DecodeError> {
557 let v = VarInt::decode(r, version)?;
558 let v = v.try_into().map_err(|_| DecodeError::BoundsExceeded)?;
559 Ok(v)
560 }
561}
562
563#[cfg(test)]
564mod tests {
565 use super::*;
566 use crate::{ietf, lite};
567 use bytes::Bytes;
568
569 #[test]
572 fn leading_ones_spec_examples() {
573 let cases: &[(&[u8], u64)] = &[
574 (&[0x25], 37),
575 (&[0x80, 0x25], 37),
576 (&[0xbb, 0xbd], 15_293),
577 (&[0xfa, 0xa1, 0xa0, 0xe4, 0x03, 0xd8], 2_893_212_287_960),
580 (
581 &[0xfe, 0xfa, 0x31, 0x8f, 0xa8, 0xe3, 0xca, 0x11],
582 70_423_237_261_249_041,
583 ),
584 (
585 &[0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff],
586 18_446_744_073_709_551_615,
587 ),
588 ];
589
590 for (bytes, expected) in cases {
591 let mut buf = Bytes::from(bytes.to_vec());
593 let decoded = VarInt::decode_leading_ones(&mut buf, ietf::Version::Draft17).expect("decode should succeed");
594 assert_eq!(
595 decoded.into_inner(),
596 *expected,
597 "decode mismatch for bytes {bytes:02x?}"
598 );
599 assert_eq!(buf.len(), 0, "all bytes should be consumed for {bytes:02x?}");
600
601 if let Some(varint) = VarInt::from_u64(*expected)
605 && (bytes.len() == 1 || *expected != 37)
606 {
607 let mut encoded = Vec::new();
608 varint
609 .encode_leading_ones(&mut encoded, ietf::Version::Draft17)
610 .expect("encode should succeed");
611 assert_eq!(&encoded, bytes, "encode mismatch for value {expected}");
612 }
613 }
614 }
615
616 #[test]
618 fn leading_ones_invalid_0xfc() {
619 let mut buf = Bytes::from_static(&[0xFC]);
620 assert!(
621 matches!(
622 VarInt::decode_leading_ones(&mut buf, ietf::Version::Draft17),
623 Err(DecodeError::InvalidValue)
624 ),
625 "0xFC should be rejected as invalid on draft-17"
626 );
627 }
628
629 #[test]
630 fn leading_ones_boundaries_round_trip() {
631 let cases = [
632 ((1u64 << 7) - 1, 1usize),
633 (1u64 << 7, 2usize),
634 ((1u64 << 14) - 1, 2usize),
635 (1u64 << 14, 3usize),
636 ((1u64 << 56) - 1, 8usize),
637 (1u64 << 56, 9usize),
638 ];
639
640 for (value, expected_len) in cases {
641 let varint = VarInt::from_u64(value).expect("value should be representable as VarInt");
642 let mut encoded = Vec::new();
643 varint
644 .encode_leading_ones(&mut encoded, ietf::Version::Draft17)
645 .expect("leading-ones encode should succeed");
646 assert_eq!(
647 encoded.len(),
648 expected_len,
649 "unexpected encoded length for value {value}"
650 );
651
652 let mut bytes = Bytes::from(encoded);
653 let decoded = VarInt::decode_leading_ones(&mut bytes, ietf::Version::Draft17)
654 .expect("leading-ones decode should succeed");
655 assert_eq!(decoded.into_inner(), value, "round-trip mismatch for value {value}");
656 }
657 }
658
659 #[test]
660 fn draft17_rejects_7_byte_varint() {
661 let bytes = Bytes::from(vec![0xFC, 0, 0, 0, 0, 0, 0]);
663 let mut buf = bytes.clone();
664 let err = VarInt::decode_leading_ones(&mut buf, ietf::Version::Draft17).unwrap_err();
665 assert!(matches!(err, DecodeError::InvalidValue));
666 }
667
668 #[test]
669 fn zigzag_roundtrip_small() {
670 for n in [-3i64, -2, -1, 0, 1, 2, 3, 100, -100] {
671 let v = VarInt::from_zigzag(n).unwrap();
672 assert_eq!(v.to_zigzag(), n, "roundtrip failed for {}", n);
673 }
674 }
675
676 #[test]
677 fn zigzag_small_values_compact() {
678 assert_eq!(VarInt::from_zigzag(0).unwrap().into_inner(), 0);
680 assert_eq!(VarInt::from_zigzag(-1).unwrap().into_inner(), 1);
681 assert_eq!(VarInt::from_zigzag(1).unwrap().into_inner(), 2);
682 assert_eq!(VarInt::from_zigzag(-2).unwrap().into_inner(), 3);
683 assert_eq!(VarInt::from_zigzag(2).unwrap().into_inner(), 4);
684 }
685
686 #[test]
687 fn zigzag_roundtrip_boundary() {
688 let max = (1i64 << 61) - 1;
690 let min = -(1i64 << 61);
691 let mid = (1i64 << 30) + 17;
692
693 for n in [max, min, mid, -mid] {
694 let v = VarInt::from_zigzag(n).unwrap();
695 assert_eq!(v.to_zigzag(), n);
696 }
697 }
698
699 #[test]
700 fn zigzag_out_of_range_rejected() {
701 assert!(VarInt::from_zigzag(1i64 << 61).is_err());
703 assert!(VarInt::from_zigzag(-(1i64 << 61) - 1).is_err());
704 assert!(VarInt::from_zigzag(i64::MAX).is_err());
705 assert!(VarInt::from_zigzag(i64::MIN).is_err());
706 }
707
708 #[test]
709 fn zigzag_quic_varint_roundtrip() {
710 for n in [-5000i64, 0, 100, -1, 1_000_000, -1_000_000] {
712 let v = VarInt::from_zigzag(n).unwrap();
713
714 let mut buf = bytes::BytesMut::new();
715 v.encode(&mut buf, lite::Version::Lite01).unwrap();
716 let mut bytes = buf.freeze();
717 let decoded = VarInt::decode(&mut bytes, lite::Version::Lite01).unwrap();
718 assert_eq!(decoded.to_zigzag(), n);
719 }
720 }
721
722 #[test]
723 fn draft18_accepts_7_byte_varint() {
724 let value: u64 = 0x1234_5678_9ABC;
726 let mut bytes = Vec::new();
727 let hi_bit = ((value >> 48) & 0x01) as u8;
730 bytes.push(0xFC | hi_bit);
731 for shift in (0..48).step_by(8).rev() {
732 bytes.push(((value >> shift) & 0xFF) as u8);
733 }
734 let mut buf = Bytes::from(bytes);
735 let decoded = VarInt::decode_leading_ones(&mut buf, ietf::Version::Draft18).unwrap();
736 assert_eq!(decoded.into_inner(), value);
737 }
738}