diff --git a/src/reader.rs b/src/reader.rs index 28d9f73..aaee8f0 100644 --- a/src/reader.rs +++ b/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() -> 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 VarIntAsyncReader for AR { async fn read_varint_async(&mut self) -> Result { let mut buf = [0 as u8; 1]; - let mut p = VarIntProcessor::default(); + let mut p = VarIntProcessor::new::(); while !p.finished() { let read = self.read(&mut buf).await?; @@ -86,7 +90,7 @@ impl VarIntAsyncReader for AR { impl VarIntReader for R { fn read_varint(&mut self) -> Result { let mut buf = [0 as u8; 1]; - let mut p = VarIntProcessor::default(); + let mut p = VarIntProcessor::new::(); while !p.finished() { let read = self.read(&mut buf)?; diff --git a/src/varint.rs b/src/varint.rs index bb5c5a2..20069fc 100644 --- a/src/varint.rs +++ b/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 VarIntMaxSize for VI { + fn varint_max_size() -> usize { + (size_of::()*8+7)/7 + } +} + macro_rules! impl_varint { ($t:ty, unsigned) => { impl VarInt for $t { diff --git a/src/varint_tests.rs b/src/varint_tests.rs index 21d1b72..432e2a0 100644 --- a/src/varint_tests.rs +++ b/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 = (0x112233 as u64).encode_var_vec(); + assert_eq!(encoded.as_slice().read_varint::().unwrap_err().kind(), std::io::ErrorKind::InvalidData); + } }