use tor_error::{bad_api_usage, into_internal};
use crate::{Error, Readable, Result};
use std::num::NonZeroUsize;
pub struct Reader<'a> {
b: &'a [u8],
off: usize,
completeness: Completeness,
}
#[derive(Copy, Clone, Debug)]
enum Completeness {
PossiblyIncomplete,
SupposedlyComplete,
}
impl<'a> Reader<'a> {
pub fn from_slice(slice: &'a [u8]) -> Self {
Reader {
b: slice,
off: 0,
completeness: Completeness::SupposedlyComplete,
}
}
pub fn from_possibly_incomplete_slice(slice: &'a [u8]) -> Self {
Reader {
b: slice,
off: 0,
completeness: Completeness::PossiblyIncomplete,
}
}
pub fn from_slice_for_test(slice: &'a [u8]) -> Self {
Self::from_possibly_incomplete_slice(slice)
}
pub fn from_bytes(b: &'a bytes::Bytes) -> Self {
Self::from_slice(b.as_ref())
}
pub fn total_len(&self) -> usize {
self.b.len()
}
pub fn remaining(&self) -> usize {
self.b.len() - self.off
}
pub fn into_rest(self) -> &'a [u8] {
&self.b[self.off..]
}
pub fn consumed(&self) -> usize {
self.off
}
pub fn advance(&mut self, n: usize) -> Result<()> {
self.peek(n)?;
self.off += n;
Ok(())
}
pub fn should_be_exhausted(&self) -> Result<()> {
if self.remaining() != 0 {
return Err(Error::ExtraneousBytes);
}
Ok(())
}
pub fn truncate(&mut self, n: usize) {
if n < self.remaining() {
self.b = &self.b[..self.off + n];
}
}
pub fn peek(&self, n: usize) -> Result<&'a [u8]> {
if let Some(deficit) = n
.checked_sub(self.remaining())
.and_then(|d| d.try_into().ok())
{
return Err(self.incomplete_error(deficit));
}
Ok(&self.b[self.off..(n + self.off)])
}
pub fn take(&mut self, n: usize) -> Result<&'a [u8]> {
let b = self.peek(n)?;
self.advance(n)?;
Ok(b)
}
pub fn take_into(&mut self, buf: &mut [u8]) -> Result<()> {
let n = buf.len();
let b = self.take(n)?;
buf.copy_from_slice(b);
Ok(())
}
pub fn take_u8(&mut self) -> Result<u8> {
let b = self.take(1)?;
Ok(b[0])
}
pub fn take_u16(&mut self) -> Result<u16> {
let b: [u8; 2] = self.extract()?;
let r = u16::from_be_bytes(b);
Ok(r)
}
pub fn take_u32(&mut self) -> Result<u32> {
let b: [u8; 4] = self.extract()?;
let r = u32::from_be_bytes(b);
Ok(r)
}
pub fn take_u64(&mut self) -> Result<u64> {
let b: [u8; 8] = self.extract()?;
let r = u64::from_be_bytes(b);
Ok(r)
}
pub fn take_u128(&mut self) -> Result<u128> {
let b: [u8; 16] = self.extract()?;
let r = u128::from_be_bytes(b);
Ok(r)
}
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.advance(1)?;
Ok(result)
}
pub fn take_rest(&mut self) -> &'a [u8] {
self.take(self.remaining())
.expect("taking remaining failed")
}
pub fn take_all_but(&mut self, n: usize) -> Result<&'a [u8]> {
match self.completeness {
Completeness::PossiblyIncomplete => {
return Err(Error::Bug(bad_api_usage!(
"Called take_all_but on a PossiblyIncomplete reader."
)));
}
Completeness::SupposedlyComplete => {}
}
let n_to_take = self.remaining().checked_sub(n).ok_or(Error::MissingData)?;
let result = self
.take(n_to_take)
.map_err(into_internal!("Subtraction misled us somehow"))?;
debug_assert_eq!(self.remaining(), n);
Ok(result)
}
pub fn extract<E: Readable>(&mut self) -> Result<E> {
let off_orig = self.off;
let result = E::take_from(self);
if result.is_err() {
self.off = off_orig;
}
result
}
pub fn extract_n<E: Readable>(&mut self, n: usize) -> Result<Vec<E>> {
let n_alloc = std::cmp::min(n, self.remaining());
let mut result = Vec::with_capacity(n_alloc);
let off_orig = self.off;
for _ in 0..n {
match E::take_from(self) {
Ok(item) => result.push(item),
Err(e) => {
self.off = off_orig;
return Err(e);
}
}
}
Ok(result)
}
pub fn read_nested_u8len<F, T>(&mut self, f: F) -> Result<T>
where
F: FnOnce(&mut Reader) -> Result<T>,
{
read_nested_generic::<u8, _, _>(self, f)
}
pub fn read_nested_u16len<F, T>(&mut self, f: F) -> Result<T>
where
F: FnOnce(&mut Reader) -> Result<T>,
{
read_nested_generic::<u16, _, _>(self, f)
}
pub fn read_nested_u32len<F, T>(&mut self, f: F) -> Result<T>
where
F: FnOnce(&mut Reader) -> Result<T>,
{
read_nested_generic::<u32, _, _>(self, f)
}
pub fn cursor(&self) -> Cursor<'a> {
Cursor {
pos: self.off,
_phantom: std::marker::PhantomData,
}
}
pub fn range(&self, start: Cursor<'a>, end: Cursor<'a>) -> &'a [u8] {
if start.pos <= end.pos && end.pos <= self.b.len() {
&self.b[start.pos..end.pos]
} else {
&self.b[..0]
}
}
pub fn incomplete_error(&self, deficit: NonZeroUsize) -> Error {
use Completeness as C;
use Error as E;
match self.completeness {
C::PossiblyIncomplete => E::Incomplete {
deficit: deficit.into(),
},
C::SupposedlyComplete => E::MissingData,
}
}
}
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
pub struct Cursor<'a> {
pos: usize,
_phantom: std::marker::PhantomData<&'a [u8]>,
}
fn read_nested_generic<L, F, T>(b: &mut Reader, f: F) -> Result<T>
where
F: FnOnce(&mut Reader) -> Result<T>,
L: Readable + Copy + Sized + TryInto<usize>,
{
let length: L = b.extract()?;
let length: usize = length.try_into().map_err(|_| Error::BadLengthValue)?;
let slice = b.take(length)?;
let mut inner = Reader::from_slice(slice);
let out = f(&mut inner)?;
inner.should_be_exhausted()?;
Ok(out)
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
#![allow(clippy::cognitive_complexity)]
use super::*;
#[test]
fn bytecursor_read_ok() {
let bytes = b"On a mountain halfway between Reno and Rome";
let mut bc = Reader::from_slice(&bytes[..]);
assert_eq!(bc.consumed(), 0);
assert_eq!(bc.remaining(), 43);
assert_eq!(bc.total_len(), 43);
assert_eq!(bc.take(3).unwrap(), &b"On "[..]);
assert_eq!(bc.consumed(), 3);
assert_eq!(bc.take_u16().unwrap(), 0x6120);
assert_eq!(bc.take_u8().unwrap(), 0x6d);
assert_eq!(bc.take_u64().unwrap(), 0x6f756e7461696e20);
assert_eq!(bc.take_u32().unwrap(), 0x68616c66);
assert_eq!(bc.consumed(), 18);
assert_eq!(bc.remaining(), 25);
assert_eq!(bc.total_len(), 43);
assert_eq!(bc.peek(7).unwrap(), &b"way bet"[..]);
assert_eq!(bc.consumed(), 18); assert_eq!(bc.remaining(), 25); assert_eq!(bc.total_len(), 43);
assert_eq!(bc.peek(7).unwrap(), &b"way bet"[..]);
assert_eq!(bc.consumed(), 18);
bc.advance(12).unwrap();
assert_eq!(bc.consumed(), 30);
assert_eq!(bc.remaining(), 13);
let rem = bc.into_rest();
assert_eq!(rem, &b"Reno and Rome"[..]);
let mut bc = Reader::from_slice(&bytes[..]);
bc.advance(22).unwrap();
assert_eq!(bc.remaining(), 21);
let rem = bc.take(21).unwrap();
assert_eq!(rem, &b"between Reno and Rome"[..]);
assert_eq!(bc.consumed(), 43);
assert_eq!(bc.remaining(), 0);
assert_eq!(bc.take(0).unwrap(), &b""[..]);
}
#[test]
fn read_u128() {
let bytes = bytes::Bytes::from(&b"irreproducibility?"[..]); let mut b = Reader::from_bytes(&bytes);
assert_eq!(b.take_u8().unwrap(), b'i');
assert_eq!(b.take_u128().unwrap(), 0x72726570726f6475636962696c697479);
assert_eq!(b.remaining(), 1);
}
#[test]
fn bytecursor_read_missing() {
let bytes = b"1234567";
let mut bc = Reader::from_slice_for_test(&bytes[..]);
assert_eq!(bc.consumed(), 0);
assert_eq!(bc.remaining(), 7);
assert_eq!(bc.total_len(), 7);
assert_eq!(bc.take_u64(), Err(Error::new_incomplete_for_test(1)));
assert_eq!(bc.take(8), Err(Error::new_incomplete_for_test(1)));
assert_eq!(bc.peek(8), Err(Error::new_incomplete_for_test(1)));
assert_eq!(bc.consumed(), 0);
assert_eq!(bc.remaining(), 7);
assert_eq!(bc.total_len(), 7);
assert_eq!(bc.take_u32().unwrap(), 0x31323334); assert_eq!(bc.take_u32(), Err(Error::new_incomplete_for_test(1)));
assert_eq!(bc.consumed(), 4);
assert_eq!(bc.remaining(), 3);
assert_eq!(bc.total_len(), 7);
assert_eq!(bc.take_u16().unwrap(), 0x3536); assert_eq!(bc.take_u16(), Err(Error::new_incomplete_for_test(1)));
assert_eq!(bc.consumed(), 6);
assert_eq!(bc.remaining(), 1);
assert_eq!(bc.total_len(), 7);
assert_eq!(bc.take_u8().unwrap(), 0x37); assert_eq!(bc.take_u8(), Err(Error::new_incomplete_for_test(1)));
assert_eq!(bc.consumed(), 7);
assert_eq!(bc.remaining(), 0);
assert_eq!(bc.total_len(), 7);
}
#[test]
fn advance_too_far() {
let bytes = b"12345";
let mut b = Reader::from_slice_for_test(&bytes[..]);
assert_eq!(b.remaining(), 5);
assert_eq!(b.advance(16), Err(Error::new_incomplete_for_test(11)));
assert_eq!(b.remaining(), 5);
assert_eq!(b.advance(5), Ok(()));
assert_eq!(b.remaining(), 0);
}
#[test]
fn truncate() {
let bytes = b"Hello universe!!!1!";
let mut b = Reader::from_slice_for_test(&bytes[..]);
assert_eq!(b.take(5).unwrap(), &b"Hello"[..]);
assert_eq!(b.remaining(), 14);
assert_eq!(b.consumed(), 5);
b.truncate(9);
assert_eq!(b.remaining(), 9);
assert_eq!(b.consumed(), 5);
assert_eq!(b.take_u8().unwrap(), 0x20);
assert_eq!(b.into_rest(), &b"universe"[..]);
}
#[test]
fn exhaust() {
let b = Reader::from_slice_for_test(&b""[..]);
assert_eq!(b.should_be_exhausted(), Ok(()));
let mut b = Reader::from_slice_for_test(&b"outis"[..]);
assert_eq!(b.should_be_exhausted(), Err(Error::ExtraneousBytes));
b.take(4).unwrap();
assert_eq!(b.should_be_exhausted(), Err(Error::ExtraneousBytes));
b.take(1).unwrap();
assert_eq!(b.should_be_exhausted(), Ok(()));
}
#[test]
fn take_rest() {
let mut b = Reader::from_slice_for_test(b"si vales valeo");
assert_eq!(b.take(3).unwrap(), b"si ");
assert_eq!(b.take_rest(), b"vales valeo");
assert_eq!(b.take_rest(), b"");
}
#[test]
fn take_until() {
let mut b = Reader::from_slice_for_test(&b"si vales valeo"[..]);
assert_eq!(b.take_until(b' ').unwrap(), &b"si"[..]);
assert_eq!(b.take_until(b' ').unwrap(), &b"vales"[..]);
assert_eq!(b.take_until(b' '), Err(Error::new_incomplete_for_test(1)));
}
#[test]
fn truncate_badly() {
let mut b = Reader::from_slice_for_test(&b"abcdefg"[..]);
b.truncate(1000);
assert_eq!(b.total_len(), 7);
assert_eq!(b.remaining(), 7);
}
#[test]
fn nested_good() {
let mut b = Reader::from_slice_for_test(b"abc\0\0\x04defghijkl");
assert_eq!(b.take(3).unwrap(), b"abc");
b.read_nested_u16len(|s| {
assert!(s.should_be_exhausted().is_ok());
Ok(())
})
.unwrap();
b.read_nested_u8len(|s| {
assert_eq!(s.take(4).unwrap(), b"defg");
assert!(s.should_be_exhausted().is_ok());
Ok(())
})
.unwrap();
assert_eq!(b.take(2).unwrap(), b"hi");
}
#[test]
fn nested_bad() {
let mut b = Reader::from_slice_for_test(b"................");
assert_eq!(
read_nested_generic::<u128, _, ()>(&mut b, |_| panic!())
.err()
.unwrap(),
Error::BadLengthValue
);
let mut b = Reader::from_slice_for_test(b"................");
assert_eq!(
b.read_nested_u32len::<_, ()>(|_| panic!()).err().unwrap(),
Error::new_incomplete_for_test(774778414 - (16 - 4))
);
}
#[test]
fn nested_inner_bad() {
let mut b = Reader::from_slice_for_test(&[1, 66]);
assert_eq!(
b.read_nested_u8len(|b| b.take_u32()),
Err(Error::MissingData),
);
}
#[test]
fn incomplete_slice() {
let mut b = Reader::from_possibly_incomplete_slice(&[]);
assert_eq!(b.take_u32(), Err(Error::new_incomplete_for_test(4)));
}
#[test]
fn extract() {
#[derive(Debug)]
struct LenEnc(Vec<u8>);
impl Readable for LenEnc {
fn take_from(b: &mut Reader<'_>) -> Result<Self> {
let length = b.take_u8()?;
let content = b.take(length as usize)?.into();
Ok(LenEnc(content))
}
}
let bytes = b"\x04this\x02is\x09sometimes\x01a\x06string!";
let mut b = Reader::from_slice_for_test(&bytes[..]);
let le: LenEnc = b.extract().unwrap();
assert_eq!(&le.0[..], &b"this"[..]);
let les: Vec<LenEnc> = b.extract_n(4).unwrap();
assert_eq!(&les[3].0[..], &b"string"[..]);
assert_eq!(b.remaining(), 1);
let le: Result<LenEnc> = b.extract();
assert_eq!(le.unwrap_err(), Error::new_incomplete_for_test(33));
assert_eq!(b.remaining(), 1);
let mut b = Reader::from_slice_for_test(&bytes[..]);
assert_eq!(b.remaining(), 28);
let les: Result<Vec<LenEnc>> = b.extract_n(10);
assert_eq!(les.unwrap_err(), Error::new_incomplete_for_test(33));
assert_eq!(b.remaining(), 28);
}
#[test]
fn cursor() -> Result<()> {
let alphabet = b"abcdefghijklmnopqrstuvwxyz";
let mut b = Reader::from_slice_for_test(&alphabet[..]);
let c1 = b.cursor();
let _ = b.take_u16()?;
let c2 = b.cursor();
let c2b = b.cursor();
b.advance(7)?;
let c3 = b.cursor();
assert_eq!(b.range(c1, c2), &b"ab"[..]);
assert_eq!(b.range(c2, c3), &b"cdefghi"[..]);
assert_eq!(b.range(c1, c3), &b"abcdefghi"[..]);
assert_eq!(b.range(c1, c1), &b""[..]);
assert_eq!(b.range(c3, c1), &b""[..]);
assert_eq!(c2, c2b);
assert!(c1 < c2);
assert!(c2 < c3);
Ok(())
}
#[test]
fn take_all_but() -> Result<()> {
let message = b"byte manipulation for fun and (non)-profit";
let mut b = Reader::from_slice(message);
assert_eq!(b.take_all_but(6)?, b"byte manipulation for fun and (non)-");
assert_eq!(b.into_rest(), b"profit");
let mut b = Reader::from_slice(message);
assert_eq!(b.take_all_but(message.len())?, b"");
assert_eq!(b.into_rest(), message);
let mut b = Reader::from_slice(message);
assert_eq!(b.take_all_but(0)?, message);
assert_eq!(b.into_rest(), b"");
let mut b = Reader::from_slice(message);
assert!(matches!(
b.take_all_but(message.len() + 1),
Err(Error::MissingData)
));
let mut b = Reader::from_possibly_incomplete_slice(message);
assert!(matches!(b.take_all_but(6), Err(Error::Bug(_))));
Ok(())
}
}