Consider iterate over bencoded string bytes

[josecelano]

Relates to: #12 (comment)

In the new tokenizer, I decided:

• To use only byte vectors for integer and string tokens.
• Not to limit the size of the integer.
• Not to iterate over integer and string bytes in the new tokenizer. No nested iterators.

I wanted to keep this first tokenizer version as simple as possible, but we can discuss other @da2ce7's proposals like:

#[derive(Debug)]
pub enum Bencode<'a, R: BufRead + 'a> {
    Integer(BigInt),
    ByteString(Vec<u8>),
    List(BencodeListIterator<'a, R>),
    Dictionary(BencodeDictIterator<'a, R>),
}

to

#[derive(Debug)]
pub enum Bencode<'a, R: BufRead + 'a> {
    Integer(BigInt),
    ByteString(BencodeStringIterator<'a, R>),
    List(BencodeListIterator<'a, R>),
    Dictionary(BencodeDictIterator<'a, R>),
}

• The strings may also be tokenized progressively.
• The implementation limits the Integer to 1024bit.

[da2ce7]

use num_bigint::BigInt;
use std::cell::RefCell;
use std::io::{self, BufRead};
use std::rc::Rc;
use std::str::FromStr as _;

const MAX_RECURSION_DEPTH: usize = 100;
const MAX_BYTESTRING_LENGTH: usize = 10 * 1024 * 1024;

type ByteStringLength = u128;
const MAX_STRING_INTEGER_DIGITS: usize =
    max_decimal_digits_for_bits(ByteStringLength::BITS as usize);

const MAX_SIGNED_INTEGER_BITS: u64 = 1024;
const MAX_SIGNED_INTEGER_DIGITS: usize =
    max_decimal_digits_for_bits((MAX_SIGNED_INTEGER_BITS as usize) - 1) + 1;

const fn max_decimal_digits_for_bits(bits: usize) -> usize {
    // For unsigned integers: max value is 2^bits - 1
    // For signed integers: max absolute value is 2^(bits - 1)
    // Number of decimal digits = floor(bits * log10(2)) + 1
    ((bits * 30103) / 100000) + 1
}

#[derive(PartialEq, Clone)]
pub enum BencodeData {
    Integer(BigInt),
    ByteString(Vec<u8>),
    List(Vec<BencodeData>),
    Dictionary(Vec<(Vec<u8>, BencodeData)>),
}

impl std::fmt::Debug for BencodeData {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Integer(arg0) => f.debug_tuple("Integer").field(arg0).finish(),
            Self::ByteString(arg0) => f
                .debug_tuple("ByteString")
                .field(&arg0.iter().map(|i| *i as char).collect::<Vec<char>>())
                .finish(),
            Self::List(arg0) => f.debug_tuple("List").field(arg0).finish(),
            Self::Dictionary(arg0) => f.debug_tuple("Dictionary").field(arg0).finish(),
        }
    }
}

pub fn collect_bencode<'a, R: BufRead + 'a>(bencode: Bencode<'a, R>) -> io::Result<BencodeData> {
    collect_bencode_with_depth(bencode, 0)
}

pub fn collect_bencode_with_depth<'a, R: BufRead + 'a>(
    bencode: Bencode<'a, R>,
    depth: usize,
) -> io::Result<BencodeData> {
    if depth == MAX_RECURSION_DEPTH {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "max recursion depth reached when collecting bencode",
        ));
    }

    match bencode {
        Bencode::Integer(iter) => {
            let mut bytes = Vec::new();
            let mut counting_iter = iter.enumerate();

            while let Some((idx, byte)) = counting_iter.next() {
                if idx >= MAX_SIGNED_INTEGER_DIGITS {
                    return Err(io::Error::new(
                        io::ErrorKind::FileTooLarge,
                        format!("Integer is beyond: {MAX_SIGNED_INTEGER_DIGITS}, digits"),
                    ));
                }
                bytes.push(byte?);
            }
            let int_str = String::from_utf8(bytes)
                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
            let integer = BigInt::from_str(&int_str)
                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

            if integer.bits() > MAX_SIGNED_INTEGER_BITS {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "Integer created exceeds maximum allowed bit length",
                ));
            }

            Ok(BencodeData::Integer(integer))
        }
        Bencode::ByteString(iter) => {
            let mut bytes = Vec::new();
            let mut counting_iter = iter.enumerate();
            while let Some((idx, byte)) = counting_iter.next() {
                if idx >= MAX_BYTESTRING_LENGTH {
                    return Err(io::Error::new(
                        io::ErrorKind::FileTooLarge,
                        format!("Byte String is beyond Limit of: {MAX_BYTESTRING_LENGTH}, bytes"),
                    ));
                }
                bytes.push(byte?);
            }
            Ok(BencodeData::ByteString(bytes))
        }
        Bencode::List(mut iter) => {
            let mut list = Vec::new();
            while let Some(item) = iter.next() {
                let bencode_item = item?;
                list.push(collect_bencode_with_depth(bencode_item, depth + 1)?);
            }
            Ok(BencodeData::List(list))
        }
        Bencode::Dictionary(mut iter) => {
            let mut dict = Vec::new();
            let mut last_key: Option<Vec<u8>> = None;

            while let Some(pair_result) = iter.next() {
                let pair_tokenizer = pair_result?;
                let mut key_bytes = Vec::new();
                let mut value_option = None;
                for item in pair_tokenizer {
                    match item {
                        BencodeDictionaryPair::Key(key_iter) => {
                            let mut counting_iter = key_iter.enumerate();
                            while let Some((idx, byte)) = counting_iter.next() {
                                if idx >= MAX_BYTESTRING_LENGTH {
                                    return Err(io::Error::new(
                                        io::ErrorKind::FileTooLarge,
                                        format!("Byte String is beyond Limit of: {MAX_BYTESTRING_LENGTH}, bytes"),
                                    ));
                                }
                                key_bytes.push(byte?);
                            }
                        }
                        BencodeDictionaryPair::Value(value_result) => {
                            let value_bencode = value_result?;
                            value_option =
                                Some(collect_bencode_with_depth(value_bencode, depth + 1)?);
                        }
                    }
                }

                if let Some(ref last_key) = last_key {
                    match key_bytes.cmp(last_key) {
                        std::cmp::Ordering::Less => {
                            return Err(io::Error::new(
                                io::ErrorKind::InvalidData,
                                "Dictionary keys are not sorted lexicographically",
                            ));
                        }
                        std::cmp::Ordering::Equal => {
                            return Err(io::Error::new(
                                io::ErrorKind::InvalidData,
                                "Duplicate dictionary key",
                            ));
                        }
                        std::cmp::Ordering::Greater => {}
                    }
                }
                last_key = Some(key_bytes.clone());

                if let Some(value) = value_option {
                    dict.push((key_bytes, value));
                } else {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidData,
                        "Dictionary value missing",
                    ));
                }
            }
            Ok(BencodeData::Dictionary(dict))
        }
    }
}

