Better size limit for #21

src/reader.rs

@@ -2,7 +2,7 @@ use std::io;
 use std::io::{Read, Result};
 
 use crate::fixed::FixedInt;
-use crate::varint::{VarInt, MSB};
+use crate::varint::{VarInt, VarIntMaxSize, MSB};
 
 #[cfg(feature = "tokio_async")]
 use tokio::io::{AsyncReadExt, AsyncRead};
@@ -34,12 +34,16 @@ pub trait VarIntAsyncReader {
 #[derive(Default)]
 pub struct VarIntProcessor {
     buf: [u8; 10],
+    maxsize: usize,
     i: usize,
 }
 
 impl VarIntProcessor {
+    fn new<VI: VarIntMaxSize>() -> VarIntProcessor {
+        VarIntProcessor { maxsize: VI::varint_max_size(), ..VarIntProcessor::default() }
+    }
     fn push(&mut self, b: u8) -> Result<()> {
-        if self.i >= 10 {
+        if self.i >= self.maxsize {
             return Err(io::Error::new(
                 io::ErrorKind::InvalidData,
                 "Unterminated varint",
@@ -62,7 +66,7 @@ impl VarIntProcessor {
 impl<AR: AsyncRead + Unpin + Send> VarIntAsyncReader for AR {
     async fn read_varint_async<VI: VarInt>(&mut self) -> Result<VI> {
         let mut buf = [0 as u8; 1];
-        let mut p = VarIntProcessor::default();
+        let mut p = VarIntProcessor::new::<VI>();
 
         while !p.finished() {
             let read = self.read(&mut buf).await?;
@@ -86,7 +90,7 @@ impl<AR: AsyncRead + Unpin + Send> VarIntAsyncReader for AR {
 impl<R: Read> VarIntReader for R {
     fn read_varint<VI: VarInt>(&mut self) -> Result<VI> {
         let mut buf = [0 as u8; 1];
-        let mut p = VarIntProcessor::default();
+        let mut p = VarIntProcessor::new::<VI>();
 
         while !p.finished() {
             let read = self.read(&mut buf)?;

src/varint.rs

@@ -1,3 +1,5 @@
+use std::mem::size_of;
+
 /// Most-significant byte, == 0x80
 pub const MSB: u8 = 0b1000_0000;
 /// All bits except for the most significant. Can be used as bitmask to drop the most-signficant
@@ -64,6 +66,16 @@ fn zigzag_decode(from: u64) -> i64 {
     ((from >> 1) ^ (-((from & 1) as i64)) as u64) as i64
 }
 
+pub(crate) trait VarIntMaxSize {
+    fn varint_max_size() -> usize;
+}
+
+impl<VI: VarInt> VarIntMaxSize for VI {
+    fn varint_max_size() -> usize {
+        (size_of::<VI>()*8+7)/7
+    }
+}
+
 macro_rules! impl_varint {
     ($t:ty, unsigned) => {
         impl VarInt for $t {

src/varint_tests.rs

@@ -200,4 +200,10 @@ mod tests {
         encoded.push(0x00);
         assert_eq!(i64::decode_var(&encoded[..]), None);
     }
+
+    #[test]
+    fn test_regression_22() {
+        let mut encoded: Vec<u8> = (0x112233 as u64).encode_var_vec();
+        assert_eq!(encoded.as_slice().read_varint::<i8>().unwrap_err().kind(), std::io::ErrorKind::InvalidData);
+    }
 }