144 lines
3.5 KiB
Rust
144 lines
3.5 KiB
Rust
use std::collections::BTreeMap;
|
|
use std::num::ParseIntError;
|
|
use std::str::{self, Utf8Error};
|
|
|
|
use bencode::{Bytes, Dict, List, Object};
|
|
use buffer::Buffer;
|
|
|
|
#[derive(Debug)]
|
|
pub struct DecodeError;
|
|
|
|
impl From<Utf8Error> for DecodeError {
|
|
fn from(_: Utf8Error) -> Self {
|
|
DecodeError
|
|
}
|
|
}
|
|
|
|
impl From<ParseIntError> for DecodeError {
|
|
fn from(_: ParseIntError) -> Self {
|
|
DecodeError
|
|
}
|
|
}
|
|
|
|
pub type DecodeResult<T> = Result<T, DecodeError>;
|
|
|
|
pub fn decode(data: &[u8]) -> DecodeResult<Object> {
|
|
let mut buf = Buffer::new(data);
|
|
decode_object(&mut buf)
|
|
}
|
|
|
|
fn decode_object(buf: &mut Buffer) -> DecodeResult<Object> {
|
|
match buf.get(0) {
|
|
Some(b'i') => {
|
|
buf.advance(1);
|
|
decode_int(buf, b'e')
|
|
},
|
|
Some(b'0' ... b'9') => {
|
|
decode_bytes(buf)
|
|
}
|
|
Some(b'l') => {
|
|
buf.advance(1);
|
|
let mut list = Vec::new();
|
|
while let Some(term) = buf.get(0) {
|
|
if term == b'e' {
|
|
buf.advance(1);
|
|
break;
|
|
}
|
|
list.push(decode_object(buf)?);
|
|
}
|
|
Ok(List::wrap(list))
|
|
}
|
|
Some(b'd') => {
|
|
buf.advance(1);
|
|
let mut dict = BTreeMap::new();
|
|
while let Some(term) = buf.get(0) {
|
|
if term == b'e' {
|
|
buf.advance(1);
|
|
break;
|
|
}
|
|
let key = _decode_bytes(buf)?;
|
|
let val = decode_object(buf)?;
|
|
dict.insert(Bytes(key), val);
|
|
}
|
|
Ok(Dict::wrap(dict))
|
|
}
|
|
_ => Err(DecodeError),
|
|
}
|
|
}
|
|
|
|
fn decode_int(buf: &mut Buffer, term: u8) -> DecodeResult<Object> {
|
|
_decode_int(buf, term).map(|num| Object::Int(num))
|
|
}
|
|
|
|
fn _decode_int(buf: &mut Buffer, term: u8) -> DecodeResult<i64> {
|
|
if let Some(end) = buf.find(term) {
|
|
let obj = {
|
|
let num = str::from_utf8(&buf[..end])?;
|
|
num.parse()?
|
|
};
|
|
buf.advance(end + 1);
|
|
Ok(obj)
|
|
} else {
|
|
Err(DecodeError)
|
|
}
|
|
}
|
|
|
|
fn decode_bytes(buf: &mut Buffer) -> DecodeResult<Object> {
|
|
_decode_bytes(buf).map(|bytes| Bytes::wrap(bytes))
|
|
}
|
|
|
|
fn _decode_bytes(buf: &mut Buffer) -> Result<Vec<u8>, DecodeError> {
|
|
let size = _decode_int(buf, b':')? as usize;
|
|
let bytes = buf[..size].to_vec();
|
|
buf.advance(size);
|
|
Ok(bytes)
|
|
}
|
|
|
|
#[test]
|
|
fn test_int_pos() {
|
|
let mut buf = Buffer::new(b"i1337e");
|
|
buf.advance(1);
|
|
|
|
assert_eq!(decode_int(&mut buf, b'e').unwrap(), Object::Int(1337));
|
|
assert_eq!(buf.pos(), 6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_int_neg() {
|
|
let mut buf = Buffer::new(b"i-1337e");
|
|
buf.advance(1);
|
|
|
|
assert_eq!(decode_int(&mut buf, b'e').unwrap(), Object::Int(-1337));
|
|
assert_eq!(buf.pos(), 7);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bytes() {
|
|
let mut buf = Buffer::new(b"5:hello");
|
|
|
|
assert_eq!(decode_bytes(&mut buf).unwrap(), Bytes::wrap(b"hello".to_vec()));
|
|
assert_eq!(buf.pos(), 7);
|
|
}
|
|
|
|
#[test]
|
|
fn test_list() {
|
|
let mut buf = Buffer::new(b"li1ei2ei3ee");
|
|
|
|
let obj = decode_object(&mut buf).unwrap();
|
|
let list = obj.as_list().unwrap();
|
|
|
|
assert_eq!(list, &List(vec![Object::Int(1), Object::Int(2), Object::Int(3)]));
|
|
assert_eq!(buf.pos(), 11);
|
|
}
|
|
|
|
#[test]
|
|
fn test_dict() {
|
|
let mut buf = Buffer::new(b"d5:helloi1337ee");
|
|
|
|
let obj = decode_object(&mut buf).unwrap();
|
|
let dict = obj.as_dict().unwrap();
|
|
|
|
assert_eq!(dict.get_int("hello").unwrap(), 1337);
|
|
assert_eq!(buf.pos(), 15);
|
|
}
|