torment/torment-bencode/src/lib.rs

#![allow(dead_code)]
use bytes::Bytes;
use std::collections::HashMap;
use std::error::Error;
use std::fmt::{Display, Formatter};
use std::num::ParseIntError;
use std::str::FromStr;
use std::string::FromUtf8Error;

#[derive(Debug, Clone)]
pub enum BencodeValue {
    Int(i64),
    Bytes(Bytes),
    Dict(BencodeDict),
    List(BencodeList),
}

#[derive(Debug)]
pub enum BencodeError {
    InputTooShort,
    UnknownType,
    FailedToParseNumber(ParseIntError),
    Expected(char),
}

impl Display for BencodeError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:?}", self)
    }
}

impl Error for BencodeError {}

impl BencodeValue {
    pub fn bytes(&self) -> Option<Bytes> {
        if let BencodeValue::Bytes(bytes) = self {
            Some(bytes.clone())
        } else {
            None
        }
    }

    pub fn string(&self) -> Option<Result<String, FromUtf8Error>> {
        if let Some(bytes) = self.bytes() {
            Some(String::from_utf8(bytes.to_vec()))
        } else {
            None
        }
    }

    pub fn int(&self) -> Option<i64> {
        if let BencodeValue::Int(num) = self {
            Some(*num)
        } else {
            None
        }
    }

    pub fn dict(&self) -> Option<&BencodeDict> {
        if let BencodeValue::Dict(dict) = self {
            Some(dict)
        } else {
            None
        }
    }

    pub fn list(&self) -> Option<&BencodeList> {
        if let BencodeValue::List(list) = self {
            Some(list)
        } else {
            None
        }
    }

    pub fn decode<T: Into<Bytes>>(bytes: T) -> Result<Self, BencodeError> {
        let bytes = bytes.into();
        Ok(Self::consume(bytes)?.1)
    }

    fn consume(bytes: Bytes) -> Result<(usize, Self), BencodeError> {
        if bytes.len() < 2 {
            return Err(BencodeError::InputTooShort);
        }

        match bytes[0] {
            b'd' => {
                let (consumed, item) = BencodeDict::consume(bytes)?;
                Ok((consumed, BencodeValue::Dict(item)))
            }

            b'0'..=b'9' => {
                let (consumed, item) = BencodeValue::consume_bytes(bytes)?;
                Ok((consumed, BencodeValue::Bytes(item)))
            }

            b'i' => {
                let (consumed, item) = BencodeValue::consume_int(bytes)?;
                Ok((consumed, BencodeValue::Int(item)))
            }

            b'l' => {
                let (consumed, list) = BencodeList::consume(bytes)?;
                Ok((consumed, BencodeValue::List(list)))
            }

            _ => Err(BencodeError::UnknownType),
        }
    }

    fn consume_int(bytes: Bytes) -> Result<(usize, i64), BencodeError> {
        let mut offset = 1;
        while bytes.len() > offset && bytes[offset] != b'e' {
            offset += 1;
        }

        if bytes.len() <= offset {
            return Err(BencodeError::Expected('e'));
        }

        let length = String::from_utf8_lossy(&bytes[1..offset]);
        offset += 1;
        let nr = i64::from_str(&length)
            .map_err(|parse_error| BencodeError::FailedToParseNumber(parse_error))?;

        Ok((offset, nr))
    }

    fn consume_bytes(bytes: Bytes) -> Result<(usize, Bytes), BencodeError> {
        let mut offset = 1;
        while bytes.len() > offset && bytes[offset] != b':' {
            offset += 1;
        }

        if bytes.len() <= offset {
            return Err(BencodeError::Expected(':'));
        }
        let length = String::from_utf8_lossy(&bytes[0..offset]);
        offset += 1;
        let length = usize::from_str(&length)
            .map_err(|parse_error| BencodeError::FailedToParseNumber(parse_error))?;
        let result = bytes.slice(offset..offset + length);
        Ok((offset + length, result))
    }
}

#[derive(Debug, Clone)]
pub struct BencodeDict {
    bytes: Bytes,
    map: HashMap<Bytes, BencodeValue>,
}

impl BencodeDict {
    pub fn buffer(&self) -> Bytes {
        self.bytes.clone()
    }

    fn consume(bytes: Bytes) -> Result<(usize, Self), BencodeError> {
        let mut offset = 1;
        let mut map = HashMap::new();
        while bytes[offset] != b'e' {
            // consume key
            let (len, str) = BencodeValue::consume_bytes(bytes.slice(offset..))?;
            offset += len;
            let (len, value) = BencodeValue::consume(bytes.slice(offset..))?;
            offset += len;

            map.insert(str, value);
        }

        offset += 1;

        Ok((
            offset,
            BencodeDict {
                bytes: bytes.slice(0..offset),
                map,
            },
        ))
    }

    pub fn contains_key(&self, key: &[u8]) -> bool {
        self.map.contains_key(key)
    }

    pub fn get(&self, key: &[u8]) -> Option<&BencodeValue> {
        self.map.get(key)
    }
}

#[derive(Debug, Clone)]
pub struct BencodeList {
    bytes: Bytes,
    list: Vec<BencodeValue>,
}

impl BencodeList {
    pub fn buffer(&self) -> Bytes {
        self.bytes.clone()
    }

    fn consume(bytes: Bytes) -> Result<(usize, Self), BencodeError> {
        let mut offset = 1;
        let mut list = vec![];

        while bytes.len() > offset && bytes[offset] != b'e' {
            let (consumed, item) = BencodeValue::consume(bytes.slice(offset..))?;
            offset += consumed;
            list.push(item);
        }

        offset += 1;

        Ok((
            offset,
            BencodeList {
                bytes: bytes.slice(0..offset),
                list,
            },
        ))
    }

    pub fn len(&self) -> usize {
        self.list.len()
    }

    pub fn iter(&self) -> impl Iterator<Item = &BencodeValue> {
        self.list.iter()
    }

    pub fn get(&self, index: usize) -> Option<&BencodeValue> {
        self.list.get(index)
    }
}

#[cfg(test)]
mod tests {
    use crate::BencodeValue;

    #[test]
    fn test_string() {
        let input = b"4:spam";
        let value = BencodeValue::decode(input.as_ref()).unwrap();
        if let BencodeValue::Bytes(bytes) = value {
            assert_eq!(bytes.as_ref(), b"spam");
        } else {
            panic!(":(")
        }
    }

    #[test]
    fn test_integer() {
        let input = b"i42e";
        let value = BencodeValue::decode(input.as_ref()).unwrap();
        if let BencodeValue::Int(nr) = value {
            assert_eq!(nr, 42);
        } else {
            panic!(":(")
        }
    }

    #[test]
    fn test_dict() {
        let input = b"d4:spam3:egg3:eggdee";
        let value = BencodeValue::decode(input.as_ref()).unwrap();
        if let BencodeValue::Dict(dict) = value {
            assert!(dict.contains_key(b"spam"));
            assert_eq!(dict.get(b"spam").unwrap().bytes().unwrap().as_ref(), b"egg");
        } else {
            panic!(":(")
        }
    }

    #[test]
    fn test_list() {
        let input = b"l4:spam3:egge";
        let value = BencodeValue::decode(input.as_ref()).unwrap();
        if let BencodeValue::List(list) = value {
            assert_eq!(list.get(0).unwrap().bytes().unwrap().as_ref(), b"spam");
            assert_eq!(list.get(1).unwrap().bytes().unwrap().as_ref(), b"egg");
        } else {
            panic!(":(")
        }
    }
}