use std::fmt;
pub const MAX_BENCODE_DEPTH: usize = 32;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BencodeStructureError {
TooDeep { limit: usize },
Truncated,
UnbalancedEnd,
BadLength,
InvalidByte(u8),
Empty,
}
impl fmt::Display for BencodeStructureError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::TooDeep { limit } => {
write!(f, "bencode nesting deeper than the {} level limit", limit)
}
Self::Truncated => write!(f, "bencode value is truncated"),
Self::UnbalancedEnd => write!(f, "bencode 'e' with no open container"),
Self::BadLength => write!(f, "bencode length prefix is invalid or exceeds the input"),
Self::InvalidByte(b) => {
write!(f, "byte 0x{:02x} cannot begin a bencode value", b)
}
Self::Empty => write!(f, "bencode input is empty"),
}
}
}
impl std::error::Error for BencodeStructureError {}
pub fn check_bencode_structure(data: &[u8]) -> Result<(), BencodeStructureError> {
check_bencode_structure_with_limit(data, MAX_BENCODE_DEPTH)
}
pub fn check_bencode_structure_with_limit(
data: &[u8],
max_depth: usize,
) -> Result<(), BencodeStructureError> {
if data.is_empty() {
return Err(BencodeStructureError::Empty);
}
let mut i = 0usize;
let mut depth = 0usize;
loop {
let Some(&byte) = data.get(i) else {
return Err(BencodeStructureError::Truncated);
};
match byte {
b'd' | b'l' => {
depth += 1;
if depth > max_depth {
return Err(BencodeStructureError::TooDeep { limit: max_depth });
}
i += 1;
}
b'e' => {
if depth == 0 {
return Err(BencodeStructureError::UnbalancedEnd);
}
depth -= 1;
i += 1;
if depth == 0 {
return Ok(());
}
}
b'i' => {
i = scan_integer(data, i)?;
if depth == 0 {
return Ok(());
}
}
b'0'..=b'9' => {
i = scan_byte_string(data, i)?;
if depth == 0 {
return Ok(());
}
}
other => return Err(BencodeStructureError::InvalidByte(other)),
}
}
}
fn scan_integer(data: &[u8], start: usize) -> Result<usize, BencodeStructureError> {
debug_assert_eq!(data[start], b'i');
let mut i = start + 1;
let mut digits = 0usize;
if data.get(i) == Some(&b'-') {
i += 1;
}
while let Some(&b) = data.get(i) {
match b {
b'0'..=b'9' => {
digits += 1;
i += 1;
}
b'e' => {
if digits == 0 {
return Err(BencodeStructureError::BadLength);
}
return Ok(i + 1);
}
_ => return Err(BencodeStructureError::BadLength),
}
}
Err(BencodeStructureError::Truncated)
}
fn scan_byte_string(data: &[u8], start: usize) -> Result<usize, BencodeStructureError> {
let mut i = start;
let mut len: usize = 0;
let mut digits = 0usize;
while let Some(&b) = data.get(i) {
match b {
b'0'..=b'9' => {
len = len
.checked_mul(10)
.and_then(|v| v.checked_add((b - b'0') as usize))
.ok_or(BencodeStructureError::BadLength)?;
digits += 1;
i += 1;
if len > data.len() {
return Err(BencodeStructureError::BadLength);
}
}
b':' => {
if digits == 0 {
return Err(BencodeStructureError::BadLength);
}
let payload_start = i + 1;
let end = payload_start
.checked_add(len)
.ok_or(BencodeStructureError::BadLength)?;
if end > data.len() {
return Err(BencodeStructureError::Truncated);
}
return Ok(end);
}
_ => return Err(BencodeStructureError::BadLength),
}
}
Err(BencodeStructureError::Truncated)
}