diff options
| -rw-r--r-- | crates/tor-bytes/src/err.rs | 9 | ||||
| -rw-r--r-- | crates/tor-bytes/src/reader.rs | 91 |
2 files changed, 90 insertions, 10 deletions
diff --git a/crates/tor-bytes/src/err.rs b/crates/tor-bytes/src/err.rs index ed1672b27..3c015de5d 100644 --- a/crates/tor-bytes/src/err.rs +++ b/crates/tor-bytes/src/err.rs @@ -111,6 +111,15 @@ pub enum Error { #[error("Internal error")] #[deftly(never_eq)] // an internal error is equal to nothing, not even itself. Bug(#[from] tor_error::Bug), + /// Using for reader.rs::take_until_with_limit + /// If we reach the limit before finding the terminator, we return this error. + #[error("Reached limit of {limit} bytes while searching for terminator 0x{terminator:02x}")] + LimitExceeded { + /// Maximum number of bytes that were allowed to read + limit: usize, + /// The terminator byte we were looking for + terminator: u8, + }, } impl Error { diff --git a/crates/tor-bytes/src/reader.rs b/crates/tor-bytes/src/reader.rs index a1f1235e8..ae5db2751 100644 --- a/crates/tor-bytes/src/reader.rs +++ b/crates/tor-bytes/src/reader.rs @@ -2,7 +2,10 @@ use tor_error::{bad_api_usage, into_internal}; -use crate::{Error, Readable, Result}; +use crate::{ + Error::{self}, + Readable, Result, +}; use std::num::NonZeroUsize; /// A type for reading messages from a slice of bytes. @@ -295,15 +298,52 @@ impl<'a> Reader<'a> { /// # Result::Ok(()) /// ``` pub fn take_until(&mut self, term: u8) -> Result<&'a [u8]> { - let pos = - self.b[self.off..] - .iter() - .position(|b| *b == term) - .ok_or(self.incomplete_error( - // - 1.try_into().expect("1 == 0"), - ))?; - let result = self.take(pos)?; + self.take_until_with_limit(term, usize::MAX) + // Since take_until_with_limit used if max_len > self.remaining() + // So when max_len is usize::MAX, it will always be greater than self.remaining(), + // and thus we will always return the error from incomplete_error + } + /// Try to consume and return bytes from this buffer until we + /// encounter a terminating byte equal to `term`. + /// Or we have read `max_len` bytes without finding the terminator. + /// The maximum value that will be returned is max_len - 1, + /// because we need to remove the terminator. + /// If we didn't find the terminator within the first `max_len` bytes, + /// we will return an error. + /// + /// # Example + /// ```rust + /// use tor_bytes::{Reader,Result}; + /// let m = b"Hello\0wrld"; + /// let mut b = Reader::from_slice(m); + /// assert_eq!(b.take_until_with_limit(0, 10)?, b"Hello"); + /// assert_eq!(b.into_rest(), b"wrld"); + /// # Result::Ok(()) + /// ``` + pub fn take_until_with_limit(&mut self, term: u8, max_len: usize) -> Result<&'a [u8]> { + let limit: usize = std::cmp::min(max_len, self.remaining()); + let pos = match self.b[self.off..limit + self.off] + .iter() + .position(|b| *b == term) + { + Some(p) => p, + None => { + if max_len > self.remaining() { + // since user asked for more than we have, we should return Incomplete, not LimitExceeded + return Err(self.incomplete_error( + // + 1.try_into().expect("1 == 0"), + )); + } else { + return Err(Error::LimitExceeded { + limit, + terminator: term, + }); + } + } + }; + + let result: &[u8] = self.take(pos)?; self.advance(1)?; Ok(result) } @@ -680,6 +720,37 @@ mod tests { } #[test] + fn take_until_with_limit() { + let mut b = Reader::from_slice_for_test(&b"si vales valeo"[..]); + assert_eq!(b.take_until_with_limit(b' ', 10).unwrap(), &b"si"[..]); + assert_eq!(b.take_until_with_limit(b' ', 10).unwrap(), &b"vales"[..]); + assert_eq!( + b.take_until_with_limit(b' ', 100), + Err(Error::new_incomplete_for_test(1)), + ); + let mut b = Reader::from_slice_for_test(&b"Hello\0World"[..]); + assert_eq!( + b.take_until_with_limit(b'\0', 1), + Err(Error::LimitExceeded { + limit: 1, + terminator: b'\0' + }) + ); + + // Test the case where the terminator is exactly at the max_len + let mut b = Reader::from_slice_for_test(&b"si vales valeo"[..]); + assert_eq!(b.take_until_with_limit(b' ', 3).unwrap(), &b"si"[..]); + // Test the case where the terminator is exactly one more than the limit + assert_eq!( + b.take_until_with_limit(b' ', 5), + Err(Error::LimitExceeded { + limit: 5, + terminator: b' ' + }) + ); + } + + #[test] fn truncate_badly() { let mut b = Reader::from_slice_for_test(&b"abcdefg"[..]); b.truncate(1000); |