bencode/

lib.rs

//! Bencode — the BitTorrent encoding ([BEP-3]) as a typed AST.
//!
//! Bencode is the wire format under every BitTorrent surface: `.torrent`
//! metainfo, tracker announces, the DHT's KRPC, the extension protocol.
//! This crate is the **keystone** of the pleme-io ENXAME suite
//! (`theory/ENXAME.md`): nothing else parses without it.
//!
//! It is built in the pleme-io idiom — the **TYPED-SPEC + INTERPRETER
//! TRIPLET** (the same shape as mado's `CsiCommand`, engawa's IR,
//! sui-spec's domains):
//!
//! * **Typed border** — [`Bencode`] is the parse-don't-validate AST.
//!   Bytes become one of four typed shapes; an ill-formed stream is a
//!   typed [`Error`], never a panic.
//! * **Parser** — [`parse`] / [`parse_prefix`]: bytes → [`Bencode`].
//! * **Emitter** — [`Bencode::to_bytes`]: AST → bytes through a typed
//!   builder, never `format!()` of the wire syntax (the ★★ TYPED
//!   EMISSION rule). Dict keys serialize in sorted order — the
//!   canonical contract BitTorrent hashing depends on — guaranteed by
//!   the [`std::collections::BTreeMap`] backing [`Bencode::Dict`].
//!
//! The round-trip law `parse(x.to_bytes()) == x` and the canonical law
//! `parse(b).map(to_bytes) == b` for canonical input are both pinned by
//! property tests.
//!
//! [BEP-3]: https://www.bittorrent.org/beps/bep_0003.html
//! [bencoding]: https://en.wikipedia.org/wiki/Bencode

#![forbid(unsafe_code)]

use std::collections::BTreeMap;

/// A bencoded value — the four shapes the grammar admits.
///
/// `Bytes` is the only string type: bencode has no text/binary
/// distinction on the wire, so the AST keeps raw bytes and offers
/// [`Bencode::as_str`] for the UTF-8 view when a field is known textual.
/// `Dict` is a [`BTreeMap`] so iteration — and therefore
/// [`Bencode::to_bytes`] — is always in the sorted-key canonical order
/// the protocol requires.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Bencode {
    /// `i<decimal>e` — a 64-bit signed integer.
    Int(i64),
    /// `<len>:<bytes>` — a length-prefixed byte string.
    Bytes(Vec<u8>),
    /// `l<elements>e` — an ordered list.
    List(Vec<Bencode>),
    /// `d<key><value>…e` — keys are byte strings, sorted + unique.
    Dict(BTreeMap<Vec<u8>, Bencode>),
}

/// A typed parse failure. Every malformed-stream condition is one of
/// these — the AST is never built from invalid bytes.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Error {
    /// The stream ended mid-value.
    UnexpectedEnd,
    /// A byte that can't begin a value (not `i`/`l`/`d`/digit), with its
    /// offset.
    UnexpectedByte { offset: usize, byte: u8 },
    /// An integer with a leading zero (`i03e`), a bare/negative zero
    /// (`i-0e`), or non-digit body.
    InvalidInteger { offset: usize },
    /// An integer that overflows `i64`.
    IntegerOverflow { offset: usize },
    /// A byte-string length that isn't valid base-10 or overruns the
    /// buffer.
    InvalidLength { offset: usize },
    /// A dict whose keys are out of order or duplicated (canonical
    /// parsing requires strictly-ascending unique keys).
    UnsortedOrDuplicateKey { offset: usize },
    /// A dict key that wasn't a byte string.
    NonStringKey { offset: usize },
    /// [`parse`] found valid trailing bytes after a complete value.
    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 {}

/// Parse exactly one bencoded value from `input`, requiring that the
/// value consumes the **entire** buffer.
///
/// # Errors
/// [`Error::TrailingData`] if bytes remain after the value; any parse
/// error otherwise.
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() })
    }
}

/// Parse one bencoded value from the front of `input`, returning the
/// value and the unconsumed tail. Used when bencoded values are framed
/// inside a larger stream (the peer-wire extension handshake, KRPC).
///
/// # Errors
/// Any of [`Error`].
pub fn parse_prefix(input: &[u8]) -> Result<(Bencode, &[u8]), Error> {
    Parser { full: input, pos: 0 }.value().map(|v| {
        // SAFETY-free: pos is always ≤ len.
        (v.0, &input[v.1..])
    })
}

struct Parser<'a> {
    full: &'a [u8],
    pos: usize,
}

