use core::fmt;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ByteErrorKind {
Truncated,
LengthOverflow,
BudgetExceeded,
OverlongModifiedUtf8,
IllegalZero,
InvalidModifiedUtf8,
MalformedSurrogate,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ByteError {
pub kind: ByteErrorKind,
pub offset: usize,
pub message: String,
}
impl ByteError {
pub(crate) fn new(kind: ByteErrorKind, offset: usize, message: impl Into<String>) -> Self {
Self {
kind,
offset,
message: message.into(),
}
}
}
impl fmt::Display for ByteError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(formatter, "{} at byte {}", self.message, self.offset)
}
}
impl std::error::Error for ByteError {}
#[derive(Clone, Debug)]
pub struct ByteReader<'a> {
bytes: &'a [u8],
position: usize,
origin: usize,
allocation_budget: usize,
}
impl<'a> ByteReader<'a> {
pub fn new(bytes: &'a [u8], allocation_budget: usize) -> Self {
Self {
bytes,
position: 0,
origin: 0,
allocation_budget,
}
}
fn with_origin(bytes: &'a [u8], allocation_budget: usize, origin: usize) -> Self {
Self {
bytes,
position: 0,
origin,
allocation_budget,
}
}
pub fn offset(&self) -> usize {
self.origin + self.position
}
pub fn remaining(&self) -> usize {
self.bytes.len() - self.position
}
pub fn allocation_budget(&self) -> usize {
self.allocation_budget
}
pub fn preflight_allocation(&self, amount: usize) -> Result<(), ByteError> {
if amount > self.allocation_budget {
return Err(ByteError::new(
ByteErrorKind::BudgetExceeded,
self.offset(),
format!(
"declared allocation {amount} exceeds budget {}",
self.allocation_budget
),
));
}
Ok(())
}
pub fn read_u1(&mut self) -> Result<u8, ByteError> {
Ok(self.take(1)?[0])
}
pub fn read_u2(&mut self) -> Result<u16, ByteError> {
let value: [u8; 2] = self.take(2)?.try_into().expect("exact length");
Ok(u16::from_be_bytes(value))
}
pub fn read_u4(&mut self) -> Result<u32, ByteError> {
let value: [u8; 4] = self.take(4)?.try_into().expect("exact length");
Ok(u32::from_be_bytes(value))
}
pub fn take(&mut self, length: usize) -> Result<&'a [u8], ByteError> {
let start = self.position;
let end = start.checked_add(length).ok_or_else(|| {
ByteError::new(
ByteErrorKind::LengthOverflow,
self.offset(),
"byte length overflow",
)
})?;
let result = self.bytes.get(start..end).ok_or_else(|| {
ByteError::new(
ByteErrorKind::Truncated,
self.origin + self.bytes.len(),
format!("needed {length} bytes, only {} remain", self.remaining()),
)
})?;
self.position = end;
Ok(result)
}
pub fn sub_reader(&mut self, length: usize) -> Result<ByteReader<'a>, ByteError> {
let origin = self.offset();
let bytes = self.take(length)?;
Ok(Self::with_origin(bytes, self.allocation_budget, origin))
}
}
#[derive(Clone, Debug)]
pub struct ByteWriter {
bytes: Vec<u8>,
budget: usize,
}
impl ByteWriter {
pub fn new(budget: usize) -> Self {
Self {
bytes: Vec::new(),
budget,
}
}
fn reserve_for(&mut self, additional: usize) -> Result<(), ByteError> {
let target = self.bytes.len().checked_add(additional).ok_or_else(|| {
ByteError::new(
ByteErrorKind::LengthOverflow,
self.bytes.len(),
"output length overflow",
)
})?;
if target > self.budget {
return Err(ByteError::new(
ByteErrorKind::BudgetExceeded,
self.bytes.len(),
format!("output length {target} exceeds budget {}", self.budget),
));
}
self.bytes.try_reserve_exact(additional).map_err(|error| {
ByteError::new(
ByteErrorKind::BudgetExceeded,
self.bytes.len(),
format!("output allocation failed: {error}"),
)
})
}
pub fn write_u1(&mut self, value: u8) -> Result<(), ByteError> {
self.write_bytes(&[value])
}
pub fn write_u2(&mut self, value: u16) -> Result<(), ByteError> {
self.write_bytes(&value.to_be_bytes())
}
pub fn write_u4(&mut self, value: u32) -> Result<(), ByteError> {
self.write_bytes(&value.to_be_bytes())
}
pub fn write_bytes(&mut self, value: &[u8]) -> Result<(), ByteError> {
self.reserve_for(value.len())?;
self.bytes.extend_from_slice(value);
Ok(())
}
pub fn as_slice(&self) -> &[u8] {
&self.bytes
}
pub fn into_bytes(self) -> Vec<u8> {
self.bytes
}
}