1use std::{borrow::Cow, cell::Cell, cmp::min, convert::TryFrom, fmt, rc::Rc};
4
5use ntex_bytes::{Buf, BufMut, BytePages, Bytes};
6
7pub const MIN_TAG: u32 = 1;
8pub const MAX_TAG: u32 = (1 << 29) - 1;
9
10#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
11#[repr(u8)]
12pub enum WireType {
13 Varint = 0,
14 SixtyFourBit = 1,
15 LengthDelimited = 2,
16 StartGroup = 3,
17 EndGroup = 4,
18 ThirtyTwoBit = 5,
19}
20
21impl TryFrom<u64> for WireType {
22 type Error = DecodeError;
23
24 #[inline]
25 fn try_from(value: u64) -> Result<Self, Self::Error> {
26 match value {
27 0 => Ok(WireType::Varint),
28 1 => Ok(WireType::SixtyFourBit),
29 2 => Ok(WireType::LengthDelimited),
30 3 => Ok(WireType::StartGroup),
31 4 => Ok(WireType::EndGroup),
32 5 => Ok(WireType::ThirtyTwoBit),
33 _ => Err(DecodeError::new(format!(
34 "invalid wire type value: {value}"
35 ))),
36 }
37 }
38}
39
40#[inline]
43pub fn encoded_len_varint(value: u64) -> usize {
44 ((((value | 1).leading_zeros() ^ 63) * 9 + 73) / 64) as usize
47}
48
49#[inline]
52pub fn encode_varint(mut value: u64, buf: &mut BytePages) {
53 loop {
54 if value < 0x80 {
55 buf.put_u8(value as u8);
56 break;
57 }
58 buf.put_u8(((value & 0x7F) | 0x80) as u8);
59 value >>= 7;
60 }
61}
62
63#[inline]
65pub fn decode_varint(buf: &mut Bytes) -> Result<u64, DecodeError> {
66 let bytes = buf.as_ref();
67 let len = buf.len();
68 if len == 0 {
69 return Err(DecodeError::new("invalid varint"));
70 }
71
72 let byte = bytes[0];
73 if byte < 0x80 {
74 buf.advance(1);
75 Ok(u64::from(byte))
76 } else if len > 10 || bytes[len - 1] < 0x80 {
77 let (value, advance) = decode_varint_slice(bytes)?;
78 buf.advance(advance);
79 Ok(value)
80 } else {
81 decode_varint_slow(buf)
82 }
83}
84
85#[inline]
99#[allow(clippy::assert_is_empty)]
100fn decode_varint_slice(bytes: &[u8]) -> Result<(u64, usize), DecodeError> {
101 assert!(!bytes.is_empty());
105 assert!(bytes.len() > 10 || bytes[bytes.len() - 1] < 0x80);
106
107 let mut b: u8 = unsafe { *bytes.get_unchecked(0) };
108 let mut part0: u32 = u32::from(b);
109 if b < 0x80 {
110 return Ok((u64::from(part0), 1));
111 }
112 part0 -= 0x80;
113 b = unsafe { *bytes.get_unchecked(1) };
114 part0 += u32::from(b) << 7;
115 if b < 0x80 {
116 return Ok((u64::from(part0), 2));
117 }
118 part0 -= 0x80 << 7;
119 b = unsafe { *bytes.get_unchecked(2) };
120 part0 += u32::from(b) << 14;
121 if b < 0x80 {
122 return Ok((u64::from(part0), 3));
123 }
124 part0 -= 0x80 << 14;
125 b = unsafe { *bytes.get_unchecked(3) };
126 part0 += u32::from(b) << 21;
127 if b < 0x80 {
128 return Ok((u64::from(part0), 4));
129 }
130 part0 -= 0x80 << 21;
131 let value = u64::from(part0);
132
133 b = unsafe { *bytes.get_unchecked(4) };
134 let mut part1: u32 = u32::from(b);
135 if b < 0x80 {
136 return Ok((value + (u64::from(part1) << 28), 5));
137 }
138 part1 -= 0x80;
139 b = unsafe { *bytes.get_unchecked(5) };
140 part1 += u32::from(b) << 7;
141 if b < 0x80 {
142 return Ok((value + (u64::from(part1) << 28), 6));
143 }
144 part1 -= 0x80 << 7;
145 b = unsafe { *bytes.get_unchecked(6) };
146 part1 += u32::from(b) << 14;
147 if b < 0x80 {
148 return Ok((value + (u64::from(part1) << 28), 7));
149 }
150 part1 -= 0x80 << 14;
151 b = unsafe { *bytes.get_unchecked(7) };
152 part1 += u32::from(b) << 21;
153 if b < 0x80 {
154 return Ok((value + (u64::from(part1) << 28), 8));
155 }
156 part1 -= 0x80 << 21;
157 let value = value + ((u64::from(part1)) << 28);
158
159 b = unsafe { *bytes.get_unchecked(8) };
160 let mut part2: u32 = u32::from(b);
161 if b < 0x80 {
162 return Ok((value + (u64::from(part2) << 56), 9));
163 }
164 part2 -= 0x80;
165 b = unsafe { *bytes.get_unchecked(9) };
166 part2 += u32::from(b) << 7;
167 if b < 0x02 {
170 return Ok((value + (u64::from(part2) << 56), 10));
171 }
172
173 Err(DecodeError::new("invalid varint"))
176}
177
178#[inline(never)]
185#[cold]
186fn decode_varint_slow(buf: &mut Bytes) -> Result<u64, DecodeError> {
187 let mut value = 0;
188 for count in 0..min(10, buf.remaining()) {
189 let byte = buf.get_u8();
190 value |= u64::from(byte & 0x7F) << (count * 7);
191 if byte <= 0x7F {
192 return if count == 9 && byte >= 0x02 {
195 Err(DecodeError::new("invalid varint"))
196 } else {
197 Ok(value)
198 };
199 }
200 }
201
202 Err(DecodeError::new("invalid varint"))
203}
204
205#[inline]
208pub fn encode_key(tag: u32, wire_type: WireType, buf: &mut BytePages) {
209 debug_assert!((MIN_TAG..=MAX_TAG).contains(&tag));
210 let key = (tag << 3) | wire_type as u32;
211 encode_varint(u64::from(key), buf);
212}
213
214#[inline]
217pub fn decode_key(buf: &mut Bytes) -> Result<(u32, WireType), DecodeError> {
218 let key = decode_varint(buf)?;
219 if key > u64::from(u32::MAX) {
220 return Err(DecodeError::new(format!("invalid key value: {key}")));
221 }
222 let wire_type = WireType::try_from(key & 0x07)?;
223 let tag = key as u32 >> 3;
224
225 if tag < MIN_TAG {
226 return Err(DecodeError::new("invalid tag value: 0"));
227 }
228
229 Ok((tag, wire_type))
230}
231
232#[inline]
235pub fn key_len(tag: u32) -> usize {
236 encoded_len_varint(u64::from(tag << 3))
237}
238
239#[inline]
242pub fn check_wire_type(expected: WireType, actual: WireType) -> Result<(), DecodeError> {
243 if expected != actual {
244 return Err(DecodeError::new(format!(
245 "invalid wire type: {actual:?} (expected {expected:?})",
246 )));
247 }
248 Ok(())
249}
250
251const GROUP_DEPTH_LIMIT: usize = 100;
254
255pub(crate) const RECURSION_LIMIT: u32 = 100;
258
259thread_local! {
260 static DEPTH: Cell<u32> = const { Cell::new(0) };
262}
263
264pub(crate) fn merge_nested<F>(f: F) -> Result<(), DecodeError>
267where
268 F: FnOnce() -> Result<(), DecodeError>,
269{
270 struct Guard;
271
272 impl Drop for Guard {
273 fn drop(&mut self) {
274 DEPTH.with(|d| d.set(d.get() - 1));
275 }
276 }
277
278 let depth = DEPTH.with(|d| {
279 let depth = d.get() + 1;
280 d.set(depth);
281 depth
282 });
283 let _guard = Guard;
284 if depth > RECURSION_LIMIT {
285 Err(DecodeError::new("recursion limit reached"))
286 } else {
287 f()
288 }
289}
290
291pub fn split_group(tag: u32, buf: &mut Bytes) -> Result<Bytes, DecodeError> {
298 let mut open = [0u32; GROUP_DEPTH_LIMIT];
301 let mut depth = 0;
302 let mut rest = buf.clone();
303 loop {
304 let body_len = buf.len() - rest.len();
305 let (inner_tag, inner_wire_type) = decode_key(&mut rest)?;
306 match inner_wire_type {
307 WireType::StartGroup => {
308 if depth == GROUP_DEPTH_LIMIT {
309 return Err(DecodeError::new("groups nested too deeply"));
310 }
311 open[depth] = inner_tag;
312 depth += 1;
313 }
314 WireType::EndGroup => {
315 let expected = if depth == 0 { tag } else { open[depth - 1] };
316 if inner_tag != expected {
317 return Err(DecodeError::new("unexpected end group tag"));
318 }
319 if depth == 0 {
320 let body = buf.split_to(body_len);
321 *buf = rest;
322 return Ok(body);
323 }
324 depth -= 1;
325 }
326 _ => skip_field(inner_wire_type, inner_tag, &mut rest)?,
327 }
328 }
329}
330
331pub fn skip_field(wire_type: WireType, tag: u32, buf: &mut Bytes) -> Result<(), DecodeError> {
332 let len = match wire_type {
333 WireType::Varint => decode_varint(buf).map(|_| 0)?,
334 WireType::ThirtyTwoBit => 4,
335 WireType::SixtyFourBit => 8,
336 WireType::LengthDelimited => decode_varint(buf)?,
337 WireType::StartGroup => split_group(tag, buf).map(|_| 0)?,
338 WireType::EndGroup => return Err(DecodeError::new("unexpected end group tag")),
339 };
340
341 buf.split_to_checked(len as usize)
342 .ok_or_else(DecodeError::incomplete)?;
343 Ok(())
344}
345
346#[derive(Clone, PartialEq, Eq)]
348pub struct DecodeError {
349 inner: Rc<Inner>,
350}
351
352#[derive(Clone, PartialEq, Eq)]
353struct Inner {
354 description: Cow<'static, str>,
356 stack: Vec<(&'static str, &'static str)>,
360}
361
362impl DecodeError {
363 #[doc(hidden)]
367 #[cold]
368 pub fn new(description: impl Into<Cow<'static, str>>) -> DecodeError {
369 DecodeError {
370 inner: Rc::new(Inner {
371 description: description.into(),
372 stack: Vec::new(),
373 }),
374 }
375 }
376
377 #[doc(hidden)]
381 #[must_use]
382 pub fn push(mut self, message: &'static str, field: &'static str) -> Self {
383 let inner = if let Some(inner) = Rc::get_mut(&mut self.inner) {
384 inner
385 } else {
386 self.inner = Rc::new(Inner {
387 description: self.inner.description.clone(),
388 stack: self.inner.stack.clone(),
389 });
390 Rc::get_mut(&mut self.inner).unwrap()
391 };
392 inner.stack.push((message, field));
393 self
394 }
395
396 pub(crate) fn incomplete() -> Self {
397 Self::new("Not enough data")
398 }
399}
400
401impl fmt::Debug for DecodeError {
402 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
403 f.debug_struct("DecodeError")
404 .field("description", &self.inner.description)
405 .field("stack", &self.inner.stack)
406 .finish()
407 }
408}
409
410impl fmt::Display for DecodeError {
411 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
412 f.write_str("failed to decode Protobuf message: ")?;
413 for &(message, field) in &self.inner.stack {
414 write!(f, "{message}.{field}: ")?;
415 }
416 f.write_str(&self.inner.description)
417 }
418}
419
420impl std::error::Error for DecodeError {}
421
422#[cfg(test)]
423mod tests {
424 use super::*;
425
426 fn encode(value: u64) -> Bytes {
427 let mut buf = BytePages::default();
428 encode_varint(value, &mut buf);
429 buf.freeze()
430 }
431
432 #[test]
433 fn varint_roundtrip() {
434 let cases: [(u64, usize); 18] = [
435 (0, 1),
436 (1, 1),
437 (0x7f, 1),
438 (0x80, 2),
439 (300, 2),
440 (0x3fff, 2),
441 (0x4000, 3),
442 (0x1f_ffff, 3),
443 (0x20_0000, 4),
444 (0x0fff_ffff, 4),
445 (0x1000_0000, 5),
446 (u64::from(u32::MAX), 5),
447 (1 << 35, 6),
448 (1 << 42, 7),
449 (1 << 49, 8),
450 (1 << 56, 9),
451 (1 << 63, 10),
452 (u64::MAX, 10),
453 ];
454
455 for (value, len) in cases {
456 let bytes = encode(value);
457 assert_eq!(bytes.len(), len, "{value}");
458 assert_eq!(encoded_len_varint(value), len, "{value}");
459
460 let mut buf = bytes.clone();
461 assert_eq!(decode_varint(&mut buf).unwrap(), value);
462 assert!(buf.is_empty());
463
464 assert_eq!(decode_varint_slice(&bytes).unwrap(), (value, len));
466 }
467 }
468
469 #[test]
470 fn varint_decode_errors() {
471 assert!(decode_varint(&mut Bytes::new()).is_err());
473
474 let mut overflow = vec![0xff_u8; 9];
476 overflow.push(0x02);
477 assert!(decode_varint(&mut Bytes::from(overflow.clone())).is_err());
478 assert!(decode_varint_slice(&overflow).is_err());
479
480 assert!(decode_varint(&mut Bytes::from(vec![0xff_u8; 11])).is_err());
482
483 assert!(decode_varint(&mut Bytes::from_static(&[0x80])).is_err());
485 assert!(decode_varint(&mut Bytes::from(vec![0xff_u8; 10])).is_err());
486 }
487
488 #[test]
489 fn varint_slow_path_stops_at_terminator() {
490 let mut buf = Bytes::from_static(&[0x80, 0x01, 0x80]);
492 assert_eq!(decode_varint(&mut buf).unwrap(), 128);
493 assert_eq!(buf, Bytes::from_static(&[0x80]));
494 }
495
496 #[test]
497 fn wire_type_from_u64() {
498 let all = [
499 WireType::Varint,
500 WireType::SixtyFourBit,
501 WireType::LengthDelimited,
502 WireType::StartGroup,
503 WireType::EndGroup,
504 WireType::ThirtyTwoBit,
505 ];
506 for (value, wire_type) in all.iter().enumerate() {
507 assert_eq!(WireType::try_from(value as u64).unwrap(), *wire_type);
508 }
509
510 let err = WireType::try_from(6).unwrap_err();
511 assert!(err.to_string().contains("invalid wire type value: 6"));
512 assert!(WireType::try_from(u64::MAX).is_err());
513 }
514
515 #[test]
516 fn key_roundtrip() {
517 let cases = [
518 (1_u32, WireType::Varint),
519 (2, WireType::SixtyFourBit),
520 (3, WireType::LengthDelimited),
521 (4, WireType::StartGroup),
522 (5, WireType::EndGroup),
523 (16, WireType::ThirtyTwoBit),
524 (2047, WireType::Varint),
525 (MAX_TAG, WireType::LengthDelimited),
526 ];
527
528 for (tag, wire_type) in cases {
529 let mut buf = BytePages::default();
530 encode_key(tag, wire_type, &mut buf);
531 let bytes = buf.freeze();
532 assert_eq!(bytes.len(), key_len(tag), "{tag}");
533
534 let mut src = bytes;
535 assert_eq!(decode_key(&mut src).unwrap(), (tag, wire_type));
536 assert!(src.is_empty());
537 }
538 }
539
540 #[test]
541 fn key_decode_errors() {
542 let err = decode_key(&mut encode(u64::from(u32::MAX) + 1)).unwrap_err();
544 assert!(err.to_string().contains("invalid key value"));
545
546 for wire_type in [6_u64, 7] {
548 let err = decode_key(&mut encode((8 << 3) | wire_type)).unwrap_err();
549 assert!(err.to_string().contains("invalid wire type value"));
550 }
551
552 let err = decode_key(&mut Bytes::from_static(&[0x02])).unwrap_err();
554 assert!(err.to_string().contains("invalid tag value: 0"));
555
556 assert!(decode_key(&mut Bytes::new()).is_err());
557 }
558
559 #[test]
560 fn check_wire_type_mismatch() {
561 assert!(check_wire_type(WireType::Varint, WireType::Varint).is_ok());
562
563 let err = check_wire_type(WireType::Varint, WireType::LengthDelimited).unwrap_err();
564 assert_eq!(
565 err.to_string(),
566 "failed to decode Protobuf message: invalid wire type: LengthDelimited (expected Varint)"
567 );
568 }
569
570 #[test]
571 fn skip_simple_fields() {
572 let cases: [(WireType, &[u8]); 4] = [
573 (WireType::Varint, &[0x96, 0x01]),
574 (WireType::ThirtyTwoBit, &[1, 2, 3, 4]),
575 (WireType::SixtyFourBit, &[1, 2, 3, 4, 5, 6, 7, 8]),
576 (WireType::LengthDelimited, &[0x03, b'a', b'b', b'c']),
577 ];
578
579 for (wire_type, data) in cases {
580 let mut src = BytePages::default();
581 src.extend_from_slice(data);
582 src.extend_from_slice(&[0xff]);
583 let mut buf = src.freeze();
584
585 skip_field(wire_type, 1, &mut buf).unwrap();
586 assert_eq!(buf, Bytes::from_static(&[0xff]), "{wire_type:?}");
587 }
588 }
589
590 #[test]
591 fn skip_group_field() {
592 let mut buf = BytePages::default();
594 encode_key(1, WireType::Varint, &mut buf);
595 encode_varint(1, &mut buf);
596 encode_key(2, WireType::StartGroup, &mut buf);
597 encode_key(1, WireType::Varint, &mut buf);
598 encode_varint(7, &mut buf);
599 encode_key(2, WireType::EndGroup, &mut buf);
600 encode_key(3, WireType::EndGroup, &mut buf);
601 buf.extend_from_slice(&[0xff]);
602
603 let mut buf = buf.freeze();
604 skip_field(WireType::StartGroup, 3, &mut buf).unwrap();
605 assert_eq!(buf, Bytes::from_static(&[0xff]));
606 }
607
608 #[test]
609 fn split_group_body() {
610 let mut buf = BytePages::default();
612 encode_key(1, WireType::Varint, &mut buf);
613 encode_varint(1, &mut buf);
614 encode_key(2, WireType::StartGroup, &mut buf);
615 encode_key(2, WireType::EndGroup, &mut buf);
616 encode_key(3, WireType::EndGroup, &mut buf);
617 buf.extend_from_slice(&[0xff]);
618
619 let mut buf = buf.freeze();
620 let body = split_group(3, &mut buf).unwrap();
621 assert_eq!(body, Bytes::from_static(&[0x08, 0x01, 0x13, 0x14]));
622 assert_eq!(buf, Bytes::from_static(&[0xff]));
623
624 let mut buf = Bytes::from_static(&[0x1c]);
626 assert!(split_group(3, &mut buf).unwrap().is_empty());
627 assert!(buf.is_empty());
628
629 let mut buf = Bytes::from_static(&[0x08, 0x01, 0x14]);
631 assert!(split_group(3, &mut buf).is_err());
632 assert_eq!(buf.len(), 3);
633 }
634
635 fn nested_groups(depth: usize) -> Bytes {
636 let mut buf = BytePages::default();
637 for _ in 0..depth {
638 encode_key(5, WireType::StartGroup, &mut buf);
639 }
640 for _ in 0..depth {
641 encode_key(5, WireType::EndGroup, &mut buf);
642 }
643 encode_key(3, WireType::EndGroup, &mut buf);
644 buf.freeze()
645 }
646
647 #[test]
648 fn group_depth_limit() {
649 let mut buf = nested_groups(GROUP_DEPTH_LIMIT);
650 skip_field(WireType::StartGroup, 3, &mut buf).unwrap();
651 assert!(buf.is_empty());
652
653 let mut buf = nested_groups(GROUP_DEPTH_LIMIT + 1);
654 let err = skip_field(WireType::StartGroup, 3, &mut buf).unwrap_err();
655 assert!(err.to_string().contains("groups nested too deeply"));
656
657 let mut buf = nested_groups(100_000);
659 assert!(split_group(3, &mut buf).is_err());
660
661 let mut buf = BytePages::default();
663 encode_key(5, WireType::StartGroup, &mut buf);
664 encode_key(6, WireType::EndGroup, &mut buf);
665 encode_key(3, WireType::EndGroup, &mut buf);
666 let err = split_group(3, &mut buf.freeze()).unwrap_err();
667 assert!(err.to_string().contains("unexpected end group tag"));
668 }
669
670 #[test]
671 fn skip_field_errors() {
672 let mut buf = BytePages::default();
674 encode_key(4, WireType::EndGroup, &mut buf);
675 let err = skip_field(WireType::StartGroup, 3, &mut buf.freeze()).unwrap_err();
676 assert!(err.to_string().contains("unexpected end group tag"));
677
678 let err = skip_field(WireType::EndGroup, 1, &mut Bytes::new()).unwrap_err();
680 assert!(err.to_string().contains("unexpected end group tag"));
681
682 let mut buf = Bytes::from_static(&[0x05, 1, 2]);
684 let err = skip_field(WireType::LengthDelimited, 1, &mut buf).unwrap_err();
685 assert!(err.to_string().contains("Not enough data"));
686
687 let mut buf = Bytes::from_static(&[0x08]);
689 assert!(skip_field(WireType::StartGroup, 3, &mut buf).is_err());
690 }
691
692 #[test]
693 fn decode_error_stack() {
694 let err = DecodeError::new("boom");
695 assert_eq!(err.to_string(), "failed to decode Protobuf message: boom");
696
697 let err = err.push("Msg", "field");
699 assert_eq!(
700 err.to_string(),
701 "failed to decode Protobuf message: Msg.field: boom"
702 );
703
704 let shared = err.clone();
706 let err = err.push("Outer", "inner");
707 assert_eq!(
708 err.to_string(),
709 "failed to decode Protobuf message: Msg.field: Outer.inner: boom"
710 );
711 assert_eq!(
712 shared.to_string(),
713 "failed to decode Protobuf message: Msg.field: boom"
714 );
715 assert_ne!(err, shared);
716 assert_eq!(shared, shared.clone());
717
718 let dbg = format!("{err:?}");
719 assert!(dbg.contains("DecodeError"));
720 assert!(dbg.contains("boom"));
721 assert!(dbg.contains("Outer"));
722 }
723
724 #[test]
725 fn decode_error_incomplete() {
726 assert_eq!(
727 DecodeError::incomplete().to_string(),
728 "failed to decode Protobuf message: Not enough data"
729 );
730 }
731}