From 72d4dfc0b8cc92feda491af234e80bf60070d393 Mon Sep 17 00:00:00 2001 From: Simon Bernier St-Pierre Date: Sun, 11 Dec 2016 22:00:31 -0500 Subject: [PATCH] refactor and simplify the bencode objects --- src/bencode/decode.rs | 12 ++-- src/bencode/encode.rs | 45 +++++++----- src/bencode/mod.rs | 157 ++++++++++++++++++++++++++++++------------ src/metainfo.rs | 45 ++++++------ 4 files changed, 171 insertions(+), 88 deletions(-) diff --git a/src/bencode/decode.rs b/src/bencode/decode.rs index ac33fa6..928bc1f 100644 --- a/src/bencode/decode.rs +++ b/src/bencode/decode.rs @@ -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 { } 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 { 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 { } fn decode_bytes(buf: &mut Buffer) -> DecodeResult { - _decode_bytes(buf).map(|bytes| Bytes::wrap(bytes)) + _decode_bytes(buf).map(Into::into) } fn _decode_bytes(buf: &mut Buffer) -> Result, 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(); diff --git a/src/bencode/encode.rs b/src/bencode/encode.rs index dca2622..c5b5ad0 100644 --- a/src/bencode/encode.rs +++ b/src/bencode/encode.rs @@ -1,10 +1,28 @@ -use std::borrow::Borrow; +use std::borrow::Cow; use bencode::{Dict, Object}; -pub fn encode(obj: B) -> Vec where B: Borrow { +/// 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 { + 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 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, 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") } diff --git a/src/bencode/mod.rs b/src/bencode/mod.rs index 10fd9c7..cde9120 100644 --- a/src/bencode/mod.rs +++ b/src/bencode/mod.rs @@ -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(bytes: A) -> Object where A: Into> { - 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; -} - -impl AsBytes for Bytes { - fn as_bytes(&self) -> Cow { - Cow::Borrowed(self) - } -} - -impl<'a> AsBytes for &'a str { - fn as_bytes(&self) -> Cow { - Cow::Owned(Bytes(str::as_bytes(self).to_vec())) - } -} - -impl<'a> AsBytes for &'a [u8] { - fn as_bytes(&self) -> Cow { - Cow::Owned(Bytes(self.to_vec())) - } -} - #[derive(Clone, Debug, Eq, PartialEq)] pub struct List(Vec); @@ -148,8 +121,8 @@ impl List { List(Vec::new()) } - pub fn wrap(list: Vec) -> Object { - Object::List(List(list)) + pub fn push(&mut self, val: V) where V: Into { + self.0.push(val.into()) } } @@ -175,36 +148,36 @@ impl Dict { Dict(BTreeMap::new()) } - pub fn wrap(dict: BTreeMap) -> Object { - Object::Dict(Dict(dict)) + pub fn insert(&mut self, key: K, val: V) where K: Into, V: Into { + self.0.insert(key.into(), val.into()); } - pub fn get_object(&self, key: A) -> Option<&Object> where A: AsBytes { - self.0.get(&key.as_bytes()) + pub fn get_object(&self, key: A) -> Option<&Object> where A: Into { + self.0.get(&key.into()) } - pub fn get_int(&self, key: A) -> Option where A: AsBytes { - self.0.get(&key.as_bytes()).and_then(|o| o.as_int()) + pub fn get_int(&self, key: A) -> Option where A: Into { + self.0.get(&key.into()).and_then(|o| o.as_int()) } - pub fn get_bytes(&self, key: A) -> Option<&Bytes> where A: AsBytes { - self.0.get(&key.as_bytes()).and_then(|o| o.as_bytes()) + pub fn get_bytes(&self, key: A) -> Option<&Bytes> where A: Into { + self.0.get(&key.into()).and_then(|o| o.as_bytes()) } - pub fn get_str(&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(&self, key: A) -> Option<&str> where A: Into { + self.0.get(&key.into()).and_then(|o| o.as_bytes().and_then(|b| str::from_utf8(b).ok())) } - pub fn get_string(&self, key: A) -> Option where A: AsBytes { + pub fn get_string(&self, key: A) -> Option where A: Into { self.get_str(key).map(|s| s.to_string()) } - pub fn get_list(&self, key: A) -> Option<&List> where A: AsBytes { - self.0.get(&key.as_bytes()).and_then(|o| o.as_list()) + pub fn get_list(&self, key: A) -> Option<&List> where A: Into { + self.0.get(&key.into()).and_then(|o| o.as_list()) } - pub fn get_dict(&self, key: A) -> Option<&Dict> where A: AsBytes { - self.0.get(&key.as_bytes()).and_then(|o| o.as_dict()) + pub fn get_dict(&self, key: A) -> Option<&Dict> where A: Into { + 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 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> for Bytes { + fn from(from: Vec) -> 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> for Object { + fn from(from: Vec) -> Self { + Object::Bytes(Bytes(from)) + } +} + +impl From for Object { + fn from(from: Bytes) -> Self { + Object::Bytes(from) + } +} + +// list + +impl From> for List where T: Into { + fn from(from: Vec) -> Self { + List(from.into_iter().map(Into::into).collect()) + } +} + +impl From> for Object where T: Into { + fn from(from: Vec) -> Self { + Object::List(List(from.into_iter().map(Into::into).collect())) + } +} + +impl From for Object { + fn from(from: List) -> Self { + Object::List(from) + } +} + +// dict + +impl From> for Dict { + fn from(from: BTreeMap) -> Self { + Dict(from) + } +} + +impl From> for Object { + fn from(from: BTreeMap) -> Self { + Object::Dict(Dict(from)) + } +} + +impl From for Object { + fn from(from: Dict) -> Self { + Object::Dict(from) + } +} diff --git a/src/metainfo.rs b/src/metainfo.rs index 98bbd03..69879ba 100644 --- a/src/metainfo.rs +++ b/src/metainfo.rs @@ -24,7 +24,8 @@ pub struct Metainfo { } impl Metainfo { - pub fn from_bencode(obj: Object) -> Option { + pub fn from_bencode(obj: O) -> Option where O: Into { + 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);