sbstp/bt
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

refactor and simplify the bencode objects

Browse Source
This commit is contained in:
Simon Bernier St-Pierre
2016-12-11 22:00:31 -05:00
parent 881c1e2b01
commit 72d4dfc0b8
4 changed files with 171 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/bencode/decode.rs
View File

@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
use std::num::ParseIntError;
use std::str::{self, Utf8Error};
use bencode::{Bytes, Dict, List, Object};
use bencode::{Bytes, Object};
use buffer::Buffer;
#[derive(Debug)]
@@ -46,7 +46,7 @@ fn decode_object(buf: &mut Buffer) -> DecodeResult<Object> {
}
list.push(decode_object(buf)?);
}
Ok(List::wrap(list))
Ok(list.into())
}
Some(b'd') => {
buf.advance(1);
@@ -60,7 +60,7 @@ fn decode_object(buf: &mut Buffer) -> DecodeResult<Object> {
let val = decode_object(buf)?;
dict.insert(Bytes(key), val);
}
Ok(Dict::wrap(dict))
Ok(dict.into())
}
_ => Err(DecodeError),
}
@@ -84,7 +84,7 @@ fn _decode_int(buf: &mut Buffer, term: u8) -> DecodeResult<i64> {
}
fn decode_bytes(buf: &mut Buffer) -> DecodeResult<Object> {
_decode_bytes(buf).map(|bytes| Bytes::wrap(bytes))
_decode_bytes(buf).map(Into::into)
}
fn _decode_bytes(buf: &mut Buffer) -> Result<Vec<u8>, DecodeError> {
@@ -116,12 +116,14 @@ fn test_int_neg() {
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!(decode_bytes(&mut buf).unwrap(), "hello".into());
assert_eq!(buf.pos(), 7);
}
#[test]
fn test_list() {
use bencode::List;
let mut buf = Buffer::new(b"li1ei2ei3ee");
let obj = decode_object(&mut buf).unwrap();

45
src/bencode/encode.rs
View File

@@ -1,10 +1,28 @@
use std::borrow::Borrow;
use std::borrow::Cow;
use bencode::{Dict, Object};
pub fn encode<B>(obj: B) -> Vec<u8> where B: Borrow<Object> {
/// Allows the encode function to take something that can be converted into an object,
/// or a reference to an object.
pub trait Encodable<'a> {
fn get_object(self) -> Cow<'a, Object>;
}
impl<'a, T> Encodable<'a> for T where T: Into<Object> {
fn get_object(self) -> Cow<'a, Object> {
Cow::Owned(self.into())
}
}
impl<'a> Encodable<'a> for &'a Object {
fn get_object(self) -> Cow<'a, Object> {
Cow::Borrowed(self)
}
}
pub fn encode<'a, O>(obj: O) -> Vec<u8> where O: Encodable<'a> {
let mut buff = Vec::new();
encode_object(&mut buff, obj.borrow());
encode_object(&mut buff, &obj.get_object());
buff
}
@@ -45,35 +63,30 @@ fn encode_dict(buff: &mut Vec<u8>, dict: &Dict) {
#[test]
fn test_int_pos() {
assert_eq!(encode(Object::Int(1337)), b"i1337e");
assert_eq!(encode(1337), b"i1337e");
}
#[test]
fn test_int_neg() {
assert_eq!(encode(Object::Int(-1337)), b"i-1337e");
assert_eq!(encode(-1337), b"i-1337e");
}
#[test]
fn test_bytes() {
use bencode::Bytes;
assert_eq!(encode(Bytes::wrap(b"hello".to_vec())), b"5:hello");
assert_eq!(encode("hello"), b"5:hello");
}
#[test]
fn test_list() {
use bencode::List;
let list = vec![Object::Int(1), Object::Int(2), Object::Int(3)];
assert_eq!(encode(List::wrap(list)), b"li1ei2ei3ee");
assert_eq!(encode(list), b"li1ei2ei3ee");
}
#[test]
fn test_dict() {
use std::collections::BTreeMap;
use bencode::{Bytes, Dict};
use bencode::Dict;
let mut dict = BTreeMap::new();
dict.insert(Bytes(b"hello".to_vec()), Object::Int(1337));
assert_eq!(encode(Dict::wrap(dict)), b"d5:helloi1337ee")
let mut dict = Dict::new();
dict.insert("hello", 1337);
assert_eq!(encode(dict), b"d5:helloi1337ee")
}

157
src/bencode/mod.rs
View File

@@ -1,14 +1,13 @@
mod decode;
mod encode;
use std::borrow::Cow;
use std::collections::BTreeMap;
use std::fmt;
use std::ops::{Deref, DerefMut};
use std::str;
pub use self::decode::{decode, DecodeError, DecodeResult};
pub use self::encode::encode;
pub use self::encode::{encode, Encodable};
#[derive(Clone, Eq, PartialEq)]
pub enum Object {
@@ -91,10 +90,6 @@ impl Bytes {
Bytes(bytes.into())
}
pub fn wrap<A>(bytes: A) -> Object where A: Into<Vec<u8>> {
Object::Bytes(Bytes(bytes.into()))
}
pub fn str(&self) -> Option<&str> {
str::from_utf8(&self.0).ok()
}
@@ -118,28 +113,6 @@ impl DerefMut for Bytes {
}
}
pub trait AsBytes {
fn as_bytes(&self) -> Cow<Bytes>;
}
impl AsBytes for Bytes {
fn as_bytes(&self) -> Cow<Bytes> {
Cow::Borrowed(self)
}
}
impl<'a> AsBytes for &'a str {
fn as_bytes(&self) -> Cow<Bytes> {
Cow::Owned(Bytes(str::as_bytes(self).to_vec()))
}
}
impl<'a> AsBytes for &'a [u8] {
fn as_bytes(&self) -> Cow<Bytes> {
Cow::Owned(Bytes(self.to_vec()))
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct List(Vec<Object>);
@@ -148,8 +121,8 @@ impl List {
List(Vec::new())
}
pub fn wrap(list: Vec<Object>) -> Object {
Object::List(List(list))
pub fn push<V>(&mut self, val: V) where V: Into<Object> {
self.0.push(val.into())
}
}
@@ -175,36 +148,36 @@ impl Dict {
Dict(BTreeMap::new())
}
pub fn wrap(dict: BTreeMap<Bytes, Object>) -> Object {
Object::Dict(Dict(dict))
pub fn insert<K, V>(&mut self, key: K, val: V) where K: Into<Bytes>, V: Into<Object> {
self.0.insert(key.into(), val.into());
}
pub fn get_object<A>(&self, key: A) -> Option<&Object> where A: AsBytes {
self.0.get(&key.as_bytes())
pub fn get_object<A>(&self, key: A) -> Option<&Object> where A: Into<Bytes> {
self.0.get(&key.into())
}
pub fn get_int<A>(&self, key: A) -> Option<i64> where A: AsBytes {
self.0.get(&key.as_bytes()).and_then(|o| o.as_int())
pub fn get_int<A>(&self, key: A) -> Option<i64> where A: Into<Bytes> {
self.0.get(&key.into()).and_then(|o| o.as_int())
}
pub fn get_bytes<A>(&self, key: A) -> Option<&Bytes> where A: AsBytes {
self.0.get(&key.as_bytes()).and_then(|o| o.as_bytes())
pub fn get_bytes<A>(&self, key: A) -> Option<&Bytes> where A: Into<Bytes> {
self.0.get(&key.into()).and_then(|o| o.as_bytes())
}
pub fn get_str<A>(&self, key: A) -> Option<&str> where A: AsBytes {
self.0.get(&key.as_bytes()).and_then(|o| o.as_bytes().and_then(|b| str::from_utf8(b).ok()))
pub fn get_str<A>(&self, key: A) -> Option<&str> where A: Into<Bytes> {
self.0.get(&key.into()).and_then(|o| o.as_bytes().and_then(|b| str::from_utf8(b).ok()))
}
pub fn get_string<A>(&self, key: A) -> Option<String> where A: AsBytes {
pub fn get_string<A>(&self, key: A) -> Option<String> where A: Into<Bytes> {
self.get_str(key).map(|s| s.to_string())
}
pub fn get_list<A>(&self, key: A) -> Option<&List> where A: AsBytes {
self.0.get(&key.as_bytes()).and_then(|o| o.as_list())
pub fn get_list<A>(&self, key: A) -> Option<&List> where A: Into<Bytes> {
self.0.get(&key.into()).and_then(|o| o.as_list())
}
pub fn get_dict<A>(&self, key: A) -> Option<&Dict> where A: AsBytes {
self.0.get(&key.as_bytes()).and_then(|o| o.as_dict())
pub fn get_dict<A>(&self, key: A) -> Option<&Dict> where A: Into<Bytes> {
self.0.get(&key.into()).and_then(|o| o.as_dict())
}
}
@@ -221,3 +194,97 @@ impl DerefMut for Dict {
&mut self.0
}
}
// convertion utilities
// int
impl From<i64> for Object {
fn from(from: i64) -> Self {
Object::Int(from)
}
}
// bytes
impl<'a> From<&'a [u8]> for Bytes {
fn from(from: &'a [u8]) -> Self {
Bytes(from.to_vec())
}
}
impl<'a> From<&'a str> for Bytes {
fn from(from: &'a str) -> Self {
Bytes(from.as_bytes().to_vec())
}
}
impl From<Vec<u8>> for Bytes {
fn from(from: Vec<u8>) -> Self {
Bytes(from)
}
}
impl<'a> From<&'a [u8]> for Object {
fn from(from: &'a [u8]) -> Self {
Object::Bytes(Bytes(from.to_vec()))
}
}
impl<'a> From<&'a str> for Object {
fn from(from: &'a str) -> Self {
Object::Bytes(Bytes(from.as_bytes().to_vec()))
}
}
impl From<Vec<u8>> for Object {
fn from(from: Vec<u8>) -> Self {
Object::Bytes(Bytes(from))
}
}
impl From<Bytes> for Object {
fn from(from: Bytes) -> Self {
Object::Bytes(from)
}
}
// list
impl<T> From<Vec<T>> for List where T: Into<Object> {
fn from(from: Vec<T>) -> Self {
List(from.into_iter().map(Into::into).collect())
}
}
impl<T> From<Vec<T>> for Object where T: Into<Object> {
fn from(from: Vec<T>) -> Self {
Object::List(List(from.into_iter().map(Into::into).collect()))
}
}
impl From<List> for Object {
fn from(from: List) -> Self {
Object::List(from)
}
}
// dict
impl From<BTreeMap<Bytes, Object>> for Dict {
fn from(from: BTreeMap<Bytes, Object>) -> Self {
Dict(from)
}
}
impl From<BTreeMap<Bytes, Object>> for Object {
fn from(from: BTreeMap<Bytes, Object>) -> Self {
Object::Dict(Dict(from))
}
}
impl From<Dict> for Object {
fn from(from: Dict) -> Self {
Object::Dict(from)
}
}

45
src/metainfo.rs
View File

@@ -24,7 +24,8 @@ pub struct Metainfo {
}
impl Metainfo {
pub fn from_bencode(obj: Object) -> Option<Metainfo> {
pub fn from_bencode<O>(obj: O) -> Option<Metainfo> where O: Into<Object> {
let obj = obj.into();
let metainfo = ts!(obj.as_dict());
let info = ts!(metainfo.get_dict("info"));
@@ -106,23 +107,23 @@ fn piece_from_slice(src: &[u8]) -> [u8; 20] {
mod tests {
use std::path::Path;
use bencode::{Bytes, Dict, List, Object};
use bencode::{Dict, List, Object};
use metainfo::Metainfo;
#[test]
fn test_single_file() {
let mut meta = Dict::new();
meta.insert(Bytes::new("announce"), Bytes::wrap("http://ubuntu.com/tracker:6969"));
meta.insert("announce", "http://ubuntu.com/tracker:6969");
let mut info = Dict::new();
info.insert(Bytes::new("name"), Bytes::wrap("ubuntu-16.04-desktop.iso"));
info.insert(Bytes::new("length"), Object::Int(1024 * 1024 * 1024));
info.insert(Bytes::new("piece length"), Object::Int(1024 * 512));
info.insert(Bytes::new("pieces"), Bytes::wrap(""));
info.insert("name", "ubuntu-16.04-desktop.iso");
info.insert("length", 1024 * 1024 * 1024);
info.insert("piece length", 1024 * 512);
info.insert("pieces", "");
meta.insert(Bytes::new("info"), Object::Dict(info));
meta.insert("info", info);
let metainfo = Metainfo::from_bencode(Object::Dict(meta)).unwrap();
let metainfo = Metainfo::from_bencode(meta).unwrap();
assert_eq!(metainfo.announce, "http://ubuntu.com/tracker:6969");
assert_eq!(metainfo.piece_length, 1024 * 512);
@@ -134,27 +135,27 @@ mod tests {
#[test]
fn test_multiple_files() {
let mut meta = Dict::new();
meta.insert(Bytes::new("announce"), Bytes::wrap("http://ubuntu.com/tracker:6969"));
meta.insert("announce", "http://ubuntu.com/tracker:6969");
let mut info = Dict::new();
info.insert(Bytes::new("name"), Bytes::wrap("base_folder"));
info.insert(Bytes::new("piece length"), Object::Int(1024 * 512));
info.insert(Bytes::new("pieces"), Bytes::wrap(""));
info.insert("name", "base_folder");
info.insert("piece length", 1024 * 512);
info.insert("pieces", "");
let mut files = List::new();
let mut file1 = Dict::new();
file1.insert(Bytes::new("path"), List::wrap(vec![Bytes::wrap("folder2"), Bytes::wrap("image.txt")]));
file1.insert(Bytes::new("length"), Object::Int(512));
files.push(Object::Dict(file1));
file1.insert("path", vec!["folder2", "image.txt"]);
file1.insert("length", 512);
files.push(file1);
let mut file2 = Dict::new();
file2.insert(Bytes::new("path"), List::wrap(vec![Bytes::wrap("txt.jpg")]));
file2.insert(Bytes::new("length"), Object::Int(64));
files.push(Object::Dict(file2));
info.insert(Bytes::new("files"), Object::List(files));
file2.insert("path", vec!["txt.jpg"]);
file2.insert("length", 64);
files.push(file2);
info.insert("files", Object::List(files));
meta.insert(Bytes::new("info"), Object::Dict(info));
meta.insert("info", info);
let metainfo = Metainfo::from_bencode(Object::Dict(meta)).unwrap();
let metainfo = Metainfo::from_bencode(meta).unwrap();
assert_eq!(metainfo.announce, "http://ubuntu.com/tracker:6969");
assert_eq!(metainfo.piece_length, 1024 * 512);