From d9423c924600a834bdbcd55706a9e28e764a4b71 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20Sinan=20A=C4=9Facan?= <omeragacan@gmail.com>
Date: Mon, 28 Jun 2021 13:49:40 +0300
Subject: [PATCH] Fix sleb128_decode overflow checking

Fixes #2625

New tests added for checking edge cases + overflows.
---
 rts/motoko-rts-tests/src/leb128.rs | 55 +++++++++++++++++++++++++++++-
 rts/motoko-rts/src/leb128.rs       | 31 +++++++++++------
 2 files changed, 75 insertions(+), 11 deletions(-)

diff --git a/rts/motoko-rts-tests/src/leb128.rs b/rts/motoko-rts-tests/src/leb128.rs
index 07cd0f84405..af827fd4145 100644
--- a/rts/motoko-rts-tests/src/leb128.rs
+++ b/rts/motoko-rts-tests/src/leb128.rs
@@ -1,12 +1,47 @@
 use motoko_rts::buf::Buf;
-use motoko_rts::leb128::{leb128_decode, leb128_encode, sleb128_decode, sleb128_encode};
+use motoko_rts::leb128::{
+    leb128_decode, leb128_decode_checked, leb128_encode, sleb128_decode, sleb128_decode_checked,
+    sleb128_encode,
+};
 
 use quickcheck::{quickcheck, TestResult};
 
 pub unsafe fn test() {
     println!("Testing (s)leb128 encode-decode roundtrip ...");
+
+    assert!(!roundtrip_signed(i32::MIN).is_failure());
+    assert!(!roundtrip_signed(i32::MAX).is_failure());
+
+    assert!(!roundtrip_unsigned(u32::MIN).is_failure());
+    assert!(!roundtrip_unsigned(u32::MAX).is_failure());
+
     quickcheck(roundtrip_signed as fn(i32) -> TestResult);
     quickcheck(roundtrip_unsigned as fn(u32) -> TestResult);
+
+    // Check overflows
+    check_signed_decode_overflow(&[
+        0b1111_1111,
+        0b1111_1111,
+        0b1111_1111,
+        0b1111_1111,
+        0b0111_0111,
+    ]); // i32::MIN - 1
+
+    check_signed_decode_overflow(&[
+        0b1000_0000,
+        0b1000_0000,
+        0b1000_0000,
+        0b1000_0000,
+        0b0000_1000,
+    ]); // i32::MAX + 1
+
+    check_unsigned_decode_overflow(&[
+        0b1000_0000,
+        0b1000_0000,
+        0b1000_0000,
+        0b1000_0000,
+        0b0001_0000,
+    ]); // u32::MAX + 1
 }
 
 fn roundtrip_signed(val: i32) -> TestResult {
@@ -36,3 +71,21 @@ fn roundtrip_unsigned(val: u32) -> TestResult {
         TestResult::from_bool(leb128_decode(&mut buf_) == val)
     }
 }
+
+unsafe fn check_signed_decode_overflow(buf: &[u8]) {
+    let mut buf_ = Buf {
+        ptr: buf.as_ptr() as *mut _,
+        end: buf.as_ptr().add(buf.len()) as *mut _,
+    };
+
+    assert_eq!(sleb128_decode_checked(&mut buf_), None);
+}
+
+unsafe fn check_unsigned_decode_overflow(buf: &[u8]) {
+    let mut buf_ = Buf {
+        ptr: buf.as_ptr() as *mut _,
+        end: buf.as_ptr().add(buf.len()) as *mut _,
+    };
+
+    assert_eq!(leb128_decode_checked(&mut buf_), None);
+}
diff --git a/rts/motoko-rts/src/leb128.rs b/rts/motoko-rts/src/leb128.rs
index f5828eee35b..7df17efa77d 100644
--- a/rts/motoko-rts/src/leb128.rs
+++ b/rts/motoko-rts/src/leb128.rs
@@ -35,6 +35,11 @@ pub unsafe extern "C" fn sleb128_encode(mut val: i32, mut buf: *mut u8) {
 
 #[no_mangle]
 pub unsafe extern "C" fn leb128_decode(buf: *mut Buf) -> u32 {
+    leb128_decode_checked(buf).expect("leb128_decode: overflow")
+}
+
+/// Returns `None` on overflow
+pub unsafe fn leb128_decode_checked(buf: *mut Buf) -> Option<u32> {
     let mut result = 0;
     let mut shift = 0;
 
@@ -46,21 +51,26 @@ pub unsafe extern "C" fn leb128_decode(buf: *mut Buf) -> u32 {
         // The 5th byte needs to be the last, and it must contribute at most 4 bits, otherwise we
         // have an overflow
         if shift == 28 && (byte & 0b1111_0000) != 0 {
-            panic!("leb128_decode: overflow");
+            return None;
         }
 
+        shift += 7;
+
         if byte & 0b1000_0000 == 0 {
             break;
         }
-
-        shift += 7;
     }
 
-    result
+    Some(result)
 }
 
 #[no_mangle]
 pub unsafe extern "C" fn sleb128_decode(buf: *mut Buf) -> i32 {
+    sleb128_decode_checked(buf).expect("sleb128_decode: overflow")
+}
+
+/// Returns `None` on overflow
+pub unsafe fn sleb128_decode_checked(buf: *mut Buf) -> Option<i32> {
     let mut result = 0;
     let mut shift = 0;
 
@@ -68,14 +78,15 @@ pub unsafe extern "C" fn sleb128_decode(buf: *mut Buf) -> i32 {
         let byte = read_byte(buf);
 
         result |= ((byte & 0b0111_1111) as i32) << shift;
-        shift += 7;
 
-        // The 5th byte needs to be the last, and it must contribute at most 4 bits, otherwise we
-        // have an overflow
-        if shift == 28 && (byte & 0b1111_0000 != 0 || byte & 0b1110_0000 != 0) {
-            panic!("sleb128_decode: overflow");
+        // Overflow check ported from Wasm reference implementation:
+        // https://github.com/WebAssembly/spec/blob/f9770eb75117cac0c878feaa5eaf4a4d9dda61f5/interpreter/binary/decode.ml#L89-L98
+        if shift == 28 && (byte & 0b0111_1000 != 0 && byte & 0b0111_1000 != 0b0111_1000) {
+            return None;
         }
 
+        shift += 7;
+
         if byte & 0b1000_0000 == 0 {
             break byte;
         }
@@ -86,5 +97,5 @@ pub unsafe extern "C" fn sleb128_decode(buf: *mut Buf) -> i32 {
         result |= !0 << shift;
     }
 
-    result
+    Some(result)
 }