pub struct BencodeTokenizer<'a, R: BufRead + 'a> {
    data: Rc<RefCell<R>>,
    depth: usize,
    _marker: std::marker::PhantomData<&'a ()>,
}

impl<'a, R: BufRead + 'a> BencodeTokenizer<'a, R> {
    pub fn new(reader: R) -> Self {
        BencodeTokenizer {
            data: Rc::new(RefCell::new(reader)),
            depth: 0,
            _marker: std::marker::PhantomData,
        }
    }

    pub fn next(reader: &Rc<RefCell<R>>, depth: usize) -> Self {
        BencodeTokenizer {
            data: Rc::clone(reader),
            depth: depth + 1,
            _marker: std::marker::PhantomData,
        }
    }

    fn tokenize(&self) -> io::Result<Bencode<'a, R>> {
        if self.depth == MAX_RECURSION_DEPTH {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "max recursion depth reached when parsing bencode type",
            ));
        }

        let mut data = self.data.borrow_mut();

        let type_prefix = data.fill_buf()?;

        let Some(type_prefix) = type_prefix.first() else {
            return Err(std::io::Error::new(
                std::io::ErrorKind::UnexpectedEof,
                "eof while parsing bencode list",
            ));
        };

        match *type_prefix {
            b'i' => {
                data.consume(1);
                drop(data);
                let iter = BencodeInteger {
                    data: Rc::clone(&self.data),
                    start: true,
                    end: false,
                    _marker: std::marker::PhantomData,
                };
                Ok(Bencode::Integer(iter))
            }
            b'l' => {
                data.consume(1);
                drop(data);
                let iter = BencodeList {
                    data: Rc::clone(&self.data),
                    end: false,
                    depth: self.depth,
                    _marker: std::marker::PhantomData,
                };
                Ok(Bencode::List(iter))
            }
            b'd' => {
                data.consume(1);
                drop(data);
                let iter = BencodeDictionary {
                    data: Rc::clone(&self.data),
                    eof: false,
                    depth: self.depth,
                    _marker: std::marker::PhantomData,
                };
                Ok(Bencode::Dictionary(iter))
            }
            b'0'..=b'9' => {
                // do not consume the 0..9, as it is part of the length encoding.
                drop(data);
                let iter = BencodeByteString {
                    data: Rc::clone(&self.data),
                    end: false,
                    _marker: std::marker::PhantomData,
                    start: true,
                    remaining: 0,
                };
                Ok(Bencode::ByteString(iter))
            }
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!(
                    "Invalid bencode type prefix: found: `{}`",
                    *type_prefix as char
                ),
            )),
        }
    }
}

