From 87cea1d49a4e5df1d4bfd3c26045a444c5f37955 Mon Sep 17 00:00:00 2001
From: Andrew Guerrero <37567782+ajguerrer@users.noreply.github.com>
Date: Mon, 20 Sep 2021 09:49:30 -0600
Subject: [PATCH] Protect encoding::decode_varint from overflow (#527)

---
 src/encoding.rs | 33 +++++++++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/encoding.rs b/src/encoding.rs
index 0c3d37669..093083111 100644
--- a/src/encoding.rs
+++ b/src/encoding.rs
@@ -66,7 +66,8 @@ where
 /// Decodes a LEB128-encoded variable length integer from the slice, returning the value and the
 /// number of bytes read.
 ///
-/// Based loosely on [`ReadVarint64FromArray`][1].
+/// Based loosely on [`ReadVarint64FromArray`][1] with a varint overflow check from
+/// [`ConsumeVarint`][2].
 ///
 /// ## Safety
 ///
@@ -74,6 +75,7 @@ where
 /// element in bytes is < `0x80`.
 ///
 /// [1]: https://github.com/google/protobuf/blob/3.3.x/src/google/protobuf/io/coded_stream.cc#L365-L406
+/// [2]: https://github.com/protocolbuffers/protobuf-go/blob/v1.27.1/encoding/protowire/wire.go#L358
 #[inline]
 fn decode_varint_slice(bytes: &[u8]) -> Result<(u64, usize), DecodeError> {
     // Fully unrolled varint decoding loop. Splitting into 32-bit pieces gives better performance.
@@ -146,16 +148,23 @@ fn decode_varint_slice(bytes: &[u8]) -> Result<(u64, usize), DecodeError> {
     part2 -= 0x80;
     b = unsafe { *bytes.get_unchecked(9) };
     part2 += u32::from(b) << 7;
-    if b < 0x80 {
+    // Check for u64::MAX overflow. See [`ConsumeVarint`][1] for details.
+    // [1]: https://github.com/protocolbuffers/protobuf-go/blob/v1.27.1/encoding/protowire/wire.go#L358
+    if b < 0x02 {
         return Ok((value + (u64::from(part2) << 56), 10));
     };
 
-    // We have overrun the maximum size of a varint (10 bytes). Assume the data is corrupt.
+    // We have overrun the maximum size of a varint (10 bytes) or the final byte caused an overflow.
+    // Assume the data is corrupt.
     Err(DecodeError::new("invalid varint"))
 }
 
 /// Decodes a LEB128-encoded variable length integer from the buffer, advancing the buffer as
 /// necessary.
+///
+/// Contains a varint overflow check from [`ConsumeVarint`][1].
+///
+/// [1]: https://github.com/protocolbuffers/protobuf-go/blob/v1.27.1/encoding/protowire/wire.go#L358
 #[inline(never)]
 #[cold]
 fn decode_varint_slow<B>(buf: &mut B) -> Result<u64, DecodeError>
@@ -167,7 +176,13 @@ where
         let byte = buf.get_u8();
         value |= u64::from(byte & 0x7F) << (count * 7);
         if byte <= 0x7F {
-            return Ok(value);
+            // Check for u64::MAX overflow. See [`ConsumeVarint`][1] for details.
+            // [1]: https://github.com/protocolbuffers/protobuf-go/blob/v1.27.1/encoding/protowire/wire.go#L358
+            if count == 9 && byte >= 0x02 {
+                return Err(DecodeError::new("invalid varint"));
+            } else {
+                return Ok(value);
+            }
         }
     }
 
@@ -1648,6 +1663,16 @@ mod test {
         );
     }
 
+    #[test]
+    fn varint_overflow() {
+        let mut u64_max_plus_one: &[u8] =
+            &[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x02];
+
+        decode_varint(&mut u64_max_plus_one).expect_err("decoding u64::MAX + 1 succeeded");
+        decode_varint_slow(&mut u64_max_plus_one)
+            .expect_err("slow decoding u64::MAX + 1 succeeded");
+    }
+
     /// This big bowl o' macro soup generates an encoding property test for each combination of map
     /// type, scalar map key, and value type.
     /// TODO: these tests take a long time to compile, can this be improved?