#![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 { if let BencodeValue::Bytes(bytes) = self { Some(bytes.clone()) } else { None } } pub fn string(&self) -> Option> { if let Some(bytes) = self.bytes() { Some(String::from_utf8(bytes.to_vec())) } else { None } } pub fn int(&self) -> Option { 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>(bytes: T) -> Result { 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, } 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, } 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 { 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!(":(") } } }