use std::collections::BTreeMap;
use super::{Error, ErrorKind, Limits, Node, Value};
pub fn decode(input: &[u8], limits: Limits) -> Result<Node<'_>, Error> {
let (node, consumed) = decode_prefix(input, limits)?;
if consumed != input.len() {
return Err(Error::new(consumed, ErrorKind::TrailingBytes));
}
Ok(node)
}
pub fn decode_prefix(input: &[u8], limits: Limits) -> Result<(Node<'_>, usize), Error> {
let limits = limits.validate()?;
if input.len() > limits.max_input_len {
return Err(Error::new(0, ErrorKind::InputLimit));
}
let mut parser = Parser {
input,
limits,
position: 0,
nodes: 0,
};
let node = parser.value(0)?;
Ok((node, parser.position))
}
struct Parser<'a> {
input: &'a [u8],
limits: Limits,
position: usize,
nodes: usize,
}
impl<'a> Parser<'a> {
fn error(&self, kind: ErrorKind) -> Error {
Error::new(self.position, kind)
}
fn peek(&self) -> Result<u8, Error> {
self.input
.get(self.position)
.copied()
.ok_or_else(|| self.error(ErrorKind::UnexpectedEnd))
}
fn count_node(&mut self) -> Result<(), Error> {
if self.nodes >= self.limits.max_nodes {
return Err(self.error(ErrorKind::NodeLimit));
}
self.nodes += 1;
Ok(())
}
fn value(&mut self, depth: usize) -> Result<Node<'a>, Error> {
self.count_node()?;
let start = self.position;
let value = match self.peek()? {
b'i' => Value::Integer(self.integer()?),
b'0'..=b'9' => Value::Bytes(self.string()?),
b'l' | b'd' if depth >= self.limits.max_depth => {
return Err(self.error(ErrorKind::DepthLimit));
}
b'l' => {
self.position += 1;
let mut entries = Vec::new();
while self.peek()? != b'e' {
entries.push(self.value(depth + 1)?);
}
self.position += 1;
Value::List(entries)
}
b'd' => {
self.position += 1;
let mut entries = BTreeMap::new();
while self.peek()? != b'e' {
self.count_node()?;
if !self.peek()?.is_ascii_digit() {
return Err(self.error(ErrorKind::DictionaryKey));
}
let key = self.string()?;
let value = self.value(depth + 1)?;
entries.insert(key, value);
}
self.position += 1;
Value::Dictionary(entries)
}
_ => return Err(self.error(ErrorKind::InvalidToken)),
};
Ok(Node {
value,
raw: Some(&self.input[start..self.position]),
})
}
fn integer(&mut self) -> Result<i64, Error> {
self.position += 1; let start = self.position;
let negative = self.peek()? == b'-';
if negative {
self.position += 1;
}
let digits = self.position;
while self.peek()?.is_ascii_digit() {
self.position += 1;
}
if self.peek()? != b'e' || self.position == digits {
return Err(self.error(ErrorKind::InvalidInteger));
}
if self.input[digits] == b'0' && (negative || self.position - digits > 1) {
return Err(Error::new(start, ErrorKind::InvalidInteger));
}
let text = std::str::from_utf8(&self.input[start..self.position])
.map_err(|_| self.error(ErrorKind::InvalidInteger))?;
let value = text
.parse::<i64>()
.map_err(|_| Error::new(start, ErrorKind::IntegerOverflow))?;
self.position += 1; Ok(value)
}
fn string(&mut self) -> Result<&'a [u8], Error> {
let start = self.position;
while self.peek()?.is_ascii_digit() {
self.position += 1;
}
if self.position == start || self.peek()? != b':' {
return Err(self.error(ErrorKind::InvalidStringLength));
}
let text = std::str::from_utf8(&self.input[start..self.position])
.map_err(|_| self.error(ErrorKind::InvalidStringLength))?;
let length = text
.parse::<u32>()
.map_err(|_| Error::new(start, ErrorKind::StringLengthOverflow))?
as usize;
if length > self.limits.max_string_len {
return Err(Error::new(start, ErrorKind::StringLimit));
}
self.position += 1; let end = self
.position
.checked_add(length)
.filter(|end| *end <= self.input.len())
.ok_or_else(|| self.error(ErrorKind::UnexpectedEnd))?;
let bytes = &self.input[self.position..end];
self.position = end;
Ok(bytes)
}
}