use crate::decompressor::error::{DecodeError, InvalidDataKind};
use crate::shared::bitpack;
use crate::shared::varint::{self, VarintError};
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) enum Parse {
Truncated,
Failed(DecodeError),
}
impl From<DecodeError> for Parse {
fn from(error: DecodeError) -> Self {
Self::Failed(error)
}
}
#[derive(Debug)]
pub(crate) struct Reader<'a> {
bytes: &'a [u8],
pos: usize,
kind: InvalidDataKind,
}
impl<'a> Reader<'a> {
pub(crate) const fn new(bytes: &'a [u8], kind: InvalidDataKind) -> Self {
Self {
bytes,
pos: 0,
kind,
}
}
pub(crate) const fn pos(&self) -> usize {
self.pos
}
pub(crate) fn consumed(&self) -> &'a [u8] {
&self.bytes[..self.pos]
}
pub(crate) const fn invalid(&self) -> Parse {
Parse::Failed(DecodeError::invalid(self.kind))
}
pub(crate) fn byte(&mut self) -> Result<u8, Parse> {
let byte = *self.bytes.get(self.pos).ok_or(Parse::Truncated)?;
self.pos += 1;
Ok(byte)
}
pub(crate) fn take(&mut self, len: usize) -> Result<&'a [u8], Parse> {
let end = self.pos.checked_add(len).ok_or(Parse::Truncated)?;
let bytes = self.bytes.get(self.pos..end).ok_or(Parse::Truncated)?;
self.pos = end;
Ok(bytes)
}
pub(crate) fn varint(&mut self) -> Result<u64, Parse> {
match varint::decode(&self.bytes[self.pos..]) {
Ok((value, len)) => {
self.pos += len;
Ok(value)
}
Err(VarintError::Truncated) => Err(Parse::Truncated),
Err(VarintError::Overlong) => Err(self.invalid()),
}
}
pub(crate) fn packed(
&mut self,
count: usize,
bits: u32,
sink: impl FnMut(u32),
) -> Result<(), Parse> {
let len =
bitpack::unpack(&self.bytes[self.pos..], count, bits, sink).ok_or(Parse::Truncated)?;
self.pos += len;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn running_out_is_truncation_and_overlong_varints_are_invalid() {
let mut reader = Reader::new(&[0x81], InvalidDataKind::ChunkHeader);
assert_eq!(reader.varint(), Err(Parse::Truncated));
assert_eq!(reader.take(2), Err(Parse::Truncated));
assert_eq!(reader.byte(), Ok(0x81));
assert_eq!(reader.byte(), Err(Parse::Truncated));
assert_eq!(reader.consumed(), [0x81]);
let overlong = [0xffu8; 10];
let mut reader = Reader::new(&overlong, InvalidDataKind::FrameHeader);
assert_eq!(
reader.varint(),
Err(Parse::Failed(DecodeError::invalid(
InvalidDataKind::FrameHeader
)))
);
let mut reader = Reader::new(&[0b101], InvalidDataKind::ChunkHeader);
let mut seen = alloc::vec::Vec::new();
reader.packed(3, 1, |v| seen.push(v)).unwrap();
assert_eq!((seen, reader.pos()), (alloc::vec![1, 0, 1], 1));
assert_eq!(reader.packed(9, 1, drop), Err(Parse::Truncated));
assert_eq!(
Parse::from(DecodeError::InvalidState),
Parse::Failed(DecodeError::InvalidState)
);
}
}