#![forbid(unsafe_code)]
use std::collections::BTreeMap;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Bencode {
Int(i64),
Bytes(Vec<u8>),
List(Vec<Bencode>),
Dict(BTreeMap<Vec<u8>, Bencode>),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Error {
UnexpectedEnd,
UnexpectedByte { offset: usize, byte: u8 },
InvalidInteger { offset: usize },
IntegerOverflow { offset: usize },
InvalidLength { offset: usize },
UnsortedOrDuplicateKey { offset: usize },
NonStringKey { offset: usize },
TrailingData { offset: usize },
}
impl std::fmt::Display for Error {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Error::UnexpectedEnd => write!(f, "unexpected end of bencode stream"),
Error::UnexpectedByte { offset, byte } => {
write!(f, "unexpected byte {byte:#04x} at offset {offset}")
}
Error::InvalidInteger { offset } => write!(f, "invalid integer at offset {offset}"),
Error::IntegerOverflow { offset } => write!(f, "integer overflows i64 at offset {offset}"),
Error::InvalidLength { offset } => write!(f, "invalid byte-string length at offset {offset}"),
Error::UnsortedOrDuplicateKey { offset } => {
write!(f, "dict keys not strictly ascending at offset {offset}")
}
Error::NonStringKey { offset } => write!(f, "non-string dict key at offset {offset}"),
Error::TrailingData { offset } => write!(f, "trailing data after value at offset {offset}"),
}
}
}
impl std::error::Error for Error {}
pub fn parse(input: &[u8]) -> Result<Bencode, Error> {
let (value, rest) = parse_prefix(input)?;
if rest.is_empty() {
Ok(value)
} else {
Err(Error::TrailingData { offset: input.len() - rest.len() })
}
}
pub fn parse_prefix(input: &[u8]) -> Result<(Bencode, &[u8]), Error> {
Parser { full: input, pos: 0 }.value().map(|v| {
(v.0, &input[v.1..])
})
}
struct Parser<'a> {
full: &'a [u8],
pos: usize,
}
impl Parser<'_> {
fn value(mut self) -> Result<(Bencode, usize), Error> {
let v = self.parse_value()?;
Ok((v, self.pos))
}
fn peek(&self) -> Result<u8, Error> {
self.full.get(self.pos).copied().ok_or(Error::UnexpectedEnd)
}
fn bump(&mut self) -> Result<u8, Error> {
let b = self.peek()?;
self.pos += 1;
Ok(b)
}
fn expect(&mut self, byte: u8) -> Result<(), Error> {
let at = self.pos;
if self.bump()? == byte {
Ok(())
} else {
Err(Error::UnexpectedByte { offset: at, byte: self.full[at] })
}
}
fn parse_value(&mut self) -> Result<Bencode, Error> {
match self.peek()? {
b'i' => self.parse_int(),
b'l' => self.parse_list(),
b'd' => self.parse_dict(),
b'0'..=b'9' => self.parse_bytes().map(Bencode::Bytes),
byte => Err(Error::UnexpectedByte { offset: self.pos, byte }),
}
}
fn parse_int(&mut self) -> Result<Bencode, Error> {
let start = self.pos;
self.expect(b'i')?;
let body_start = self.pos;
let mut end = self.pos;
while self.full.get(end).is_some_and(|&b| b != b'e') {
end += 1;
}
if end >= self.full.len() {
return Err(Error::UnexpectedEnd);
}
let body = &self.full[body_start..end];
validate_integer_body(body).map_err(|()| Error::InvalidInteger { offset: start })?;
let s = std::str::from_utf8(body).map_err(|_| Error::InvalidInteger { offset: start })?;
let n = s.parse::<i64>().map_err(|_| Error::IntegerOverflow { offset: start })?;
self.pos = end + 1; Ok(Bencode::Int(n))
}
fn parse_bytes(&mut self) -> Result<Vec<u8>, Error> {
let start = self.pos;
let mut end = self.pos;
while self.full.get(end).is_some_and(u8::is_ascii_digit) {
end += 1;
}
let digits = &self.full[start..end];
if digits.is_empty()
|| (digits.len() > 1 && digits[0] == b'0')
{
return Err(Error::InvalidLength { offset: start });
}
let len: usize = std::str::from_utf8(digits)
.ok()
.and_then(|s| s.parse().ok())
.ok_or(Error::InvalidLength { offset: start })?;
if self.full.get(end) != Some(&b':') {
return Err(Error::InvalidLength { offset: end });
}
let data_start = end + 1;
let data_end = data_start.checked_add(len).ok_or(Error::InvalidLength { offset: start })?;
if data_end > self.full.len() {
return Err(Error::UnexpectedEnd);
}
self.pos = data_end;
Ok(self.full[data_start..data_end].to_vec())
}
fn parse_list(&mut self) -> Result<Bencode, Error> {
self.expect(b'l')?;
let mut items = Vec::new();
while self.peek()? != b'e' {
items.push(self.parse_value()?);
}
self.pos += 1; Ok(Bencode::List(items))
}
fn parse_dict(&mut self) -> Result<Bencode, Error> {
self.expect(b'd')?;
let mut map: BTreeMap<Vec<u8>, Bencode> = BTreeMap::new();
let mut last_key: Option<Vec<u8>> = None;
while self.peek()? != b'e' {
let key_offset = self.pos;
if !self.peek()?.is_ascii_digit() {
return Err(Error::NonStringKey { offset: key_offset });
}
let key = self.parse_bytes()?;
if let Some(prev) = &last_key {
if key <= *prev {
return Err(Error::UnsortedOrDuplicateKey { offset: key_offset });
}
}
let value = self.parse_value()?;
last_key = Some(key.clone());
map.insert(key, value);
}
self.pos += 1; Ok(Bencode::Dict(map))
}
}
fn validate_integer_body(body: &[u8]) -> Result<(), ()> {
match body {
[] => Err(()),
[b'0'] => Ok(()),
[b'-', b'0', ..] => Err(()), [b'0', _, ..] => Err(()), _ => {
let digits = body.strip_prefix(b"-").unwrap_or(body);
if digits.is_empty() || !digits.iter().all(u8::is_ascii_digit) {
Err(())
} else {
Ok(())
}
}
}
}
impl Bencode {
#[must_use]
pub fn to_bytes(&self) -> Vec<u8> {
let mut out = Vec::new();
self.encode_into(&mut out);
out
}
fn encode_into(&self, out: &mut Vec<u8>) {
match self {
Bencode::Int(n) => {
out.push(b'i');
push_i64(out, *n);
out.push(b'e');
}
Bencode::Bytes(bytes) => {
push_usize(out, bytes.len());
out.push(b':');
out.extend_from_slice(bytes);
}
Bencode::List(items) => {
out.push(b'l');
for item in items {
item.encode_into(out);
}
out.push(b'e');
}
Bencode::Dict(map) => {
out.push(b'd');
for (key, value) in map {
push_usize(out, key.len());
out.push(b':');
out.extend_from_slice(key);
value.encode_into(out);
}
out.push(b'e');
}
}
}
#[must_use]
pub fn as_int(&self) -> Option<i64> {
match self {
Bencode::Int(n) => Some(*n),
_ => None,
}
}
#[must_use]
pub fn as_bytes(&self) -> Option<&[u8]> {
match self {
Bencode::Bytes(b) => Some(b),
_ => None,
}
}
#[must_use]
pub fn as_str(&self) -> Option<&str> {
self.as_bytes().and_then(|b| std::str::from_utf8(b).ok())
}
#[must_use]
pub fn as_list(&self) -> Option<&[Bencode]> {
match self {
Bencode::List(items) => Some(items),
_ => None,
}
}
#[must_use]
pub fn as_dict(&self) -> Option<&BTreeMap<Vec<u8>, Bencode>> {
match self {
Bencode::Dict(map) => Some(map),
_ => None,
}
}
#[must_use]
pub fn get(&self, key: &[u8]) -> Option<&Bencode> {
self.as_dict().and_then(|m| m.get(key))
}
}
fn push_i64(out: &mut Vec<u8>, n: i64) {
if n < 0 {
out.push(b'-');
push_u64(out, (n as i128).unsigned_abs() as u64);
} else {
push_u64(out, n as u64);
}
}
fn push_usize(out: &mut Vec<u8>, n: usize) {
push_u64(out, n as u64);
}
fn push_u64(out: &mut Vec<u8>, n: u64) {
if n == 0 {
out.push(b'0');
return;
}
let mut buf = [0u8; 20];
let mut i = buf.len();
let mut v = n;
while v > 0 {
i -= 1;
buf[i] = b'0' + u8::try_from(v % 10).expect("digit 0..=9");
v /= 10;
}
out.extend_from_slice(&buf[i..]);
}
#[cfg(test)]
mod tests {
use super::*;
fn dict(pairs: &[(&str, Bencode)]) -> Bencode {
Bencode::Dict(pairs.iter().map(|(k, v)| ((*k).as_bytes().to_vec(), v.clone())).collect())
}
fn bytes(s: &str) -> Bencode {
Bencode::Bytes(s.as_bytes().to_vec())
}
#[test]
fn parses_the_bep3_examples() {
assert_eq!(parse(b"i42e").unwrap(), Bencode::Int(42));
assert_eq!(parse(b"i-42e").unwrap(), Bencode::Int(-42));
assert_eq!(parse(b"i0e").unwrap(), Bencode::Int(0));
assert_eq!(parse(b"4:spam").unwrap(), bytes("spam"));
assert_eq!(parse(b"0:").unwrap(), bytes(""));
assert_eq!(
parse(b"l4:spami42ee").unwrap(),
Bencode::List(vec![bytes("spam"), Bencode::Int(42)])
);
assert_eq!(
parse(b"d3:bar4:spam3:fooi42ee").unwrap(),
dict(&[("bar", bytes("spam")), ("foo", Bencode::Int(42))])
);
}
#[test]
fn rejects_malformed_integers() {
assert_eq!(parse(b"i03e"), Err(Error::InvalidInteger { offset: 0 }));
assert_eq!(parse(b"i-0e"), Err(Error::InvalidInteger { offset: 0 }));
assert_eq!(parse(b"ie"), Err(Error::InvalidInteger { offset: 0 }));
assert_eq!(parse(b"i12x3e"), Err(Error::InvalidInteger { offset: 0 }));
assert!(matches!(parse(b"i99999999999999999999e"), Err(Error::IntegerOverflow { .. })));
}
#[test]
fn rejects_malformed_strings_and_dicts() {
assert!(matches!(parse(b"01:a"), Err(Error::InvalidLength { .. })));
assert_eq!(parse(b"5:abc"), Err(Error::UnexpectedEnd));
assert!(matches!(
parse(b"d3:fooi1e3:bari2ee"),
Err(Error::UnsortedOrDuplicateKey { .. })
));
assert!(matches!(
parse(b"d1:ai1e1:ai2ee"),
Err(Error::UnsortedOrDuplicateKey { .. })
));
assert!(matches!(parse(b"di1ei2ee"), Err(Error::NonStringKey { .. })));
}
#[test]
fn rejects_trailing_data() {
assert!(matches!(parse(b"i1ei2e"), Err(Error::TrailingData { .. })));
let (v, rest) = parse_prefix(b"i1ei2e").unwrap();
assert_eq!(v, Bencode::Int(1));
assert_eq!(rest, b"i2e");
}
#[test]
fn encodes_canonically_with_sorted_keys() {
let mut map = BTreeMap::new();
map.insert(b"foo".to_vec(), Bencode::Int(42));
map.insert(b"bar".to_vec(), bytes("spam"));
assert_eq!(Bencode::Dict(map).to_bytes(), b"d3:bar4:spam3:fooi42ee");
assert_eq!(Bencode::Int(-42).to_bytes(), b"i-42e");
assert_eq!(Bencode::Int(0).to_bytes(), b"i0e");
assert_eq!(bytes("").to_bytes(), b"0:");
}
#[test]
fn encodes_i64_extremes_without_format() {
assert_eq!(Bencode::Int(i64::MAX).to_bytes(), b"i9223372036854775807e");
assert_eq!(Bencode::Int(i64::MIN).to_bytes(), b"i-9223372036854775808e");
}
proptest::proptest! {
#[test]
fn parse_of_encode_is_identity(v in arb_bencode(4)) {
let bytes = v.to_bytes();
proptest::prop_assert_eq!(parse(&bytes).unwrap(), v);
}
#[test]
fn encode_of_parse_is_canonical_fixed_point(v in arb_bencode(4)) {
let bytes = v.to_bytes();
let reparsed = parse(&bytes).unwrap();
proptest::prop_assert_eq!(reparsed.to_bytes(), bytes);
}
}
fn arb_bencode(depth: u32) -> impl proptest::strategy::Strategy<Value = Bencode> {
use proptest::collection::{btree_map, vec};
use proptest::prelude::*;
let leaf = prop_oneof![
any::<i64>().prop_map(Bencode::Int),
vec(any::<u8>(), 0..16).prop_map(Bencode::Bytes),
];
leaf.prop_recursive(depth, 64, 8, |inner| {
prop_oneof![
vec(inner.clone(), 0..6).prop_map(Bencode::List),
btree_map(vec(any::<u8>(), 0..8), inner, 0..6).prop_map(Bencode::Dict),
]
})
}
}