use super::*;
use crate::LavaTorrentError;
use std::fs::File;
use std::hash::BuildHasher;
use std::io::{BufWriter, Write};
use std::path::Path;
pub fn write_string<S, W>(string: S, dst: &mut W) -> Result<(), LavaTorrentError>
where
S: AsRef<str>,
W: Write,
{
let string = string.as_ref();
dst.write_all(&string.len().to_string().into_bytes())?;
dst.write_all(&[STRING_DELIMITER])?;
dst.write_all(string.as_bytes())?;
Ok(())
}
pub fn write_bytes<B, W>(bytes: B, dst: &mut W) -> Result<(), LavaTorrentError>
where
B: AsRef<[u8]>,
W: Write,
{
let bytes = bytes.as_ref();
dst.write_all(&bytes.len().to_string().into_bytes())?;
dst.write_all(&[STRING_DELIMITER])?;
dst.write_all(bytes)?;
Ok(())
}
pub fn write_integer<W>(int: i64, dst: &mut W) -> Result<(), LavaTorrentError>
where
W: Write,
{
dst.write_all(&[INTEGER_PREFIX])?;
dst.write_all(int.to_string().as_bytes())?;
dst.write_all(&[INTEGER_POSTFIX])?;
Ok(())
}
pub fn write_list<L, W>(list: L, dst: &mut W) -> Result<(), LavaTorrentError>
where
L: AsRef<[BencodeElem]>,
W: Write,
{
let list = list.as_ref();
dst.write_all(&[LIST_PREFIX])?;
for item in list {
item.write_into(dst)?;
}
dst.write_all(&[LIST_POSTFIX])?;
Ok(())
}
pub fn write_dictionary<W, S>(
dict: &HashMap<String, BencodeElem, S>,
dst: &mut W,
) -> Result<(), LavaTorrentError>
where
W: Write,
S: BuildHasher,
{
let mut sorted = dict.iter().collect::<Vec<(&String, &BencodeElem)>>();
sorted.sort_by_key(|&(key, _)| key.as_bytes());
dst.write_all(&[DICTIONARY_PREFIX])?;
for (key, val) in sorted {
write_string(key, dst)?;
val.write_into(dst)?;
}
dst.write_all(&[DICTIONARY_POSTFIX])?;
Ok(())
}
pub fn write_raw_dictionary<W, S>(
dict: &HashMap<Vec<u8>, BencodeElem, S>,
dst: &mut W,
) -> Result<(), LavaTorrentError>
where
W: Write,
S: BuildHasher,
{
let mut sorted = dict.iter().collect::<Vec<(&Vec<u8>, &BencodeElem)>>();
sorted.sort_by_key(|&(key, _)| key);
dst.write_all(&[DICTIONARY_PREFIX])?;
for (key, val) in sorted {
write_bytes(key, dst)?;
val.write_into(dst)?;
}
dst.write_all(&[DICTIONARY_POSTFIX])?;
Ok(())
}
pub fn encode_string<S>(string: S) -> Vec<u8>
where
S: AsRef<str>,
{
let string = string.as_ref();
let mut encoded = Vec::with_capacity(string.len() + 2);
write_string(string, &mut encoded).expect("Write to vec failed!");
encoded
}
pub fn encode_bytes<B>(bytes: B) -> Vec<u8>
where
B: AsRef<[u8]>,
{
let bytes = bytes.as_ref();
let mut encoded = Vec::with_capacity(bytes.len() + 2);
write_bytes(bytes, &mut encoded).expect("Write to vec failed!");
encoded
}
pub fn encode_integer(int: i64) -> Vec<u8> {
let mut encoded = Vec::new();
write_integer(int, &mut encoded).expect("Write to vec failed!");
encoded
}
pub fn encode_list<L>(list: L) -> Vec<u8>
where
L: AsRef<[BencodeElem]>,
{
let mut encoded = Vec::new();
write_list(list, &mut encoded).expect("Write to vec failed!");
encoded
}
pub fn encode_dictionary<S>(dict: &HashMap<String, BencodeElem, S>) -> Vec<u8>
where
S: BuildHasher,
{
let mut encoded = Vec::new();
write_dictionary(dict, &mut encoded).expect("Write to vec failed!");
encoded
}
pub fn encode_raw_dictionary<S>(dict: &HashMap<Vec<u8>, BencodeElem, S>) -> Vec<u8>
where
S: BuildHasher,
{
let mut encoded = Vec::new();
write_raw_dictionary(dict, &mut encoded).expect("Write to vec failed!");
encoded
}
impl BencodeElem {
pub fn write_into<W>(&self, dst: &mut W) -> Result<(), LavaTorrentError>
where
W: Write,
{
match *self {
BencodeElem::String(ref string) => write_string(string, dst),
BencodeElem::Bytes(ref bytes) => write_bytes(bytes, dst),
BencodeElem::Integer(int) => write_integer(int, dst),
BencodeElem::List(ref list) => write_list(list, dst),
BencodeElem::Dictionary(ref dict) => write_dictionary(dict, dst),
BencodeElem::RawDictionary(ref dict) => write_raw_dictionary(dict, dst),
}
}
pub fn write_into_file<P>(&self, path: P) -> Result<(), LavaTorrentError>
where
P: AsRef<Path>,
{
let file = File::create(&path)?;
self.write_into(&mut BufWriter::new(&file))?;
file.sync_all()?;
Ok(())
}
pub fn encode(&self) -> Vec<u8> {
match *self {
BencodeElem::String(ref string) => encode_string(string),
BencodeElem::Bytes(ref bytes) => encode_bytes(bytes),
BencodeElem::Integer(int) => encode_integer(int),
BencodeElem::List(ref list) => encode_list(list),
BencodeElem::Dictionary(ref dict) => encode_dictionary(dict),
BencodeElem::RawDictionary(ref dict) => encode_raw_dictionary(dict),
}
}
}
#[cfg(test)]
mod bencode_elem_write_tests {
use super::*;
use std::collections::hash_map::RandomState;
use std::iter::FromIterator;
#[test]
fn write_string_ok() {
let mut vec = Vec::new();
write_string("spam", &mut vec).unwrap();
assert_eq!(vec, "4:spam".as_bytes().to_vec());
}
#[test]
fn write_bytes_ok() {
let mut vec = Vec::new();
write_bytes([0x01, 0x02, 0x03, 0x04], &mut vec).unwrap();
assert_eq!(vec, vec![b'4', b':', 0x01, 0x02, 0x03, 0x04]);
}
#[test]
fn write_integer_ok() {
let mut vec = Vec::new();
write_integer(42, &mut vec).unwrap();
assert_eq!(vec, vec![b'i', b'4', b'2', b'e']);
}
#[test]
fn write_list_ok() {
let mut vec = Vec::new();
write_list(&vec![bencode_elem!(42), bencode_elem!("spam")], &mut vec).unwrap();
assert_eq!(
vec,
vec![b'l', b'i', b'4', b'2', b'e', b'4', b':', b's', b'p', b'a', b'm', b'e']
);
}
#[test]
fn write_dictionary_ok() {
let mut vec = Vec::new();
write_dictionary::<_, RandomState>(
&HashMap::from_iter(
vec![
("spam".to_owned(), bencode_elem!(42)),
("cow".to_owned(), bencode_elem!("moo")),
]
.into_iter(),
),
&mut vec,
)
.unwrap();
assert_eq!(
vec,
vec![
b'd', b'3', b':', b'c', b'o', b'w', b'3', b':', b'm', b'o', b'o', b'4', b':', b's',
b'p', b'a', b'm', b'i', b'4', b'2', b'e', b'e',
]
);
}
#[test]
fn encode_string_ok() {
assert_eq!(encode_string("spam"), "4:spam".as_bytes().to_vec(),)
}
#[test]
fn encode_bytes_ok() {
assert_eq!(
encode_bytes([0x01, 0x02, 0x03, 0x04]),
vec![b'4', b':', 0x01, 0x02, 0x03, 0x04],
)
}
#[test]
fn encode_integer_ok() {
assert_eq!(encode_integer(42), vec![b'i', b'4', b'2', b'e'])
}
#[test]
fn encode_list_ok() {
assert_eq!(
encode_list(vec![bencode_elem!(42), bencode_elem!("spam")]),
vec![b'l', b'i', b'4', b'2', b'e', b'4', b':', b's', b'p', b'a', b'm', b'e'],
)
}
#[test]
fn encode_dictionary_ok() {
assert_eq!(
encode_dictionary::<RandomState>(&HashMap::from_iter(
vec![
("spam".to_owned(), bencode_elem!(42)),
("cow".to_owned(), bencode_elem!("moo")),
]
.into_iter()
)),
vec![
b'd', b'3', b':', b'c', b'o', b'w', b'3', b':', b'm', b'o', b'o', b'4', b':', b's',
b'p', b'a', b'm', b'i', b'4', b'2', b'e', b'e',
],
)
}
#[test]
fn bencode_elem_write_string_ok() {
let mut vec = Vec::new();
bencode_elem!("spam").write_into(&mut vec).unwrap();
assert_eq!(vec, "4:spam".as_bytes().to_vec());
}
#[test]
fn bencode_elem_write_bytes_ok() {
let mut vec = Vec::new();
bencode_elem!((0x01, 0x02, 0x03, 0x04))
.write_into(&mut vec)
.unwrap();
assert_eq!(vec, vec![b'4', b':', 0x01, 0x02, 0x03, 0x04]);
}
#[test]
fn bencode_elem_write_integer_ok() {
let mut vec = Vec::new();
bencode_elem!(42).write_into(&mut vec).unwrap();
assert_eq!(vec, vec![b'i', b'4', b'2', b'e']);
}
#[test]
fn bencode_elem_write_list_ok() {
let mut vec = Vec::new();
bencode_elem!([42, "spam"]).write_into(&mut vec).unwrap();
assert_eq!(
vec,
vec![b'l', b'i', b'4', b'2', b'e', b'4', b':', b's', b'p', b'a', b'm', b'e']
);
}
#[test]
fn bencode_elem_write_dictionary_ok() {
let mut vec = Vec::new();
bencode_elem!({ ("spam", 42), ("cow", "moo") })
.write_into(&mut vec)
.unwrap();
assert_eq!(
vec,
vec![
b'd', b'3', b':', b'c', b'o', b'w', b'3', b':', b'm', b'o', b'o', b'4', b':', b's',
b'p', b'a', b'm', b'i', b'4', b'2', b'e', b'e',
]
);
}
#[test]
fn bencode_elem_encode_string_ok() {
assert_eq!(bencode_elem!("spam").encode(), "4:spam".as_bytes().to_vec(),)
}
#[test]
fn bencode_elem_encode_bytes_ok() {
assert_eq!(
bencode_elem!((0x01, 0x02, 0x03, 0x04)).encode(),
vec![b'4', b':', 0x01, 0x02, 0x03, 0x04],
)
}
#[test]
fn bencode_elem_encode_integer_ok() {
assert_eq!(bencode_elem!(42).encode(), vec![b'i', b'4', b'2', b'e'])
}
#[test]
fn bencode_elem_encode_list_ok() {
assert_eq!(
bencode_elem!([42, "spam"]).encode(),
vec![b'l', b'i', b'4', b'2', b'e', b'4', b':', b's', b'p', b'a', b'm', b'e'],
)
}
#[test]
fn bencode_elem_encode_dictionary_ok() {
assert_eq!(
bencode_elem!({ ("spam", 42), ("cow", "moo") }).encode(),
vec![
b'd', b'3', b':', b'c', b'o', b'w', b'3', b':', b'm', b'o', b'o', b'4', b':', b's',
b'p', b'a', b'm', b'i', b'4', b'2', b'e', b'e',
],
)
}
}