impl Parser<'_> {
    /// Parse one value starting at `self.pos`; return it + the new pos.
    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;
        // Read up to the terminating 'e'.
        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; // consume 'e'
        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];
        // No leading zeros except a bare "0"; must be non-empty.
        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; // consume 'e'
        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;
            // Keys MUST be byte strings.
            if !self.peek()?.is_ascii_digit() {
                return Err(Error::NonStringKey { offset: key_offset });
            }
            let key = self.parse_bytes()?;
            // Canonical: strictly-ascending unique keys.
            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; // consume 'e'
        Ok(Bencode::Dict(map))
    }
}

/// A bencoded integer body is valid iff: non-empty, optional leading
/// `-`, no leading zeros (except a bare `0`), no `-0`, only digits.
fn validate_integer_body(body: &[u8]) -> Result<(), ()> {
    match body {
        [] => Err(()),
        [b'0'] => Ok(()),
        [b'-', b'0', ..] => Err(()),       // "-0", "-012"
        [b'0', _, ..] => Err(()),          // "0…" with more digits
        _ => {
            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 {
    /// Encode to canonical bencode bytes through a typed builder — never
    /// a `format!()` of the wire syntax (the ★★ TYPED EMISSION rule).
    /// Dict keys emit in sorted order ([`BTreeMap`] iteration), the
    /// canonicalization BitTorrent info-hashing relies on.
    #[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');
            }
        }
    }

    // ── Typed accessors (the consume-don't-stringly-match surface) ──

    /// The integer, if this is an [`Bencode::Int`].
    #[must_use]
    pub fn as_int(&self) -> Option<i64> {
        match self {
            Bencode::Int(n) => Some(*n),
            _ => None,
        }
    }

    /// The raw bytes, if this is a [`Bencode::Bytes`].
    #[must_use]
    pub fn as_bytes(&self) -> Option<&[u8]> {
        match self {
            Bencode::Bytes(b) => Some(b),
            _ => None,
        }
    }

    /// The UTF-8 string view of a byte string, if valid UTF-8.
    #[must_use]
    pub fn as_str(&self) -> Option<&str> {
        self.as_bytes().and_then(|b| std::str::from_utf8(b).ok())
    }

    /// The elements, if this is a [`Bencode::List`].
    #[must_use]
    pub fn as_list(&self) -> Option<&[Bencode]> {
        match self {
            Bencode::List(items) => Some(items),
            _ => None,
        }
    }

    /// The map, if this is a [`Bencode::Dict`].
    #[must_use]
    pub fn as_dict(&self) -> Option<&BTreeMap<Vec<u8>, Bencode>> {
        match self {
            Bencode::Dict(map) => Some(map),
            _ => None,
        }
    }

    /// Look up a key in a dict (by byte key). `None` if not a dict or the
    /// key is absent.
    #[must_use]
    pub fn get(&self, key: &[u8]) -> Option<&Bencode> {
        self.as_dict().and_then(|m| m.get(key))
    }
}

/// Push the base-10 bytes of an `i64` (TYPED EMISSION — a typed integer
/// rendered to its decimal value, not a `format!()` of escape syntax).
fn push_i64(out: &mut Vec<u8>, n: i64) {
    if n < 0 {
        out.push(b'-');
        // Negate into u64 to handle i64::MIN without overflow.
        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;
    }
    // Render digits into a small stack buffer, then copy in order. u64
    // is at most 20 decimal digits.
    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() {
        // leading-zero length
        assert!(matches!(parse(b"01:a"), Err(Error::InvalidLength { .. })));
        // length overruns buffer
        assert_eq!(parse(b"5:abc"), Err(Error::UnexpectedEnd));
        // out-of-order dict keys
        assert!(matches!(
            parse(b"d3:fooi1e3:bari2ee"),
            Err(Error::UnsortedOrDuplicateKey { .. })
        ));
        // duplicate key
        assert!(matches!(
            parse(b"d1:ai1e1:ai2ee"),
            Err(Error::UnsortedOrDuplicateKey { .. })
        ));
        // non-string key
        assert!(matches!(parse(b"di1ei2ee"), Err(Error::NonStringKey { .. })));
    }

    #[test]
    fn rejects_trailing_data() {
        assert!(matches!(parse(b"i1ei2e"), Err(Error::TrailingData { .. })));
        // parse_prefix keeps the tail instead
        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() {
        // Insert keys out of order; emission must sort them.
        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! {
        /// Round-trip law: parsing what we encoded yields the original.
        #[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);
        }

        /// Canonical law: re-encoding a parsed canonical stream is a
        /// byte-for-byte fixed point (what info-hashing relies on).
        #[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);
        }
    }

    /// A recursive arbitrary `Bencode` with bounded depth.
    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),
            ]
        })
    }
}