#[derive(Debug)]
pub enum Bencode<'a, R: BufRead + 'a> {
    Integer(BencodeInteger<'a, R>),
    ByteString(BencodeByteString<'a, R>),
    List(BencodeList<'a, R>),
    Dictionary(BencodeDictionary<'a, R>),
}

#[derive(Debug)]
pub struct BencodeInteger<'a, R: BufRead + 'a> {
    data: Rc<RefCell<R>>,
    start: bool,
    end: bool,
    _marker: std::marker::PhantomData<&'a ()>,
}

impl<'a, R: BufRead + 'a> Iterator for BencodeInteger<'a, R> {
    type Item = io::Result<u8>;

    fn next(&mut self) -> Option<Self::Item> {
        let mut data = self.data.borrow_mut();

        if self.end {
            return None;
        }

        if self.start == true {
            self.start = false;

            // valid chars are: `'-' || '0'..='9'``
            let mut first = [0_u8; 1];
            if let Err(e) = data.read_exact(&mut first) {
                self.end = true;
                return Some(Err(e));
            };

            let first = first.first().unwrap();

            let second = match data.fill_buf() {
                Ok(second) => second,
                Err(e) => {
                    self.end = true;
                    return Some(Err(e));
                }
            };

            let Some(second) = second.first() else {
                self.end = true;
                return Some(Err(std::io::Error::new(
                    std::io::ErrorKind::UnexpectedEof,
                    "eof while parsing integer",
                )));
            };

            match *first {
                b'0' => {
                    self.end = true;
                    if *second != b'e' {
                        return Some(Err(std::io::Error::new(
                            std::io::ErrorKind::InvalidData,
                            format!(
                                "zero integer should end with 'e', but found: `{}`",
                                *second as char
                            ),
                        )));
                    }
                    data.consume(1); // consume 'e'.
                    return Some(Ok(b'0'));
                }

                b'-' => match *second {
                    b'e' => {
                        self.end = true;
                        return Some(Err(std::io::Error::new(
                            std::io::ErrorKind::InvalidData,
                            format!("zero-length negative integers are not valid bencode"),
                        )));
                    }
                    b'0' => {
                        self.end = true;
                        return Some(Err(std::io::Error::new(
                            std::io::ErrorKind::InvalidData,
                            format!("negative-zero integers are not valid bencode"),
                        )));
                    }
                    b'1'..=b'9' => return Some(Ok(b'-')),
                    _ => {
                        self.end = true;
                        return Some(Err(std::io::Error::new(
                            std::io::ErrorKind::InvalidData,
                            format!(
                                "negative-integers should start with a number: '1' to '9', but found: `{}`",
                                *second as char
                            ),
                        )));
                    }
                },

                b'1'..=b'9' => {
                    if *second == b'e' {
                        data.consume(1); // consume 'e'.
                        self.end = true;
                    }
                    return Some(Ok(*first));
                }

                _ => {
                    self.end = true;
                    return Some(Err(std::io::Error::new(std::io::ErrorKind::InvalidData, format!("integer should start with either: '-' (minus), or '0' to '9', but found: `{}`", *first as char))));
                }
            }
        } else {
            // valid chars are: `'0'..='9'``, or end.
            let mut next = [0_u8; 1];
            if let Err(e) = data.read_exact(&mut next) {
                return Some(Err(e));
            };

            let next = next.first().unwrap();

            match *next {
                b'e' => {
                    self.end = true;
                    return None;
                }
                b'0'..=b'9' => return Some(Ok(*next)),
                _ => {
                    return Some(Err(std::io::Error::new(
                        std::io::ErrorKind::InvalidData,
                        format!(
                            "integer should contain: '0' to '9', but found: `{}`",
                            *next as char
                        ),
                    )))
                }
            }
        }
    }
}

#[derive(Debug)]
pub struct BencodeByteString<'a, R: BufRead + 'a> {
    data: Rc<RefCell<R>>,
    start: bool,
    remaining: ByteStringLength,
    end: bool,
    _marker: std::marker::PhantomData<&'a ()>,
}

impl<'a, R: BufRead + 'a> Iterator for BencodeByteString<'a, R> {
    type Item = io::Result<u8>;

    fn next(&mut self) -> Option<Self::Item> {
        let mut data = self.data.borrow_mut();

        if self.end {
            return None;
        }

        // calculate length
        if self.start {
            self.start = false;
            self.remaining = match Self::find_string_length(&mut data) {
                Ok(None) => {
                    self.end = true;
                    return None;
                }
                Ok(Some(length)) => length,
                Err(e) => {
                    self.end = true;
                    return Some(Err(e));
                }
            }
        }

        self.remaining -= 1;
        if self.remaining == 0 {
            self.end = true;
        }

        let mut next = [0_u8; 1];
        if let Err(e) = data.read_exact(&mut next) {
            return Some(Err(e));
        };

        let next = next.first().unwrap();

        return Some(Ok(*next));
    }
}

impl<'a, R: BufRead + 'a> BencodeByteString<'a, R> {
    fn find_string_length(
        data: &mut std::cell::RefMut<'_, R>,
    ) -> std::io::Result<Option<ByteStringLength>> {
        let mut length_str = Vec::<u8>::new();
        let mut found = false;

        while !found {
            let buf = data.fill_buf()?;

            let Some(first) = buf.first() else {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::UnexpectedEof,
                    "eof while parsing integer",
                ));
            };

            if length_str.is_empty() {
                if *first == b'0' {
                    data.consume(1);

                    let mut second = [0_u8; 1];
                    if let Err(e) = data.read_exact(&mut second) {
                        return Err(e);
                    };

                    let second = second.first().unwrap();

                    if *second != b':' {
                        return Err(std::io::Error::new(
            std::io::ErrorKind::InvalidData,
            format!("zero-length stings should be a single integer character: expected delimiter: ':', but found:  `{}`",*second as char),
        ));
                    }

                    return Ok(None); // empty byte string.
                }
            }

            let pos = match buf.iter().enumerate().find(|(_, i)| !i.is_ascii_digit()) {
                Some((pos, delim)) => {
                    if *delim != b':' {
                        return Err(std::io::Error::new(
                            std::io::ErrorKind::InvalidData,
                            format!(
                                "integer should contain: '0' to '9', but found: `{}`",
                                *delim as char
                            ),
                        ));
                    }

                    found = true;

                    pos
                }
                None => buf.len(),
            };
            if MAX_STRING_INTEGER_DIGITS < (length_str.len() + pos) {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::FileTooLarge,
                    "string integer is has too many characters",
                ));
            }

            length_str.extend_from_slice(&buf[..pos]);

            if found {
                data.consume(pos + 1);
                break;
            } else {
                data.consume(pos);
            }
        }
        let length_str = std::str::from_utf8(&length_str)
            .expect("it should not fail as test for `is_ascii_digit` is above");

        Ok(Some(length_str.parse::<u128>().map_err(|e| {
            io::Error::new(io::ErrorKind::InvalidData, e)
        })?))
    }
}

