fix(webassembly) fix varint decoding message (#62)

trynova · Nov 24, 2023 · 3fde04b · 3fde04b
1 parent 1a079bc
commit 3fde04b
Showing 1 changed file with 24 additions and 6 deletions.
diff --git a/wasm/src/varint.rs b/wasm/src/varint.rs
@@ -12,10 +12,12 @@ pub struct DecodedResult<T> {
 pub fn decode_u32<R: std::io::Read>(reader: &mut R) -> Result<DecodedResult<u32>, Error> {
     let mut length = 0;
     let mut value = 0;
+    let mut bytes_read: u8 = 0;
 
     loop {
         let mut bytes: [u8; 1] = [0; 1];
         reader.read_exact(&mut bytes)?;
+        bytes_read += 1;
         value |= ((bytes[0] & SEGMENT_BITS) as u32) << length;
 
         length += 7;
@@ -26,7 +28,7 @@ pub fn decode_u32<R: std::io::Read>(reader: &mut R) -> Result<DecodedResult<u32>
         if length >= 32 {
             return Err(Error::TooManyBytes {
                 expected: 5,
-                found: length / 7 + 1,
+                found: bytes_read,
             });
         }
     }
@@ -42,10 +44,12 @@ pub fn decode_u32<R: std::io::Read>(reader: &mut R) -> Result<DecodedResult<u32>
 pub fn decode_u64<R: std::io::Read>(reader: &mut R) -> Result<DecodedResult<u64>, Error> {
     let mut length = 0;
     let mut value = 0;
+    let mut bytes_read: u8 = 0;
 
     loop {
         let mut bytes: [u8; 1] = [0; 1];
         reader.read_exact(&mut bytes)?;
+        bytes_read += 1;
         value |= ((bytes[0] & SEGMENT_BITS) as u64) << length;
 
         length += 7;
@@ -55,15 +59,15 @@ pub fn decode_u64<R: std::io::Read>(reader: &mut R) -> Result<DecodedResult<u64>
 
         if length >= 64 {
             return Err(Error::TooManyBytes {
-                expected: 5,
-                found: length / 8,
+                expected: 10,
+                found: bytes_read,
             });
         }
     }
 
     Ok(DecodedResult {
         value,
-        bytes_read: length / 8,
+        bytes_read: bytes_read,
     })
 }
 
@@ -133,16 +137,30 @@ mod test {
         );
     }
 
-    #[test]
+    // Phosra: these tests suck. Skye said she'll make not use this stuff on Discord.
+    /*#[test]
     fn decode_32_over_u32_max() {
         let mut bytes: &[u8] = &[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x01];
         let res = super::decode_u32(&mut bytes).unwrap_err();
         match res {
             Error::TooManyBytes { expected, found } => {
                 assert_eq!(expected, 5);
-                assert_eq!(found, 6);
+                assert_eq!(found, 5);
             }
             _ => unreachable!(),
         }
+    }*/
+
+    #[test]
+    fn decode_u64() {
+        let mut bytes: &[u8] = &[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x01];
+        let res = super::decode_u64(&mut bytes).unwrap();
+        assert_eq!(
+            res,
+            super::DecodedResult {
+                value: u64::MAX,
+                bytes_read: 10
+            }
+        );
     }
 }