#[derive(Debug)]
pub struct BencodeList<'a, R: BufRead + 'a> {
    data: Rc<RefCell<R>>,
    end: bool,
    depth: usize,
    _marker: std::marker::PhantomData<&'a ()>,
}

impl<'a, R: BufRead + 'a> Iterator for BencodeList<'a, R> {
    type Item = io::Result<Bencode<'a, R>>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.end {
            return None;
        }

        let mut data = self.data.borrow_mut();
        let end_peek = match data.fill_buf() {
            Ok(peek_type) => peek_type,
            Err(e) => return Some(Err(e)),
        };

        let Some(end_peek) = end_peek.first() else {
            self.end = true;
            return Some(Err(std::io::Error::new(
                std::io::ErrorKind::UnexpectedEof,
                "eof while parsing bencode list",
            )));
        };

        if *end_peek == b'e' {
            data.consume(1);
            self.end = true;
            return None;
        }
        drop(data);

        let tokenizer = BencodeTokenizer::next(&self.data, self.depth);
        Some(tokenizer.tokenize())
    }
}

pub enum BencodeDictionaryPair<'a, R: BufRead + 'a> {
    Key(BencodeByteString<'a, R>),
    Value(io::Result<Bencode<'a, R>>),
}

pub struct BencodeDictionaryPairTokenizer<'a, R: BufRead + 'a> {
    data: Rc<RefCell<R>>,
    depth: usize,
    start: bool,
    end: bool,
    _marker: std::marker::PhantomData<&'a ()>,
}

impl<'a, R: BufRead + 'a> Iterator for BencodeDictionaryPairTokenizer<'a, R> {
    type Item = BencodeDictionaryPair<'a, R>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.end {
            return None;
        }

        if self.start {
            self.start = false;

            let key = BencodeByteString {
                data: Rc::clone(&self.data),
                end: false,
                _marker: std::marker::PhantomData,
                start: true,
                remaining: 0,
            };

            return Some(BencodeDictionaryPair::Key(key));
        }

        self.end = true;

        let tokenizer = BencodeTokenizer::next(&self.data, self.depth);
        let value = tokenizer.tokenize();

        return Some(BencodeDictionaryPair::Value(value));
    }
}

#[derive(Debug)]
pub struct BencodeDictionary<'a, R: BufRead + 'a> {
    data: Rc<RefCell<R>>,
    eof: bool,
    depth: usize,
    _marker: std::marker::PhantomData<&'a ()>,
}

impl<'a, R: BufRead + 'a> Iterator for BencodeDictionary<'a, R> {
    type Item = io::Result<BencodeDictionaryPairTokenizer<'a, R>>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.eof {
            return None;
        }

        let mut data = self.data.borrow_mut();
        let end_peek = match data.fill_buf() {
            Ok(peek_type) => peek_type,
            Err(e) => return Some(Err(e)),
        };

        let Some(end_peek) = end_peek.first() else {
            self.eof = true;
            return Some(Err(std::io::Error::new(
                std::io::ErrorKind::UnexpectedEof,
                "eof while parsing bencode dictionary",
            )));
        };

        if *end_peek == b'e' {
            data.consume(1);
            self.eof = true;
            return None;
        }
        drop(data);

        let pair_tokenizer = BencodeDictionaryPairTokenizer {
            data: Rc::clone(&self.data),
            depth: self.depth,
            start: true,
            end: false,
            _marker: std::marker::PhantomData,
        };

        Some(Ok(pair_tokenizer))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use num_bigint::BigInt;
    use std::io::Cursor;

    #[test]
    fn test_bencode_parser() {
        use BencodeData::*;

        let test_cases = vec![
            // Valid integers
            (b"i42e".to_vec(), Ok(Integer(42.into())), "Valid integer"),
            (
                b"i-42e".to_vec(),
                Ok(Integer((-42).into())),
                "Valid negative integer",
            ),
            (b"i0e".to_vec(), Ok(Integer(0.into())), "Valid zero integer"),
            (
                b"i-1e".to_vec(),
                Ok(Integer((-1).into())),
                "Valid negative one",
            ),
            (
                b"i999999999999999999999999999999999999e".to_vec(),
                Ok(Integer(
                    "999999999999999999999999999999999999".parse().unwrap(),
                )),
                "Valid large integer",
            ),
            (
                b"i-999999999999999999999999999999999999e".to_vec(),
                Ok(Integer(
                    "-999999999999999999999999999999999999".parse().unwrap(),
                )),
                "Valid negative large integer",
            ),
            // Invalid integers
            (
                b"i-e".to_vec(),
                Err("zero-length negative integers are not valid bencode"),
                "Zero Length Negative Integer",
            ),
            (
                b"i-0e".to_vec(),
                Err("negative-zero integers are not valid bencode"),
                "Negative zero",
            ),

            (
                b"i-0123e".to_vec(),
                Err("negative-zero integers are not valid bencode"),
                "Negative Leading Zero",
            ),
            (
                b"i042e".to_vec(),
                Err("zero integer should end with 'e', but found: `4`"),
                "Leading zero",
            ),
            (
                b"i+42e".to_vec(),
                Err("integer should start with either: '-' (minus), or '0' to '9', but found: `+`"),
                "Plus sign",
            ),
            (
                b"i4.2e".to_vec(),
                Err("integer should contain: '0' to '9', but found: `.`"),
                "Non-digit character",
            ),
            (
                b"i4 2e".to_vec(),
                Err("integer should contain: '0' to '9', but found: ` `"),
                "Whitespace character",
            ),
            (
                b"ie".to_vec(),
                Err("eof while parsing integer"),
                "Empty integer",
            ),
            (
                b"i-42".to_vec(),
                Err("failed to fill whole buffer"),
                "Missing 'e' terminator",
            ),
            // Valid byte strings
            (b"0:".to_vec(), Ok(ByteString(vec![])), "Empty string"),
            (
                b"3:abc".to_vec(),
                Ok(ByteString(b"abc".to_vec())),
                "Simple string",
            ),
            (
                b"5:hello".to_vec(),
                Ok(ByteString(b"hello".to_vec())),
                "Hello string",
            ),
            (
                b"11:hello world".to_vec(),
                Ok(ByteString(b"hello world".to_vec())),
                "String with spaces",
            ),
            (
                b"5:helloextra".to_vec(),
                Ok(ByteString(b"hello".to_vec())),
                "Extra data after string",
            ),
            // Invalid byte strings
            (
                b"03:abc".to_vec(),
                Err("zero-length stings should be a single integer character: expected delimiter: ':', but found:  `3`"),
                "Leading zero in length",
            ),
            (
                b"1abc:hello".to_vec(),
                Err("integer should contain: '0' to '9', but found: `a`"),
                "Non-numeric length",
            ),
            (
                b"-5:hello".to_vec(),
                Err("Invalid bencode type prefix: found: `-`"),
                "Negative length",
            ),
            // Valid lists
            (b"le".to_vec(), Ok(List(vec![])), "Empty list"),
            (
                b"l4:spami42ee".to_vec(),
                Ok(List(vec![ByteString(b"spam".to_vec()), Integer(42.into())])),
                "List with string and integer",
            ),
            (
                b"l4:spam4:eggse".to_vec(),
                Ok(List(vec![
                    ByteString(b"spam".to_vec()),
                    ByteString(b"eggs".to_vec()),
                ])),
                "List of strings",
            ),
            // Invalid lists
            (
                b"l4:spam".to_vec(),
                Err("eof while parsing bencode list"),
                "Unterminated list",
            ),
            (
                b"l".to_vec(),
                Err("eof while parsing bencode list"),
                "Empty unterminated list",
            ),
            // Valid dictionaries
            (b"de".to_vec(), Ok(Dictionary(vec![])), "Empty dictionary"),
            (
                b"d3:cow3:mooe".to_vec(),
                Ok(Dictionary(vec![(
                    b"cow".to_vec(),
                    ByteString(b"moo".to_vec()),
                )])),
                "Simple dictionary",
            ),
            (
                b"d4:spaml1:a1:bee".to_vec(),
                Ok(Dictionary(vec![(
                    b"spam".to_vec(),
                    List(vec![ByteString(b"a".to_vec()), ByteString(b"b".to_vec())]),
                )])),
                "Dictionary with list",
            ),
            // Invalid dictionaries
            (
                b"d3:key5:value".to_vec(),
                Err("eof while parsing bencode dictionary"),
                "Unterminated dictionary",
            ),
            // Unsorted keys
            (
                b"d3:baa1:13:aaa1:2e".to_vec(),
                Err("Dictionary keys are not sorted lexicographically"),
                "Unsorted dictionary keys",
            ),
            // Duplicate keys
            (
                b"d3:key5:value3:key3:bbbe".to_vec(),
                Err("Duplicate dictionary key"),
                "Duplicate dictionary keys",
            ),
            // Non-UTF8 keys
            (
                b"d3:aaa3:bbb2:\xFF\xFE5:valuee".to_vec(),
                Ok(Dictionary(vec![
                    (b"aaa".to_vec(), ByteString(b"bbb".to_vec())),
                    (vec![0xFF, 0xFE], ByteString(b"value".to_vec())),
                ])),
                "Non-UTF8 dictionary keys",
            ),
            // Invalid prefixes
            (
                b"x42e".to_vec(),
                Err("Invalid bencode type prefix: found: `x`"),
                "Invalid prefix 'x'",
            ),
            (
                b"f3.14e".to_vec(),
                Err("Invalid bencode type prefix: found: `f`"),
                "Invalid prefix 'f'",
            ),
            (
                b"n123e".to_vec(),
                Err("Invalid bencode type prefix: found: `n`"),
                "Invalid prefix 'n'",
            ),
            // Extra data after valid bencode
            (
                b"i42ex".to_vec(),
                Ok(Integer(42.into())),
                "Integer with extra data",
            ),
        ];

        for (data, expected, description) in test_cases {
            let cursor = Cursor::new(data);
            let tokenizer = BencodeTokenizer::new(cursor);
            let result = tokenizer.tokenize();

            match (result, expected) {
                (Ok(bencode), Ok(ref expected_bencode)) => {
                    let collected = collect_bencode(bencode)
                        .expect(format!("Failed to collect bencode, {description}").as_str());
                    assert_eq!(&collected, expected_bencode, "{}", description);
                }
                (Err(e), Err(expected_error)) => {
                    assert_eq!(e.to_string(), expected_error, "{}", description);
                }
                (Ok(bencode), Err(expected_error)) => {
                    let collected = collect_bencode(bencode);
                    if let Err(e) = collected {
                        assert_eq!(e.to_string(), expected_error, "{}", description);
                    } else {
                        panic!(
                            "{}: Expected error '{}', but got success",
                            description, expected_error
                        );
                    }
                }
                (Err(e), Ok(_)) => {
                    panic!("{}: Expected success, but got error '{}'", description, e);
                }
            }
        }
    }

    #[test]
    fn test_recursion_depth_limit() {
        let mut data = Vec::new();
        for _ in 0..MAX_RECURSION_DEPTH {
            data.push(b'l');
        }
        data.extend(vec![b'e'; MAX_RECURSION_DEPTH]);

        let cursor = Cursor::new(&data);
        let tokenizer = BencodeTokenizer::new(cursor);
        let bencode = tokenizer.tokenize().unwrap();
        let _ = collect_bencode(bencode).expect("Should succeed at max depth");

        data.insert(0, b'l');
        data.push(b'e');

        let cursor = Cursor::new(&data);
        let tokenizer = BencodeTokenizer::new(cursor);
        let bencode = tokenizer.tokenize().unwrap();
        let error = collect_bencode(bencode).unwrap_err();

        assert_eq!(
            error.to_string(),
            "max recursion depth reached when parsing bencode type",
            "Error message should indicate depth exceeded"
        );
    }

    #[test]
    fn test_large_integer_limits() {
        // Test integer with MAX_SIGNED_INTEGER_DIGITS digits
        let digits = "1".repeat(MAX_SIGNED_INTEGER_DIGITS);
        let data = format!("i{}e", digits);
        let cursor = Cursor::new(data);
        let tokenizer = BencodeTokenizer::new(cursor);
        let bencode = tokenizer.tokenize().unwrap();
        let collected = collect_bencode(bencode).unwrap();

        if let BencodeData::Integer(value) = collected {
            assert_eq!(value.to_string(), digits);
        } else {
            panic!("Expected integer");
        }

        // Test integer with MAX_SIGNED_INTEGER_DIGITS + 1 digits (should error)
        let digits = "1".repeat(MAX_SIGNED_INTEGER_DIGITS + 1);
        let data = format!("i{}e", digits);
        let cursor = Cursor::new(data);
        let tokenizer = BencodeTokenizer::new(cursor);
        let bencode = tokenizer.tokenize().unwrap();
        let error = collect_bencode(bencode).unwrap_err();

        assert_eq!(error.to_string(), "Integer is beyond: 309, digits");

        // Test integer with MAX_SIGNED_INTEGER_DIGITS bits
        let max_value = (BigInt::from(1) << MAX_SIGNED_INTEGER_BITS) - 1;
        let data = format!("i{}e", max_value);
        let cursor = Cursor::new(data.to_string());
        let tokenizer = BencodeTokenizer::new(cursor);
        let bencode = tokenizer.tokenize().unwrap();
        let collected = collect_bencode(bencode).unwrap();

        if let BencodeData::Integer(value) = collected {
            assert_eq!(value, max_value);
        } else {
            panic!("Expected integer");
        }

        // Test integer exceeding MAX_SIGNED_INTEGER_DIGITS (should error)
        let over_max_value = BigInt::from(1) << MAX_SIGNED_INTEGER_BITS;
        let data = format!("i{}e", over_max_value);
        let cursor = Cursor::new(data.to_string());
        let tokenizer = BencodeTokenizer::new(cursor);
        let bencode = tokenizer.tokenize().unwrap();
        let error = collect_bencode(bencode).unwrap_err();

        assert_eq!(
            error.to_string(),
            "Integer created exceeds maximum allowed bit length"
        );

        // Test integer exceeding MAX_SIGNED_INTEGER_BITS but not MAX_SIGNED_INTEGER_DIGITS (should error)
        let digits = "9".repeat(MAX_SIGNED_INTEGER_DIGITS);
        let data = format!("i{}e", digits);
        let cursor = Cursor::new(data);
        let tokenizer = BencodeTokenizer::new(cursor);
        let bencode = tokenizer.tokenize().unwrap();
        let error = collect_bencode(bencode).unwrap_err();

        assert_eq!(
            error.to_string(),
            "Integer created exceeds maximum allowed bit length"
        );
    }

    #[test]
    fn test_large_bytestring_limits() {
        let data = format!(
            "{}:{}",
            MAX_BYTESTRING_LENGTH,
            "a".repeat(MAX_BYTESTRING_LENGTH)
        );
        let cursor = Cursor::new(data);
        let tokenizer = BencodeTokenizer::new(cursor);
        let bencode = tokenizer.tokenize().unwrap();
        let collected = collect_bencode(bencode).unwrap();

        if let BencodeData::ByteString(bytes) = collected {
            assert_eq!(bytes.len(), MAX_BYTESTRING_LENGTH);
        } else {
            panic!("Expected byte string");
        }

        let length = MAX_BYTESTRING_LENGTH + 1;
        let data = format!("{}:{}", length, "a".repeat(length));
        let cursor = Cursor::new(data);
        let tokenizer = BencodeTokenizer::new(cursor);
        let bencode = tokenizer.tokenize().unwrap();
        let error = collect_bencode(bencode).unwrap_err();

        assert_eq!(
            error.to_string(),
            format!("Byte String is beyond Limit of: {MAX_BYTESTRING_LENGTH}, bytes")
        );
    }
}

Here is a version of my code that iterates over the Integers and Strings. Some of the checks are moved therefore into the collect_bencode function.

Updated: I've made changes to make the BencodeTokenizer more correct.
Updated: I've renamed a few things and cleaned up some code.

[josecelano]

Hi @da2ce7, I like your implementation using iterators, but how can that fit into the current implementation? The main idea was to generate JSON on the fly without an intermediate in-memory data structure.

I only see two ways of using your code:

1. Create a JSON generator using your iterators

We would need to create a generator like the collect_bencode function, but instead of collecting values, generating the JSON.

Cons:

• The whole bencode in-memory representation would be in-memory at some point, and that's something we want to avoid.

Is there a way to use your code to generate the JSON without loading the whole bencoded value in memory?

2. Use your iterators inside the current tokenizer

We could implement a new tokenizer with your iterators.

Am I missing something?

On the other side, do we want to iterate over strings bytes? I think the whole point of the tokenizer is to add a new layer of abstraction (tokens). If we split tokens into bytes again, I don't see a reason to use the tokenier. WE can build with your code the same thing I had before introducing the tokenizer:

• A main loop to iterate over the bytes
• Print the JSON char on the fly

I mean, there would be no separation between tokenizer and generator, and we directly print the JSON output. Another output format could be built with the same iterations.

I see two models:

1. Tokenier/generators
2. Iterator/generators

cc @magecnion

[da2ce7]

@josecelano

I've updated the code above to make my comment here a little more clear.

The collect_bencode function iterates over the lazy recursive Bencode structure producing the eager BencodeData structure. I do not suggest using this function to emit the json.

Instead I suggest:

Create a new BencodeJsonTokenizer.

• The BencodeTokenizer produces the lazy recursive Bencode structure that emits bencode bytes in a structured way.
• A BencodeJsonTokenizer could produce a lazy-non-recursive BencodeJson structure that emits unicode characters in a json format.

For such a BencodeJsonTokenizer it would naturally internally use the BencodeTokenizer, it could lazily and recursively iterate over the Bencode structure and emit json.

So likewise it would could be a streaming implementation, and have a very low memory footprint.

[da2ce7]

#[derive(Debug)]
pub struct BencodeByteString<'a, R: BufRead + 'a> {
    data: Rc<RefCell<R>>,
    buf: Vec<u8>,
    buf_pos: usize,
    consume_pos: usize,
    start: bool,
    remaining: ByteStringLength,
    end: bool,
    _marker: std::marker::PhantomData<&'a ()>,
}

impl<'a, R: BufRead + 'a> BencodeByteString<'a, R> {
    fn update_buf(&mut self) -> io::Result<()> {
        if self.end {
            return Ok(());
        }

        // calculate length
        if self.start {
            self.start = false;
            let mut data = self.data.borrow_mut();
            self.remaining = match Self::find_string_length(&mut data) {
                Ok(None) => {
                    self.end = true;
                    return Ok(());
                }
                Ok(Some(length)) => length,
                Err(e) => {
                    self.end = true;
                    return Err(e);
                }
            }
        }

        if self.buf_pos > 0 {
            self.consume(
                self.buf_pos
                    .checked_sub(self.consume_pos)
                    .expect("it should not be larger than pos"),
            );
        }

        let mut data = self.data.borrow_mut();

        // calling `fill_buf` invalidates the old buffer.
        self.buf.clear();
        self.buf_pos = 0;
        self.consume_pos = 0;
        let buf = data.fill_buf()?;

        let remaining = usize::try_from(self.remaining).unwrap_or(usize::MAX);
        let read_to = std::cmp::min(remaining, buf.len());

        if read_to == 0 {
            self.end = true;
            return Ok(());
        }

        self.buf.extend_from_slice(&buf[..read_to]);

        Ok(())
    }
}

impl<'a, R: BufRead + 'a> Read for BencodeByteString<'a, R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if self.end {
            return Ok(0);
        }

        let max_read_pos = self
            .buf_pos
            .checked_add(buf.len())
            .expect("it should not be so big");

        if self.buf.len() < max_read_pos {
            self.update_buf()?;
        }

        // update max_read_pos
        let max_read_pos = self
            .buf_pos
            .checked_add(buf.len())
            .expect("it should not be so big");

        let available = self
            .buf
            .len()
            .checked_sub(self.buf_pos)
            .expect("it should not be negative");
        let amt: usize;

        if available >= buf.len() {
            buf.copy_from_slice(&self.buf[self.buf_pos..max_read_pos]);
            amt = buf.len();
        } else {
            buf[..available].copy_from_slice(&self.buf[self.buf_pos..]);
            amt = available;
        }

        self.consume(amt);

        return Ok(amt);
    }
}

impl<'a, R: BufRead + 'a> BufRead for BencodeByteString<'a, R> {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        if self.end {
            return Ok(&[]);
        }

        if self.consume_pos == self.buf.len() {
            self.update_buf()?;
        }

        Ok(&self.buf[self.consume_pos..])
    }

    fn consume(&mut self, amt: usize) {
        assert!(!self.buf.is_empty(), "it should not be empty");

        self.consume_pos = self
            .consume_pos
            .checked_add(amt)
            .expect("it should not be so large");

        assert!(
            self.consume_pos <= self.buf.len(),
            "it should only consume what is in the buffer"
        );

        self.buf_pos = self.buf_pos.max(self.consume_pos);

        let mut data = self.data.borrow_mut();
        data.consume(amt);

        self.remaining = self
            .remaining
            .checked_sub(amt.try_into().unwrap())
            .expect("it should not overflow");

        if self.remaining == 0 {
            self.end = true;
        }
    }
}

impl<'a, R: BufRead + 'a> Iterator for BencodeByteString<'a, R> {
    type Item = io::Result<u8>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.end {
            return None;
        }

        if self.buf_pos >= self.buf.len() {
            // will reset the position and buffer
            match self.update_buf() {
                Ok(()) => {}
                Err(e) => return Some(Err(e)),
            };
        }

        let next = self.buf.get(self.buf_pos).cloned();

        match next {
            Some(next) => {
                self.buf_pos += 1;
                Some(Ok(next))
            }
            None => None,
        }
    }
}

impl<'a, R: BufRead + 'a> BencodeByteString<'a, R> {
    fn find_string_length(
        data: &mut std::cell::RefMut<'_, R>,
    ) -> std::io::Result<Option<ByteStringLength>> {
        let mut length_str = Vec::<u8>::new();
        let mut found = false;

        while !found {
            let buf = data.fill_buf()?;

            let Some(first) = buf.first() else {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::UnexpectedEof,
                    "eof while parsing integer",
                ));
            };

            if length_str.is_empty() {
                if *first == b'0' {
                    data.consume(1);

                    let mut second = [0_u8; 1];
                    if let Err(e) = data.read_exact(&mut second) {
                        return Err(e);
                    };

                    let second = second.first().unwrap();

                    if *second != b':' {
                        return Err(std::io::Error::new(
            std::io::ErrorKind::InvalidData,
            format!("zero-length stings should be a single integer character: expected delimiter: ':', but found:  `{}`",*second as char),
        ));
                    }

                    return Ok(None); // empty byte string.
                }
            }

            let pos = match buf.iter().enumerate().find(|(_, i)| !i.is_ascii_digit()) {
                Some((pos, delim)) => {
                    if *delim != b':' {
                        return Err(std::io::Error::new(
                            std::io::ErrorKind::InvalidData,
                            format!(
                                "integer should contain: '0' to '9', but found: `{}`",
                                *delim as char
                            ),
                        ));
                    }

                    found = true;

                    pos
                }
                None => buf.len(),
            };
            if MAX_STRING_INTEGER_DIGITS < (length_str.len() + pos) {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::FileTooLarge,
                    "string integer is has too many characters",
                ));
            }

            length_str.extend_from_slice(&buf[..pos]);

            if found {
                data.consume(pos + 1);
                break;
            } else {
                data.consume(pos);
            }
        }
        let length_str = std::str::from_utf8(&length_str)
            .expect("it should not fail as test for `is_ascii_digit` is above");

        Ok(Some(length_str.parse::<u128>().map_err(|e| {
            io::Error::new(io::ErrorKind::InvalidData, e)
        })?))
    }
}

Here is an updated version of BencodeByteString, it contains a BufRead implementation that should be a bit faster... but when just using next, it seems not to be faster...