eater
/
torment
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

"Initial commit"

Browse Source
master
eater 4 years ago
commit b806c63df4
Signed by: eater
GPG Key ID: AD2560A0F84F0759
24 changed files with 6068 additions and 0 deletions
Show Stats Download Patch File Download Diff File

1
.gitignore vendored
Unescape Escape View File

@ -0,0 +1 @@
/target

1145
Cargo.lock generated
View File

File diff suppressed because it is too large Load Diff

9
Cargo.toml
Unescape Escape View File

@ -0,0 +1,9 @@
[workspace]
members = [
"torment-core",
"torment-cli",
"torment-dht",
"torment-dht-node",
"torment-bencode",
"torment-peer"
]

5
README.md
Unescape Escape View File

@ -0,0 +1,5 @@
# eater's Torment
A torrent client
It's written in Rust, that's the only selling point really

10
torment-bencode/Cargo.toml
Unescape Escape View File

@ -0,0 +1,10 @@
[package]
name = "torment-bencode"
version = "0.1.0"
authors = ["eater <=@eater.me>"]
edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
bytes = "0.5.6"

284
torment-bencode/src/lib.rs
Unescape Escape View File

@ -0,0 +1,284 @@
#![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<Bytes> {
if let BencodeValue::Bytes(bytes) = self {
Some(bytes.clone())
} else {
None
}
}
pub fn string(&self) -> Option<Result<String, FromUtf8Error>> {
if let Some(bytes) = self.bytes() {
Some(String::from_utf8(bytes.to_vec()))
} else {
None
}
}
pub fn int(&self) -> Option<i64> {
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<T: Into<Bytes>>(bytes: T) -> Result<Self, BencodeError> {
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<Bytes, BencodeValue>,
}
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<BencodeValue>,
}
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<Item = &BencodeValue> {
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!(":(")
}
}
}

11
torment-cli/Cargo.toml
Unescape Escape View File

@ -0,0 +1,11 @@
[package]
name = "torment-cli"
version = "0.1.0"
authors = ["eater <=@eater.me>"]
edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
clap = "2.33.3"
torment-core = { path = "../torment-core" }

39
torment-cli/src/main.rs
Unescape Escape View File

@ -0,0 +1,39 @@
use clap::{App, Arg, SubCommand};
use std::fs::File;
use std::io::Read;
use std::path::Path;
use torment_core::metainfo::Torrent;
fn main() {
let matches = App::new("torment-cli")
.version(env!("CARGO_PKG_VERSION"))
.subcommand(
SubCommand::with_name("torrent")
.about("Shows details about torrent file")
.arg(Arg::with_name("file").required(true)),
)
.get_matches();
match matches.subcommand() {
("torrent", Some(subcmd)) => {
let file = subcmd.value_of("file").unwrap();
let mut bytes: Vec<u8> = vec![];
File::open(file).unwrap().read_to_end(&mut bytes).unwrap();
match Torrent::from_bytes(
&bytes,
Path::new(file)
.file_stem()
.and_then(|file_name| file_name.to_str())
.map(|str| str.to_string()),
) {
Ok(torrent) => {
println!("{:#?}", torrent);
}
Err(err) => println!("Error: {}", err),
}
}
_ => {}
}
}

21
torment-core/Cargo.toml
Unescape Escape View File

@ -0,0 +1,21 @@
[package]
name = "torment-core"
version = "0.1.0"
authors = ["eater <=@eater.me>"]
edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
torment-bencode = { path = "../torment-bencode" }
hex = "0.4.2"
rand = "0.7.3"
bendy = "0.3.2"
serde = "1.0.115"
serde_derive = "1.0.115"
chrono = "0.4.15"
num-traits = "0.2.12"
url = "2.1.1"
bytes = "0.5.6"
ring = "0.16.15"
lazy_static = "1.4.0"

218
torment-core/src/infohash/mod.rs
Unescape Escape View File

@ -0,0 +1,218 @@
pub mod v1;
pub mod v2;
use hex::{FromHex, FromHexError, ToHex};
use std::convert::TryInto;
use std::error::Error;
use std::fmt::{Display, Formatter};
use std::iter;
pub use v1::InfoHash as InfoHashV1;
pub use v2::InfoHash as InfoHashV2;
pub trait InfoHashCapable: ToHex + FromHex<Error = InfoHashFromHashError> {
fn version(&self) -> u32;
fn to_bytes(&self) -> Box<[u8]>;
fn to_bytes_truncated(&self) -> [u8; 20] {
self.to_bytes()[0..20].try_into().unwrap()
}
/// Creates a V1 style info hash from this info hash
fn to_v1(&self) -> InfoHashV1;
fn from_bytes<T: AsRef<[u8]>>(bytes: T) -> Result<Self, InfoHashFromBytesError>;
}
#[derive(Debug)]
pub enum InfoHashFromHashError {
InvalidLength {
expected: Option<usize>,
actual: usize,
},
HexEncodingError(FromHexError),
}
impl Error for InfoHashFromHashError {}
impl Display for InfoHashFromHashError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::HexEncodingError(error) => write!(f, "Input is invalid hex: {}", error),
Self::InvalidLength { expected, actual } => write!(
f,
"Input has invalid length of {} expected {}",
actual,
expected
.map(|x| format!("{}", x))
.unwrap_or("40 or 64".to_string())
),
}
}
}
#[derive(Debug)]
pub enum InfoHashFromBytesError {
InvalidLength {
expected: Option<usize>,
actual: usize,
},
}
impl Error for InfoHashFromBytesError {}
impl Display for InfoHashFromBytesError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::InvalidLength { expected, actual } => write!(
f,
"Input has invalid length of {} expected {}",
actual,
expected
.map(|x| format!("{}", x))
.unwrap_or("20 or 32".to_string())
),
}
}
}
pub enum HybridInfoHash {
V1(InfoHashV1),
V2(InfoHashV2),
}
impl Display for HybridInfoHash {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.write_str(self.encode_hex::<String>().as_str())
}
}
impl ToHex for HybridInfoHash {
fn encode_hex<T: iter::FromIterator<char>>(&self) -> T {
match self {
Self::V1(info) => info.encode_hex(),
Self::V2(info) => info.encode_hex(),
}
}
fn encode_hex_upper<T: iter::FromIterator<char>>(&self) -> T {
match self {
Self::V1(info) => info.encode_hex_upper(),
Self::V2(info) => info.encode_hex_upper(),
}
}
}
impl FromHex for HybridInfoHash {
type Error = InfoHashFromHashError;
fn from_hex<T: AsRef<[u8]>>(hex: T) -> Result<Self, Self::Error> {
let hex = hex.as_ref();
Ok(match hex.len() {
64 => Self::V2(InfoHashV2::from_hex(hex)?),
40 => Self::V1(InfoHashV1::from_hex(hex)?),
len => {
return Err(InfoHashFromHashError::InvalidLength {
actual: len,
expected: None,
});
}
})
}
}
impl InfoHashCapable for HybridInfoHash {
fn version(&self) -> u32 {
match self {
Self::V1(_) => 1,
Self::V2(_) => 2,
}
}
fn to_bytes(&self) -> Box<[u8]> {
match self {
Self::V1(info) => info.to_bytes(),
Self::V2(info) => info.to_bytes(),
}
}
fn to_v1(&self) -> InfoHashV1 {
match self {
Self::V1(info) => info.to_v1(),
Self::V2(info) => info.to_v1(),
}
}
fn from_bytes<T: AsRef<[u8]>>(bytes: T) -> Result<Self, InfoHashFromBytesError> {
let bytes_ref = bytes.as_ref();
match bytes_ref.len() {
20 => Ok(Self::V1(InfoHashV1::from_bytes(bytes)?)),
32 => Ok(Self::V2(InfoHashV2::from_bytes(bytes)?)),
len => Err(InfoHashFromBytesError::InvalidLength {
expected: None,
actual: len,
}),
}
}
}
#[cfg(test)]
mod test {
use crate::infohash::{HybridInfoHash, InfoHashCapable};
use hex::{FromHex, ToHex};
#[test]
fn conversion() {
let info_hash_bytes: Vec<u8> = vec![
0x0a, 0xfe, 0x96, 0xa8, 0x62, 0x37, 0x8e, 0x7f, 0x69, 0x9d, 0x54, 0x6b, 0x7b, 0x8f,
0xe6, 0xfc, 0x47, 0xd0, 0xe2, 0x4a,
];
let info_hash =
HybridInfoHash::from_bytes(&info_hash_bytes).expect("Failed to parse info hash (v1)");
assert_eq!(1, info_hash.version());
assert_eq!(info_hash_bytes, info_hash.to_bytes().to_vec(), "Failed to assert that info hash (v1) could be converted to byte format and back without changing");
let info_hash_bytes: Vec<u8> = vec![
0xa3, 0x1f, 0xe9, 0x65, 0x6f, 0xc8, 0xd3, 0xa4, 0x59, 0xe6, 0x23, 0xdc, 0x82, 0x04,
0xe6, 0xd0, 0x26, 0x8f, 0x8d, 0xf5, 0x6d, 0x73, 0x4d, 0xac, 0x3c, 0xa3, 0x26, 0x2e,
0xdb, 0x5d, 0xb8, 0x83,
];
let info_hash =
HybridInfoHash::from_bytes(&info_hash_bytes).expect("Failed to parse info hash (v2)");
assert_eq!(info_hash_bytes, info_hash.to_bytes().to_vec(), "Failed to assert that info hash (v2) could be converted to byte format and back without changing");
}
#[test]
fn hex_conversion() {
let info_hash_hex = "0afe96a862378e7f699d546b7b8fe6fc47d0e24a";
let info_hash =
HybridInfoHash::from_hex(info_hash_hex).expect("Failed to parse info hash (v1)");
assert_eq!(1, info_hash.version());
assert_eq!(info_hash_hex, info_hash.encode_hex::<String>().as_str(), "Failed to assert that info hash (v1) could be converted to hex format and back without changing");
let info_hash_hex = "a31fe9656fc8d3a459e623dc8204e6d0268f8df56d734dac3ca3262edb5db883";
let info_hash =
HybridInfoHash::from_hex(info_hash_hex).expect("Failed to parse info hash (v2)");
assert_eq!(2, info_hash.version());
assert_eq!(info_hash_hex, info_hash.encode_hex::<String>().as_str(), "Failed to assert that info hash (v2) could be converted to hex format and back without changing");
}
#[test]
fn truncation() {
let info_hash_hex = "0afe96a862378e7f699d546b7b8fe6fc47d0e24a";
let info_hash =
HybridInfoHash::from_hex(info_hash_hex).expect("Failed to parse info hash (v1)");
assert_eq!(1, info_hash.version());
assert_eq!(
info_hash_hex,
info_hash.to_v1().encode_hex::<String>().as_str(),
"Failed to assert that info hash (v1) could be truncated to v1 hash without changing"
);
let info_hash_hex = "a31fe9656fc8d3a459e623dc8204e6d0268f8df56d734dac3ca3262edb5db883";
let info_hash =
HybridInfoHash::from_hex(info_hash_hex).expect("Failed to parse info hash (v2)");
assert_eq!(
&info_hash_hex[0..40],
info_hash.to_v1().encode_hex::<String>().as_str(),
"Failed to assert that info hash (v2) could be truncated to v1 hash without changing"
);
}
}

354
torment-core/src/infohash/v1.rs
Unescape Escape View File

@ -0,0 +1,354 @@
use crate::infohash::{InfoHashCapable, InfoHashFromBytesError, InfoHashFromHashError, InfoHashV1};
use bendy::decoding::{Error as DecodingError, FromBencode, Object};
use bendy::encoding::{Error as EncodingError, SingleItemEncoder, ToBencode};
use hex::{FromHex, ToHex};
use rand::random;
use serde::de::Error;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::convert::{TryFrom, TryInto};
use std::fmt::{Debug, Display, Formatter};
use std::iter;
use std::ops::{Add, BitAnd, BitOr, BitXor, Rem, Sub};
/// A make shift u160 type which uses a single u128 and u32 to keep operations speedy
/// Only really meant for ordering, bit operations and indexing
#[derive(Ord, PartialOrd, Eq, PartialEq, Copy, Clone, Hash, Default)]
pub struct U160(pub(crate) u128, pub(crate) u32);
impl Serialize for U160 {
fn serialize<S>(&self, serializer: S) -> Result<<S as Serializer>::Ok, <S as Serializer>::Error>
where
S: Serializer,
{
serializer.collect_str(&self)
}
}
impl<'de> Deserialize<'de> for U160 {
fn deserialize<D>(deserializer: D) -> Result<Self, <D as Deserializer<'de>>::Error>
where
D: Deserializer<'de>,
{
String::deserialize(deserializer)
.and_then(|x| U160::from_hex(x).map_err(|err| D::Error::custom(err.to_string())))
}
}
impl U160 {
pub const MAX: U160 = U160(u128::MAX, u32::MAX);
pub const MIN: U160 = U160(u128::MIN, u32::MIN);
pub const ONE: U160 = U160(0, 1);
pub const ZERO: U160 = U160(0, 0);
pub fn leading_zeroes(&self) -> u8 {
let mut leading_zeroes = self.0.leading_zeros() as u8;
if leading_zeroes == 128 {
leading_zeroes += self.1.leading_zeros() as u8;
}
leading_zeroes
}
}
impl ToBencode for U160 {
const MAX_DEPTH: usize = 0;
fn encode(&self, encoder: SingleItemEncoder) -> Result<(), EncodingError> {
encoder.emit_bytes(&self.to_byte_array()[..])
}
}
impl FromBencode for U160 {
fn decode_bencode_object(object: Object) -> Result<Self, DecodingError>
where
Self: Sized,
{
U160::from_bytes(object.try_into_bytes()?).map_err(|x| DecodingError::malformed_content(x))
}
}
impl<'obj, 'ser> TryFrom<Object<'obj, 'ser>> for U160 {
type Error = DecodingError;
fn try_from(value: Object<'obj, 'ser>) -> Result<Self, Self::Error> {
U160::decode_bencode_object(value)
}
}
impl Add for U160 {
type Output = U160;
fn add(self, rhs: Self) -> Self::Output {
let (lower, overflow) = self.1.overflowing_add(rhs.1);
U160(overflow as u128 + self.0 + rhs.0, lower)
}
}
impl Sub for U160 {
type Output = U160;
fn sub(self, rhs: Self) -> Self::Output {
let (lower, overflow) = self.1.overflowing_sub(rhs.1);
U160((self.0 - rhs.0) - overflow as u128, lower)
}
}
impl U160 {
pub fn random() -> U160 {
U160(random(), random())
}
pub fn half(&self) -> U160 {
U160(
self.0 / 2,
(self.1 / 2)
+ if self.0.rem(2) == 1 {
0b10000000_00000000_00000000_00000000
} else {
0
},
)
}
}
pub type InfoHash = U160;
impl Debug for U160 {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
if !f.alternate() {
f.write_str(&self.encode_hex::<String>())
} else {
f.debug_tuple("U160").field(&self.0).field(&self.1).finish()
}
}
}
impl Display for U160 {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.write_str(self.encode_hex::<String>().as_str())
}
}
impl From<&[u8; 20]> for InfoHash {
fn from(input: &[u8; 20]) -> Self {
U160(
u128::from_be_bytes(input[0..16].try_into().unwrap()),
u32::from_be_bytes(input[16..20].try_into().unwrap()),
)
}
}
impl From<[u8; 20]> for InfoHash {
fn from(input: [u8; 20]) -> Self {
Self::from(&input)
}
}
impl InfoHashCapable for InfoHash {
fn version(&self) -> u32 {
1
}
fn to_bytes(&self) -> Box<[u8]> {
self.to_byte_array().into()
}
fn to_v1(&self) -> InfoHashV1 {
*self
}
fn from_bytes<T: AsRef<[u8]>>(bytes: T) -> Result<Self, InfoHashFromBytesError> {
let bytes_ref = bytes.as_ref();
if bytes_ref.len() != 20 {
return Err(InfoHashFromBytesError::InvalidLength {
actual: bytes_ref.len(),
expected: Some(20),
});
}
let byte_array: &[u8; 20] = bytes_ref.try_into().unwrap();
Ok(Self::from(byte_array))
}
}
impl InfoHash {
pub fn to_byte_array(&self) -> [u8; 20] {
let mut bytes = [0u8; 20];
bytes[0..16].copy_from_slice(&self.0.to_be_bytes());
bytes[16..].copy_from_slice(&self.1.to_be_bytes());
bytes
}
}
impl FromHex for InfoHash {
type Error = InfoHashFromHashError;
fn from_hex<T: AsRef<[u8]>>(hex: T) -> Result<Self, Self::Error> {
let hex_ref = hex.as_ref();
if hex_ref.len() != 40 {
return Err(InfoHashFromHashError::InvalidLength {
expected: Some(40),
actual: hex.as_ref().len(),
});
}
let decoded = hex::decode(hex_ref).map_err(InfoHashFromHashError::HexEncodingError)?;
Ok(U160::from_bytes(decoded).unwrap())
}
}
impl ToHex for InfoHash {
fn encode_hex<T: iter::FromIterator<char>>(&self) -> T {
self.to_byte_array().encode_hex()
}
fn encode_hex_upper<T: iter::FromIterator<char>>(&self) -> T {
self.to_byte_array().encode_hex_upper()
}
}
impl BitXor for U160 {
type Output = U160;
fn bitxor(self, rhs: Self) -> Self::Output {
U160(self.0 ^ rhs.0, self.1 ^ rhs.1)
}
}
impl BitAnd for U160 {
type Output = U160;
fn bitand(self, rhs: Self) -> Self::Output {
U160(self.0 & rhs.0, self.1 & rhs.1)
}
}
impl BitOr for U160 {
type Output = U160;
fn bitor(self, rhs: Self) -> Self::Output {
U160(self.0 | rhs.0, self.1 | rhs.1)
}
}
#[cfg(test)]
mod test {
use crate::infohash::v1::{InfoHash, U160};
use crate::infohash::{InfoHashCapable, InfoHashFromBytesError, InfoHashFromHashError};
use hex::{FromHex, ToHex};
#[test]
fn conversion() {
let info_hash_bytes: [u8; 20] = [
0x0a, 0xfe, 0x96, 0xa8, 0x62, 0x37, 0x8e, 0x7f, 0x69, 0x9d, 0x54, 0x6b, 0x7b, 0x8f,
0xe6, 0xfc, 0x47, 0xd0, 0xe2, 0x4a,
];
let info_hash = InfoHash::from_bytes(info_hash_bytes).expect("Failed to parse info hash");
assert_eq!(1, info_hash.version());
assert_eq!(info_hash_bytes, info_hash.to_byte_array(), "Failed to assert that info hash could be converted to byte format and back without changing")
}
#[test]
fn hex_conversion() {
let info_hash_hex = "0afe96a862378e7f699d546b7b8fe6fc47d0e24a";
let info_hash = InfoHash::from_hex(info_hash_hex).expect("Failed to parse info hash");
assert_eq!(info_hash_hex, info_hash.encode_hex::<String>().as_str(), "Failed to assert that info hash could be converted to hex format and back without changing")
}
#[test]
fn formatting() {
let info_hash_hex = "0afe96a862378e7f699d546b7b8fe6fc47d0e24a";
let info_hash = InfoHash::from_hex(info_hash_hex).expect("Failed to parse info hash");
assert_eq!(
"0afe96a862378e7f699d546b7b8fe6fc47d0e24a",
format!("{}", info_hash)
);
assert_eq!(
"0afe96a862378e7f699d546b7b8fe6fc47d0e24a",
format!("{:?}", info_hash)
);
assert_eq!(
"U160(
14614179062085549902615142429809960700,
1204871754,
)",
format!("{:#?}", info_hash)
);
}
#[test]
fn bit_operations() {
let info_hash_alpha: InfoHash = [0u8; 20].into();
let info_hash_omega: InfoHash = [255u8; 20].into();
assert_eq!(
info_hash_omega ^ info_hash_omega,
info_hash_alpha,
"XOR has invalid result"
);
assert_eq!(
info_hash_alpha ^ info_hash_omega,
info_hash_omega,
"XOR has invalid result"
);
assert_eq!(
info_hash_alpha | info_hash_omega,
info_hash_omega,
"OR has invalid result"
);
assert_eq!(
info_hash_alpha | info_hash_alpha,
info_hash_alpha,
"OR has invalid result"
);
assert_eq!(
info_hash_alpha & info_hash_omega,
info_hash_alpha,
"OR has invalid result"
);
assert_eq!(
info_hash_omega & info_hash_omega,
info_hash_omega,
"OR has invalid result"
);
}
#[test]
fn ordering() {
assert!(U160(1, 0) > U160(0, u32::MAX));
assert!(U160(0, 0) < U160(0, u32::MAX));
}
#[test]
fn invalid_length() {
let info_hash_hex = "0afe96a862378e7f699d546b7b8fe6fc47d0e24a3";
let info_hash = InfoHash::from_hex(info_hash_hex).expect_err("No error thrown");
if let InfoHashFromHashError::InvalidLength { expected, actual } = info_hash {
assert_eq!(expected, Some(40));
assert_eq!(actual, 41);
} else {
panic!("Should've thrown invalid length");
}
let info_hash_hex = hex::decode("0afe96a862378e7f699d546b7b8fe6fc47d0e24aaa").unwrap();
let info_hash = InfoHash::from_bytes(info_hash_hex).expect_err("No error thrown");
let InfoHashFromBytesError::InvalidLength { expected, actual } = info_hash;
assert_eq!(expected, Some(20));
assert_eq!(actual, 21);
}
#[test]
fn u160_math() {
assert_eq!(U160::ZERO + U160::ONE, U160::ONE);
assert_eq!(U160(1, 0) - U160::ONE, U160(0, u32::MAX))
}
}

212
torment-core/src/infohash/v2.rs
Unescape Escape View File

@ -0,0 +1,212 @@
use crate::infohash::v1::U160;
use crate::infohash::{InfoHashCapable, InfoHashFromBytesError, InfoHashFromHashError, InfoHashV1};
use hex::{FromHex, ToHex};
use rand::random;
use std::convert::TryInto;
use std::fmt::{Debug, Display, Formatter};
use std::iter;
use std::ops::{BitAnd, BitOr, BitXor};
/// A make shift u265 type which uses a 2 u128's to keep operations speedy
/// Only really meant for ordering, bit operations and indexing so no math ops are implemented
#[derive(PartialOrd, Ord, Eq, PartialEq, Copy, Clone, Hash)]
pub struct U256(pub(crate) u128, pub(crate) u128);
impl U256 {
pub fn random() -> U256 {
U256(random(), random())
}
}
pub type InfoHash = U256;
impl Debug for U256 {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
if !f.alternate() {
f.write_str(&self.encode_hex::<String>())
} else {
f.debug_tuple("U256").field(&self.0).field(&self.1).finish()
}
}
}
impl Display for U256 {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.write_str(self.encode_hex::<String>().as_str())
}
}
impl From<&[u8; 32]> for InfoHash {
fn from(input: &[u8; 32]) -> Self {
U256(
u128::from_be_bytes(input[0..16].try_into().unwrap()),
u128::from_be_bytes(input[16..32].try_into().unwrap()),
)
}
}
impl From<[u8; 32]> for InfoHash {
fn from(input: [u8; 32]) -> Self {
Self::from(&input)
}
}
impl InfoHashCapable for InfoHash {
fn version(&self) -> u32 {
2
}
fn to_bytes(&self) -> Box<[u8]> {
self.to_byte_array().into()
}
fn to_v1(&self) -> InfoHashV1 {
U160(self.0, (self.1 >> 96) as u32)
}
fn from_bytes<T: AsRef<[u8]>>(bytes: T) -> Result<Self, InfoHashFromBytesError> {
let bytes_ref = bytes.as_ref();
if bytes_ref.len() != 32 {
return Err(InfoHashFromBytesError::InvalidLength {
actual: bytes_ref.len(),
expected: Some(32),
});
}
let byte_array: &[u8; 32] = bytes_ref.try_into().unwrap();
Ok(Self::from(byte_array))
}
}
impl InfoHash {
pub fn to_byte_array(&self) -> [u8; 32] {
let mut bytes = [0u8; 32];
bytes[0..16].copy_from_slice(&self.0.to_be_bytes());
bytes[16..].copy_from_slice(&self.1.to_be_bytes());
bytes
}
}
impl FromHex for InfoHash {
type Error = InfoHashFromHashError;
fn from_hex<T: AsRef<[u8]>>(hex: T) -> Result<Self, Self::Error> {
let hex_ref = hex.as_ref();
if hex_ref.len() != 64 {
return Err(InfoHashFromHashError::InvalidLength {
expected: Some(64),
actual: hex.as_ref().len(),
});
}
let decoded = hex::decode(hex_ref).map_err(InfoHashFromHashError::HexEncodingError)?;
Ok(U256::from_bytes(decoded).unwrap())
}
}
impl ToHex for U256 {
fn encode_hex<T: iter::FromIterator<char>>(&self) -> T {
self.to_byte_array().encode_hex()
}
fn encode_hex_upper<T: iter::FromIterator<char>>(&self) -> T {
self.to_byte_array().encode_hex_upper()
}
}
impl BitXor for U256 {
type Output = U256;
fn bitxor(self, rhs: Self) -> Self::Output {
U256(self.0 ^ rhs.0, self.1 ^ rhs.1)
}
}
impl BitAnd for U256 {
type Output = U256;
fn bitand(self, rhs: Self) -> Self::Output {
U256(self.0 & rhs.0, self.1 & rhs.1)
}
}
impl BitOr for U256 {
type Output = U256;
fn bitor(self, rhs: Self) -> Self::Output {
U256(self.0 | rhs.0, self.1 | rhs.1)
}
}
#[cfg(test)]
mod test {
use crate::infohash::v2::{InfoHash, U256};
use crate::infohash::InfoHashCapable;
use hex::{FromHex, ToHex};
#[test]
fn conversion() {
let info_hash_bytes: [u8; 32] = [
0xa3, 0x1f, 0xe9, 0x65, 0x6f, 0xc8, 0xd3, 0xa4, 0x59, 0xe6, 0x23, 0xdc, 0x82, 0x04,
0xe6, 0xd0, 0x26, 0x8f, 0x8d, 0xf5, 0x6d, 0x73, 0x4d, 0xac, 0x3c, 0xa3, 0x26, 0x2e,
0xdb, 0x5d, 0xb8, 0x83,
];
let info_hash = InfoHash::from_bytes(info_hash_bytes).expect("Failed to parse info hash");
assert_eq!(info_hash_bytes, info_hash.to_byte_array(), "Failed to assert that info hash could be converted to byte format and back without changing")
}
#[test]
fn hex_conversion() {
let info_hash_hex = "a31fe9656fc8d3a459e623dc8204e6d0268f8df56d734dac3ca3262edb5db883";
let info_hash = InfoHash::from_hex(info_hash_hex).expect("Failed to parse info hash");
assert_eq!(info_hash_hex, info_hash.encode_hex::<String>().as_str(), "Failed to assert that info hash could be converted to hex format and back without changing")
}
#[test]
fn bit_operations() {
let info_hash_alpha: InfoHash = [0u8; 32].into();
let info_hash_omega: InfoHash = [255u8; 32].into();
assert_eq!(
info_hash_omega ^ info_hash_omega,
info_hash_alpha,
"XOR has invalid result"
);
assert_eq!(
info_hash_alpha ^ info_hash_omega,
info_hash_omega,
"XOR has invalid result"
);
assert_eq!(
info_hash_alpha | info_hash_omega,
info_hash_omega,
"OR has invalid result"
);
assert_eq!(
info_hash_alpha | info_hash_alpha,
info_hash_alpha,
"OR has invalid result"
);
assert_eq!(
info_hash_alpha & info_hash_omega,
info_hash_alpha,
"OR has invalid result"
);
assert_eq!(
info_hash_omega & info_hash_omega,
info_hash_omega,
"OR has invalid result"
);
}
#[test]
fn ordering() {
assert!(U256(1, 0) > U256(0, u128::MAX));
assert!(U256(0, 0) < U256(0, u128::MAX));
}
}

580
torment-core/src/ip.rs
Unescape Escape View File

@ -0,0 +1,580 @@
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
/// This code is the unstable ip code from the rust std
/// Since I would like to build on the stable rust, I just copied it yolo
#[derive(Copy, PartialEq, Eq, Clone, Hash, Debug)]
pub enum Ipv6MulticastScope {
InterfaceLocal,
LinkLocal,
RealmLocal,
AdminLocal,
SiteLocal,
OrganizationLocal,
Global,
}
pub trait IpAddrExt {
fn ext_is_global(&self) -> bool;
fn ext_is_documentation(&self) -> bool;
}
impl IpAddrExt for IpAddr {
/// Returns [`true`] if the address appears to be globally routable.
///
/// See the documentation for [`Ipv4Addr::is_global`][IPv4] and
/// [`Ipv6Addr::is_global`][IPv6] for more details.
///
/// [IPv4]: ../../std/net/struct.Ipv4Addr.html#method.is_global
/// [IPv6]: ../../std/net/struct.Ipv6Addr.html#method.is_global
/// [`true`]: ../../std/primitive.bool.html
///
/// # Examples
///
/// ```
/// use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
/// use torment_core::ip::IpAddrExt;
///
/// assert_eq!(IpAddr::V4(Ipv4Addr::new(80, 9, 12, 3)).ext_is_global(), true);
/// assert_eq!(IpAddr::V6(Ipv6Addr::new(0, 0, 0x1c9, 0, 0, 0xafc8, 0, 0x1)).ext_is_global(), true);
/// ```
fn ext_is_global(&self) -> bool {
match self {
IpAddr::V4(ip) => ip.ext_is_global(),
IpAddr::V6(ip) => ip.ext_is_global(),
}
}
/// Returns [`true`] if this address is in a range designated for documentation.
///
/// See the documentation for [`Ipv4Addr::is_documentation`][IPv4] and
/// [`Ipv6Addr::is_documentation`][IPv6] for more details.
///
/// [IPv4]: ../../std/net/struct.Ipv4Addr.html#method.is_documentation
/// [IPv6]: ../../std/net/struct.Ipv6Addr.html#method.is_documentation
/// [`true`]: ../../std/primitive.bool.html
///
/// # Examples
///
/// ```
/// use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
/// use torment_core::ip::IpAddrExt;
///
/// assert_eq!(IpAddr::V4(Ipv4Addr::new(203, 0, 113, 6)).ext_is_documentation(), true);
/// assert_eq!(
/// IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0)).ext_is_documentation(),
/// true
/// );
/// ```
fn ext_is_documentation(&self) -> bool {
match self {
IpAddr::V4(ip) => ip.is_documentation(),
IpAddr::V6(ip) => ip.ext_is_documentation(),
}
}
}
pub trait Ipv4AddrExt {
fn ext_is_global(&self) -> bool;
fn ext_is_shared(&self) -> bool;
fn ext_is_ietf_protocol_assignment(&self) -> bool;
fn ext_is_benchmarking(&self) -> bool;
fn ext_is_reserved(&self) -> bool;
}
impl Ipv4AddrExt for Ipv4Addr {
/// Returns [`true`] if the address appears to be globally routable.
/// See [iana-ipv4-special-registry][ipv4-sr].
///
/// The following return false:
///
/// - private addresses (see [`is_private()`](#method.is_private))
/// - the loopback address (see [`is_loopback()`](#method.is_loopback))
/// - the link-local address (see [`is_link_local()`](#method.is_link_local))
/// - the broadcast address (see [`is_broadcast()`](#method.is_broadcast))
/// - addresses used for documentation (see [`is_documentation()`](#method.is_documentation))
/// - the unspecified address (see [`is_unspecified()`](#method.is_unspecified)), and the whole
/// 0.0.0.0/8 block
/// - addresses reserved for future protocols (see
/// [`is_ietf_protocol_assignment()`](#method.is_ietf_protocol_assignment), except
/// `192.0.0.9/32` and `192.0.0.10/32` which are globally routable
/// - addresses reserved for future use (see [`is_reserved()`](#method.is_reserved)
/// - addresses reserved for networking devices benchmarking (see
/// [`is_benchmarking`](#method.is_benchmarking))
///
/// [ipv4-sr]: https://www.iana.org/assignments/iana-ipv4-special-registry/iana-ipv4-special-registry.xhtml
/// [`true`]: ../../std/primitive.bool.html
///
/// # Examples
///
/// ```
/// use std::net::Ipv4Addr;
/// use torment_core::ip::Ipv4AddrExt;
///
/// // private addresses are not global
/// assert_eq!(Ipv4Addr::new(10, 254, 0, 0).ext_is_global(), false);
/// assert_eq!(Ipv4Addr::new(192, 168, 10, 65).ext_is_global(), false);
/// assert_eq!(Ipv4Addr::new(172, 16, 10, 65).ext_is_global(), false);
///
/// // the 0.0.0.0/8 block is not global
/// assert_eq!(Ipv4Addr::new(0, 1, 2, 3).ext_is_global(), false);
/// // in particular, the unspecified address is not global
/// assert_eq!(Ipv4Addr::new(0, 0, 0, 0).ext_is_global(), false);
///
/// // the loopback address is not global
/// assert_eq!(Ipv4Addr::new(127, 0, 0, 1).ext_is_global(), false);
///
/// // link local addresses are not global
/// assert_eq!(Ipv4Addr::new(169, 254, 45, 1).ext_is_global(), false);
///
/// // the broadcast address is not global
/// assert_eq!(Ipv4Addr::new(255, 255, 255, 255).ext_is_global(), false);
///
/// // the address space designated for documentation is not global
/// assert_eq!(Ipv4Addr::new(192, 0, 2, 255).ext_is_global(), false);
/// assert_eq!(Ipv4Addr::new(198, 51, 100, 65).ext_is_global(), false);
/// assert_eq!(Ipv4Addr::new(203, 0, 113, 6).ext_is_global(), false);
///
/// // shared addresses are not global
/// assert_eq!(Ipv4Addr::new(100, 100, 0, 0).ext_is_global(), false);
///
/// // addresses reserved for protocol assignment are not global
/// assert_eq!(Ipv4Addr::new(192, 0, 0, 0).ext_is_global(), false);
/// assert_eq!(Ipv4Addr::new(192, 0, 0, 255).ext_is_global(), false);
///
/// // addresses reserved for future use are not global
/// assert_eq!(Ipv4Addr::new(250, 10, 20, 30).ext_is_global(), false);
///
/// // addresses reserved for network devices benchmarking are not global
/// assert_eq!(Ipv4Addr::new(198, 18, 0, 0).ext_is_global(), false);
///
/// // All the other addresses are global
/// assert_eq!(Ipv4Addr::new(1, 1, 1, 1).ext_is_global(), true);
/// assert_eq!(Ipv4Addr::new(80, 9, 12, 3).ext_is_global(), true);
/// ```
fn ext_is_global(&self) -> bool {
// check if this address is 192.0.0.9 or 192.0.0.10. These addresses are the only two
// globally routable addresses in the 192.0.0.0/24 range.
if u32::from(*self) == 0xc0000009 || u32::from(*self) == 0xc000000a {
return true;
}
!self.is_private()
&& !self.is_loopback()
&& !self.is_link_local()
&& !self.is_broadcast()
&& !self.is_documentation()
&& !self.ext_is_shared()
&& !self.ext_is_ietf_protocol_assignment()
&& !self.ext_is_reserved()
&& !self.ext_is_benchmarking()
// Make sure the address is not in 0.0.0.0/8
&& self.octets()[0] != 0
}
/// Returns [`true`] if this address is part of the Shared Address Space defined in
/// [IETF RFC 6598] (`100.64.0.0/10`).
///
/// [IETF RFC 6598]: https://tools.ietf.org/html/rfc6598
/// [`true`]: ../../std/primitive.bool.html
///
/// # Examples
///
/// ```
/// use std::net::Ipv4Addr;
/// use torment_core::ip::Ipv4AddrExt;
///
/// assert_eq!(Ipv4Addr::new(100, 64, 0, 0).ext_is_shared(), true);
/// assert_eq!(Ipv4Addr::new(100, 127, 255, 255).ext_is_shared(), true);
/// assert_eq!(Ipv4Addr::new(100, 128, 0, 0).ext_is_shared(), false);
/// ```
fn ext_is_shared(&self) -> bool {
self.octets()[0] == 100 && (self.octets()[1] & 0b1100_0000 == 0b0100_0000)
}
/// Returns [`true`] if this address is part of `192.0.0.0/24`, which is reserved to
/// IANA for IETF protocol assignments, as documented in [IETF RFC 6890].
///
/// Note that parts of this block are in use:
///
/// - `192.0.0.8/32` is the "IPv4 dummy address" (see [IETF RFC 7600])
/// - `192.0.0.9/32` is the "Port Control Protocol Anycast" (see [IETF RFC 7723])
/// - `192.0.0.10/32` is used for NAT traversal (see [IETF RFC 8155])
///
/// [IETF RFC 6890]: https://tools.ietf.org/html/rfc6890
/// [IETF RFC 7600]: https://tools.ietf.org/html/rfc7600
/// [IETF RFC 7723]: https://tools.ietf.org/html/rfc7723
/// [IETF RFC 8155]: https://tools.ietf.org/html/rfc8155
/// [`true`]: ../../std/primitive.bool.html
///
/// # Examples
///
/// ```
/// use std::net::Ipv4Addr;
/// use torment_core::ip::Ipv4AddrExt;
///
/// assert_eq!(Ipv4Addr::new(192, 0, 0, 0).ext_is_ietf_protocol_assignment(), true);
/// assert_eq!(Ipv4Addr::new(192, 0, 0, 8).ext_is_ietf_protocol_assignment(), true);
/// assert_eq!(Ipv4Addr::new(192, 0, 0, 9).ext_is_ietf_protocol_assignment(), true);
/// assert_eq!(Ipv4Addr::new(192, 0, 0, 255).ext_is_ietf_protocol_assignment(), true);
/// assert_eq!(Ipv4Addr::new(192, 0, 1, 0).ext_is_ietf_protocol_assignment(), false);
/// assert_eq!(Ipv4Addr::new(191, 255, 255, 255).ext_is_ietf_protocol_assignment(), false);
/// ```
fn ext_is_ietf_protocol_assignment(&self) -> bool {
self.octets()[0] == 192 && self.octets()[1] == 0 && self.octets()[2] == 0
}
/// Returns [`true`] if this address part of the `198.18.0.0/15` range, which is reserved for
/// network devices benchmarking. This range is defined in [IETF RFC 2544] as `192.18.0.0`
/// through `198.19.255.255` but [errata 423] corrects it to `198.18.0.0/15`.
///
/// [IETF RFC 2544]: https://tools.ietf.org/html/rfc2544
/// [errata 423]: https://www.rfc-editor.org/errata/eid423
/// [`true`]: ../../std/primitive.bool.html
///
/// # Examples
///
/// ```
/// use std::net::Ipv4Addr;
/// use torment_core::ip::Ipv4AddrExt;
///
/// assert_eq!(Ipv4Addr::new(198, 17, 255, 255).ext_is_benchmarking(), false);
/// assert_eq!(Ipv4Addr::new(198, 18, 0, 0).ext_is_benchmarking(), true);
/// assert_eq!(Ipv4Addr::new(198, 19, 255, 255).ext_is_benchmarking(), true);
/// assert_eq!(Ipv4Addr::new(198, 20, 0, 0).ext_is_benchmarking(), false);
/// ```
fn ext_is_benchmarking(&self) -> bool {
self.octets()[0] == 198 && (self.octets()[1] & 0xfe) == 18
}
/// Returns [`true`] if this address is reserved by IANA for future use. [IETF RFC 1112]
/// defines the block of reserved addresses as `240.0.0.0/4`. This range normally includes the
/// broadcast address `255.255.255.255`, but this implementation explicitly excludes it, since
/// it is obviously not reserved for future use.
///
/// [IETF RFC 1112]: https://tools.ietf.org/html/rfc1112
/// [`true`]: ../../std/primitive.bool.html
///
/// # Warning
///
/// As IANA assigns new addresses, this method will be
/// updated. This may result in non-reserved addresses being
/// treated as reserved in code that relies on an outdated version
/// of this method.
///
/// # Examples
///
/// ```
/// use std::net::Ipv4Addr;
/// use torment_core::ip::Ipv4AddrExt;
///
/// assert_eq!(Ipv4Addr::new(240, 0, 0, 0).ext_is_reserved(), true);
/// assert_eq!(Ipv4Addr::new(255, 255, 255, 254).ext_is_reserved(), true);
///
/// assert_eq!(Ipv4Addr::new(239, 255, 255, 255).ext_is_reserved(), false);
/// // The broadcast address is not considered as reserved for future use by this implementation
/// assert_eq!(Ipv4Addr::new(255, 255, 255, 255).ext_is_reserved(), false);
/// ```
fn ext_is_reserved(&self) -> bool {
self.octets()[0] & 240 == 240 && !self.is_broadcast()
}
}
pub trait Ipv6AddrExt {
fn ext_is_global(&self) -> bool;
fn ext_is_unique_local(&self) -> bool;
fn ext_is_unicast_link_local_strict(&self) -> bool;
fn ext_is_unicast_link_local(&self) -> bool;
fn ext_is_unicast_site_local(&self) -> bool;
fn ext_is_unicast_global(&self) -> bool;
fn ext_is_documentation(&self) -> bool;
fn ext_multicast_scope(&self) -> Option<Ipv6MulticastScope>;
}
impl Ipv6AddrExt for Ipv6Addr {
/// Returns [`true`] if the address appears to be globally routable.
///
/// The following return [`false`]:
///
/// - the loopback address
/// - link-local and unique local unicast addresses
/// - interface-, link-, realm-, admin- and site-local multicast addresses
///
/// [`true`]: ../../std/primitive.bool.html
/// [`false`]: ../../std/primitive.bool.html
///
/// # Examples
///
/// ```
/// use std::net::Ipv6Addr;
/// use torment_core::ip::Ipv6AddrExt;
///
/// assert_eq!(Ipv6Addr::new(0, 0, 0, 0, 0, 0xffff, 0xc00a, 0x2ff).ext_is_global(), true);
/// assert_eq!(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 0x1).ext_is_global(), false);
/// assert_eq!(Ipv6Addr::new(0, 0, 0x1c9, 0, 0, 0xafc8, 0, 0x1).ext_is_global(), true);
/// ```
fn ext_is_global(&self) -> bool {
match Ipv6AddrExt::ext_multicast_scope(self) {
Some(Ipv6MulticastScope::Global) => true,
None => self.ext_is_unicast_global(),
_ => false,
}
}
/// Returns [`true`] if this is a unique local address (`fc00::/7`).
///
/// This property is defined in [IETF RFC 4193].
///
/// [IETF RFC 4193]: https://tools.ietf.org/html/rfc4193
/// [`true`]: ../../std/primitive.bool.html
///
/// # Examples
///
/// ```
/// use std::net::Ipv6Addr;
/// use torment_core::ip::Ipv6AddrExt;
///
/// assert_eq!(Ipv6Addr::new(0, 0, 0, 0, 0, 0xffff, 0xc00a, 0x2ff).ext_is_unique_local(), false);
/// assert_eq!(Ipv6Addr::new(0xfc02, 0, 0, 0, 0, 0, 0, 0).ext_is_unique_local(), true);
/// ```
fn ext_is_unique_local(&self) -> bool {
(self.segments()[0] & 0xfe00) == 0xfc00
}
/// Returns [`true`] if the address is a unicast link-local address (`fe80::/64`).
///
/// A common mis-conception is to think that "unicast link-local addresses start with
/// `fe80::`", but the [IETF RFC 4291] actually defines a stricter format for these addresses:
///
/// ```no_rust
/// | 10 |
/// | bits | 54 bits | 64 bits |
/// +----------+-------------------------+----------------------------+
/// |1111111010| 0 | interface ID |
/// +----------+-------------------------+----------------------------+
/// ```
///
/// This method validates the format defined in the RFC and won't recognize the following
/// addresses such as `fe80:0:0:1::` or `fe81::` as unicast link-local addresses for example.
/// If you need a less strict validation use [`is_unicast_link_local()`] instead.
///
/// # Examples
///
/// ```
/// use std::net::Ipv6Addr;
/// use torment_core::ip::Ipv6AddrExt;
///
/// let ip = Ipv6Addr::new(0xfe80, 0, 0, 0, 0, 0, 0, 0);
/// assert!(ip.ext_is_unicast_link_local_strict());
///
/// let ip = Ipv6Addr::new(0xfe80, 0, 0, 0, 0xffff, 0xffff, 0xffff, 0xffff);
/// assert!(ip.ext_is_unicast_link_local_strict());
///
/// let ip = Ipv6Addr::new(0xfe80, 0, 0, 1, 0, 0, 0, 0);
/// assert!(!ip.ext_is_unicast_link_local_strict());
/// assert!(ip.ext_is_unicast_link_local());
///
/// let ip = Ipv6Addr::new(0xfe81, 0, 0, 0, 0, 0, 0, 0);
/// assert!(!ip.ext_is_unicast_link_local_strict());
/// assert!(ip.ext_is_unicast_link_local());
/// ```
///
/// # See also
///
/// - [IETF RFC 4291 section 2.5.6]
/// - [RFC 4291 errata 4406]
/// - [`is_unicast_link_local()`]
///
/// [IETF RFC 4291]: https://tools.ietf.org/html/rfc4291
/// [IETF RFC 4291 section 2.5.6]: https://tools.ietf.org/html/rfc4291#section-2.5.6
/// [`true`]: ../../std/primitive.bool.html
/// [RFC 4291 errata 4406]: https://www.rfc-editor.org/errata/eid4406
/// [`is_unicast_link_local()`]: ../../std/net/struct.Ipv6Addr.html#method.is_unicast_link_local
///
fn ext_is_unicast_link_local_strict(&self) -> bool {
(self.segments()[0] & 0xffff) == 0xfe80
&& (self.segments()[1] & 0xffff) == 0
&& (self.segments()[2] & 0xffff) == 0
&& (self.segments()[3] & 0xffff) == 0
}
/// Returns [`true`] if the address is a unicast link-local address (`fe80::/10`).
///
/// This method returns [`true`] for addresses in the range reserved by [RFC 4291 section 2.4],
/// i.e. addresses with the following format:
///
/// ```no_rust
/// | 10 |
/// | bits | 54 bits | 64 bits |
/// +----------+-------------------------+----------------------------+
/// |1111111010| arbitratry value | interface ID |
/// +----------+-------------------------+----------------------------+
/// ```
///
/// As a result, this method consider addresses such as `fe80:0:0:1::` or `fe81::` to be
/// unicast link-local addresses, whereas [`is_unicast_link_local_strict()`] does not. If you
/// need a strict validation fully compliant with the RFC, use
/// [`is_unicast_link_local_strict()`].
///
/// # Examples
///
/// ```
/// use std::net::Ipv6Addr;
/// use torment_core::ip::Ipv6AddrExt;
///
/// let ip = Ipv6Addr::new(0xfe80, 0, 0, 0, 0, 0, 0, 0);
/// assert!(ip.ext_is_unicast_link_local());
///
/// let ip = Ipv6Addr::new(0xfe80, 0, 0, 0, 0xffff, 0xffff, 0xffff, 0xffff);
/// assert!(ip.ext_is_unicast_link_local());
///
/// let ip = Ipv6Addr::new(0xfe80, 0, 0, 1, 0, 0, 0, 0);
/// assert!(ip.ext_is_unicast_link_local());
/// assert!(!ip.ext_is_unicast_link_local_strict());
///
/// let ip = Ipv6Addr::new(0xfe81, 0, 0, 0, 0, 0, 0, 0);
/// assert!(ip.ext_is_unicast_link_local());
/// assert!(!ip.ext_is_unicast_link_local_strict());
/// ```
///
/// # See also
///
/// - [IETF RFC 4291 section 2.4]
/// - [RFC 4291 errata 4406]
///
/// [IETF RFC 4291 section 2.4]: https://tools.ietf.org/html/rfc4291#section-2.4
/// [`true`]: ../../std/primitive.bool.html
/// [RFC 4291 errata 4406]: https://www.rfc-editor.org/errata/eid4406
/// [`is_unicast_link_local_strict()`]: ../../std/net/struct.Ipv6Addr.html#method.is_unicast_link_local_strict
///
fn ext_is_unicast_link_local(&self) -> bool {
(self.segments()[0] & 0xffc0) == 0xfe80
}
/// Returns [`true`] if this is a deprecated unicast site-local address (fec0::/10). The
/// unicast site-local address format is defined in [RFC 4291 section 2.5.7] as:
///
/// ```no_rust
/// | 10 |
/// | bits | 54 bits | 64 bits |
/// +----------+-------------------------+----------------------------+
/// |1111111011| subnet ID | interface ID |
/// +----------+-------------------------+----------------------------+
/// ```
///
/// [`true`]: ../../std/primitive.bool.html
/// [RFC 4291 section 2.5.7]: https://tools.ietf.org/html/rfc4291#section-2.5.7
///
/// # Examples
///
/// ```
/// use std::net::Ipv6Addr;
/// use torment_core::ip::Ipv6AddrExt;
///
/// assert_eq!(
/// Ipv6Addr::new(0, 0, 0, 0, 0, 0xffff, 0xc00a, 0x2ff).ext_is_unicast_site_local(),
/// false
/// );
/// assert_eq!(Ipv6Addr::new(0xfec2, 0, 0, 0, 0, 0, 0, 0).ext_is_unicast_site_local(), true);
/// ```
///
/// # Warning
///
/// As per [RFC 3879], the whole `FEC0::/10` prefix is
/// deprecated. New software must not support site-local
/// addresses.
///
/// [RFC 3879]: https://tools.ietf.org/html/rfc3879
fn ext_is_unicast_site_local(&self) -> bool {
(self.segments()[0] & 0xffc0) == 0xfec0
}
/// Returns [`true`] if the address is a globally routable unicast address.
///
/// The following return false:
///
/// - the loopback address
/// - the link-local addresses
/// - unique local addresses
/// - the unspecified address
/// - the address range reserved for documentation
///
/// This method returns [`true`] for site-local addresses as per [RFC 4291 section 2.5.7]
///
/// ```no_rust
/// The special behavior of [the site-local unicast] prefix defined in [RFC3513] must no longer
/// be supported in new implementations (i.e., new implementations must treat this prefix as
/// Global Unicast).
/// ```
///
/// [`true`]: ../../std/primitive.bool.html
/// [RFC 4291 section 2.5.7]: https://tools.ietf.org/html/rfc4291#section-2.5.7
///
/// # Examples
///
/// ```
/// use std::net::Ipv6Addr;
/// use torment_core::ip::Ipv6AddrExt;
///
/// assert_eq!(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0).ext_is_unicast_global(), false);
/// assert_eq!(Ipv6Addr::new(0, 0, 0, 0, 0, 0xffff, 0xc00a, 0x2ff).ext_is_unicast_global(), true);
/// ```
fn ext_is_unicast_global(&self) -> bool {
!self.is_multicast()
&& !self.is_loopback()
&& !self.ext_is_unicast_link_local()
&& !self.ext_is_unique_local()
&& !self.is_unspecified()
&& !self.ext_is_documentation()
}
/// Returns [`true`] if this is an address reserved for documentation
/// (2001:db8::/32).
///
/// This property is defined in [IETF RFC 3849].
///
/// [IETF RFC 3849]: https://tools.ietf.org/html/rfc3849
/// [`true`]: ../../std/primitive.bool.html
///
/// # Examples
///
/// ```
/// use std::net::Ipv6Addr;
/// use torment_core::ip::Ipv6AddrExt;
///
/// assert_eq!(Ipv6Addr::new(0, 0, 0, 0, 0, 0xffff, 0xc00a, 0x2ff).ext_is_documentation(), false);
/// assert_eq!(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0).ext_is_documentation(), true);
/// ```
fn ext_is_documentation(&self) -> bool {
(self.segments()[0] == 0x2001) && (self.segments()[1] == 0xdb8)
}
/// Returns the address's multicast scope if the address is multicast.
///
/// # Examples
///
/// ```
/// use std::net::{Ipv6Addr};
/// use torment_core::ip::{Ipv6AddrExt, Ipv6MulticastScope};
///
/// assert_eq!(
/// Ipv6Addr::new(0xff0e, 0, 0, 0, 0, 0, 0, 0).ext_multicast_scope(),
/// Some(Ipv6MulticastScope::Global)
/// );
/// assert_eq!(Ipv6Addr::new(0, 0, 0, 0, 0, 0xffff, 0xc00a, 0x2ff).ext_multicast_scope(), None);
/// ```
fn ext_multicast_scope(&self) -> Option<Ipv6MulticastScope> {
if self.is_multicast() {
match self.segments()[0] & 0x000f {
1 => Some(Ipv6MulticastScope::InterfaceLocal),
2 => Some(Ipv6MulticastScope::LinkLocal),
3 => Some(Ipv6MulticastScope::RealmLocal),
4 => Some(Ipv6MulticastScope::AdminLocal),
5 => Some(Ipv6MulticastScope::SiteLocal),
8 => Some(Ipv6MulticastScope::OrganizationLocal),
14 => Some(Ipv6MulticastScope::Global),
_ => None,
}
} else {
None
}
}
}

116
torment-core/src/lib.rs
Unescape Escape View File

@ -0,0 +1,116 @@
#![allow(dead_code)]
use crate::infohash::v1::U160;
use crate::infohash::InfoHashCapable;
use serde_derive::{Deserialize, Serialize};
use std::convert::TryInto;
use std::fmt::{Display, Formatter};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
pub mod infohash;
pub mod ip;
pub mod metainfo;
pub mod utils;
pub trait CompactContact: Sized {
fn to_compact_contact(&self) -> Vec<u8>;
fn from_compact_contact<T: AsRef<[u8]>>(input: T) -> Result<Self, ParsingError>;
}
impl CompactContact for SocketAddr {
fn to_compact_contact(&self) -> Vec<u8> {
match self.ip() {
IpAddr::V4(ref v4) => {
let mut compact = [0u8; 6];
compact[..4].copy_from_slice(&v4.octets());
compact[4..].copy_from_slice(&self.port().to_be_bytes()[..]);
compact.to_vec()
}
IpAddr::V6(ref v6) => {
let mut compact = [0u8; 18];
compact[..16].copy_from_slice(&v6.octets());
compact[16..].copy_from_slice(&self.port().to_be_bytes());
compact.to_vec()
}
}
}
fn from_compact_contact<T: AsRef<[u8]>>(input: T) -> Result<Self, ParsingError> {
let b = input.as_ref();
if b.len() == 6 {
let ipv4: [u8; 4] = b[..4].try_into().unwrap();
Ok(SocketAddr::new(
Ipv4Addr::from(ipv4).into(),
u16::from_be_bytes(b[4..].try_into().unwrap()),
))
} else if b.len() == 18 {
let ipv6: [u8; 16] = b[..16].try_into().unwrap();
Ok(SocketAddr::new(
Ipv6Addr::from(ipv6).into(),
u16::from_be_bytes(b[16..].try_into().unwrap()),
))
} else {
Err(ParsingError)
}
}
}
#[derive(Debug)]
pub struct ParsingError;
impl std::error::Error for ParsingError {}
impl Display for ParsingError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "there was a parsing error.")
}
}
#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
pub struct ContactInfo {
pub id: U160,
pub contact: SocketAddr,
}
impl ContactInfo {
pub fn to_bytes(&self) -> Vec<u8> {
let contact = self.contact.to_compact_contact();
let mut bytes = self.id.to_bytes().to_vec();
bytes.extend(contact);
bytes
}
pub fn from_bytes<T: AsRef<[u8]>>(input: T) -> Result<ContactInfo, ParsingError> {
let b = input.as_ref();
if b.len() == 26 || b.len() == 38 {
Ok(ContactInfo {
id: U160::from_bytes(&b[..20]).map_err(|_| ParsingError)?,
contact: SocketAddr::from_compact_contact(&b[20..])?,
})
} else {
Err(ParsingError)
}
}
}
#[derive(Debug, Default)]
pub struct PeerStorage {}
#[derive(Ord, PartialOrd, Eq, PartialEq, Copy, Clone)]
pub enum LookupFilter {
IPv6,
IPv4,
All,
}
impl PeerStorage {
pub fn new() -> PeerStorage {
PeerStorage {}
}
pub fn add_peers(&mut self, _info_hash: U160, _peers: Vec<SocketAddr>) {}
pub fn get_peers(&self, _info_hash: U160, _filter: LookupFilter) -> Vec<SocketAddr> {
vec![]
}
}

236
torment-core/src/metainfo.rs
Unescape Escape View File

@ -0,0 +1,236 @@
use crate::infohash::v1::U160;
use crate::infohash::InfoHashCapable;
use bytes::Bytes;
use num_traits::cast::ToPrimitive;
use ring::digest::{digest, SHA1_FOR_LEGACY_USE_ONLY};
use serde::export::fmt::Debug;
use serde::export::Formatter;
use std::collections::HashSet;
use std::fmt::Display;
use std::string::FromUtf8Error;
use torment_bencode::{BencodeError, BencodeValue};
use url::Url;
#[derive(Debug)]
pub enum MetaInfoParsingError {
DecodingFailure(BencodeError),
Utf8Error(FromUtf8Error),
WrongType,
MissingEntry,
}
impl From<BencodeError> for MetaInfoParsingError {
fn from(err: BencodeError) -> Self {
MetaInfoParsingError::DecodingFailure(err)
}
}
impl From<std::string::FromUtf8Error> for MetaInfoParsingError {
fn from(err: FromUtf8Error) -> Self {
MetaInfoParsingError::Utf8Error(err)
}
}
impl Display for MetaInfoParsingError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "{:?}", self)
}
}
pub struct Torrent {
name: String,
announce_list: Vec<HashSet<Url>>,
info: MetaInfo,
bencoded: BencodeValue,
}
impl Debug for Torrent {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Torrent")
.field("name", &self.name)
.field("announce_list", &self.announce_list)
.field("info", &self.info)
.finish()
}
}
impl Torrent {
pub fn from_bytes<T: AsRef<[u8]>>(
bytes: T,
name: Option<String>,
) -> Result<Self, MetaInfoParsingError> {
let b = Bytes::from(bytes.as_ref().to_vec());
let bencoded = BencodeValue::decode(b.clone())?;
let dict = bencoded.dict().ok_or(MetaInfoParsingError::WrongType)?;
let info = dict
.get(b"info")
.ok_or(MetaInfoParsingError::MissingEntry)?
.dict()
.ok_or(MetaInfoParsingError::WrongType)?;
let data = info.buffer();
let info = MetaInfo::from_bytes(data)?;
let name = dict
.get(b"name")
.and_then(|name| name.bytes())
.map(|str| String::from_utf8(str.to_vec()))
.transpose()?
.or(name)
.unwrap_or(info.info_hash.to_string());
let announce_list = if let Some(announce) = dict.get(b"announce-list") {
announce
.list()
.map(|list| -> Result<Vec<HashSet<Url>>, MetaInfoParsingError> {
let mut items = vec![];
for i in 0..list.len() {
let mut urls = HashSet::new();
let item = list.get(i).unwrap();
let url_list = item.list().unwrap();
for url_i in 0..url_list.len() {
urls.insert(
Url::parse(
&*url_list.get(url_i).unwrap().string().transpose()?.unwrap(),
)
.unwrap(),
);
}
items.push(urls);
}
Ok(items)
})
.transpose()?
.unwrap_or(vec![])
} else {
vec![]
};
Ok(Torrent {
name,
announce_list,
info,
bencoded,
})
}
}
pub struct MetaInfo {
info_hash: U160,
bencoded: BencodeValue,
piece_length: usize,
pieces: Bytes,
object: MetaInfoObject,
}
#[derive(Clone, Debug)]
pub enum MetaInfoObject {
Files(Vec<MetaFile>),
File(usize),
}
impl Debug for MetaInfo {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MetaInfo")
.field("info_hash", &self.info_hash)
.field("piece_length", &self.piece_length)
.field("object", &self.object)
.finish()
}
}
impl MetaInfo {
pub fn from_bytes<T: AsRef<[u8]>>(bytes: T) -> Result<Self, MetaInfoParsingError> {
let b = Bytes::from(bytes.as_ref().to_vec());
let bencoded = BencodeValue::decode(b)?;
let dict = bencoded.dict().ok_or(MetaInfoParsingError::WrongType)?;
let piece_length = dict
.get(b"piece length")
.ok_or(MetaInfoParsingError::MissingEntry)?
.int()
.ok_or(MetaInfoParsingError::WrongType)?
.to_usize()
.ok_or(MetaInfoParsingError::WrongType)?;
let object = if let Some(files) = dict.get(b"files") {
let mut meta_files = vec![];
let files = files.list().ok_or(MetaInfoParsingError::WrongType)?;
for file_i in 0..files.len() {
let file = files.get(file_i).unwrap();
meta_files.push(MetaFile::from_bencoded(file)?);
}
MetaInfoObject::Files(meta_files)
} else {
MetaInfoObject::File(
dict.get(b"length")
.ok_or(MetaInfoParsingError::MissingEntry)?
.int()
.ok_or(MetaInfoParsingError::WrongType)?
.to_usize()
.ok_or(MetaInfoParsingError::WrongType)?,
)
};
let pieces = dict
.get(b"pieces")
.ok_or(MetaInfoParsingError::MissingEntry)?
.bytes()
.ok_or(MetaInfoParsingError::WrongType)?;
let info_hash =
U160::from_bytes(digest(&SHA1_FOR_LEGACY_USE_ONLY, dict.buffer().as_ref())).unwrap();
Ok(MetaInfo {
info_hash,
bencoded,
piece_length,
pieces,
object,
})
}
}
#[derive(Debug, Clone)]
pub struct MetaFile {
length: usize,
path: Vec<String>,
}
impl MetaFile {
fn from_bencoded(bencoded: &BencodeValue) -> Result<Self, MetaInfoParsingError> {
let dict = bencoded.dict().ok_or(MetaInfoParsingError::WrongType)?;
let path = dict
.get(b"path")
.ok_or(MetaInfoParsingError::MissingEntry)?
.list()
.ok_or(MetaInfoParsingError::WrongType)?;
let mut path_ele = vec![];
for path_i in 0..path.len() {
let path_item = path.get(path_i).unwrap();
path_ele.push(
path_item
.string()
.ok_or(MetaInfoParsingError::WrongType)??,
)
}
Ok(MetaFile {
path: path_ele,
length: dict
.get(b"length")
.ok_or(MetaInfoParsingError::MissingEntry)?
.int()
.ok_or(MetaInfoParsingError::WrongType)?
.to_usize()
.ok_or(MetaInfoParsingError::WrongType)?,
})
}
}

362
torment-core/src/utils.rs
Unescape Escape View File

@ -0,0 +1,362 @@
use std::collections::{BTreeMap, BTreeSet};
use std::fmt::{Debug, Formatter};
use std::time::Instant;
#[derive(Default)]
pub struct EphemeralMap<K: Ord + Clone, V> {
entry_map: BTreeMap<K, EphemeralNode<V>>,
expiry_map: BTreeMap<Instant, BTreeSet<K>>,
}
impl<K: Ord + Debug + Clone, V: Debug> Debug for EphemeralMap<K, V> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_map().entries(self.entry_map.iter()).finish()
}
}
impl<K: Ord + Clone + Debug, V> EphemeralMap<K, V> {
pub fn new() -> Self {
EphemeralMap {
entry_map: Default::default(),
expiry_map: Default::default(),
}
}
pub fn clean(&mut self) {
let valid = self.expiry_map.split_off(&Instant::now());
let expired = std::mem::replace(&mut self.expiry_map, valid);
for (_, keys) in expired {
for key in keys {
self.entry_map.remove(&key);
}
}
}
pub fn insert(&mut self, key: K, value: V, expires_at: Instant) -> Option<EphemeralNode<V>> {
let node = EphemeralNode { value, expires_at };
let removed = self.remove(&key);
self.add_expiry(key.clone(), expires_at);
self.entry_map.insert(key, node);
removed
}
pub fn remove(&mut self, key: &K) -> Option<EphemeralNode<V>> {
if let Some(node) = self.entry_map.remove(key) {
self.remove_expiry(key, node.expires_at);
Some(node)
} else {
None
}
}
pub fn remove_value(&mut self, key: &K) -> Option<V> {
self.remove(key).map(|x| x.value)
}
fn add_expiry(&mut self, key: K, expires_at: Instant) {
self.expiry_map
.entry(expires_at)
.or_insert_with(|| BTreeSet::new())
.insert(key);
}
fn remove_expiry(&mut self, key: &K, expires_at: Instant) {
let remove_node = if let Some(item) = self.expiry_map.get_mut(&expires_at) {
item.remove(&key);
item.len() == 0
} else {
false
};
if remove_node {
self.expiry_map.remove(&expires_at);
}
}
pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
self.entry_map.iter().filter_map(|(key, entry)| {
if entry.is_expired() {
None
} else {
Some((key, &entry.value))
}
})
}
pub fn iter_mut(&mut self) -> impl Iterator<Item = (&K, &mut V)> {
self.entry_map.iter_mut().filter_map(|(key, entry)| {
if entry.is_expired() {
None
} else {
Some((key, &mut entry.value))
}
})
}
pub fn len(&self) -> usize {
self.iter().count()
}
pub fn contains_key(&self, key: &K) -> bool {
self.get(key).is_some()
}
pub fn get(&self, key: &K) -> Option<&V> {
self.get_node(key).map(|x| &x.value)
}
pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
self.get_node_mut(key).map(|x| &mut x.value)
}
pub fn get_node(&self, key: &K) -> Option<&EphemeralNode<V>> {
let node = if let Some(node) = self.entry_map.get(&key) {
node
} else {
return None;
};
if node.is_expired() {
return None;
}
Some(node)
}
fn get_node_mut(&mut self, key: &K) -> Option<&mut EphemeralNode<V>> {
let node = if let Some(node) = self.entry_map.get_mut(&key) {
node
} else {
return None;
};
if node.is_expired() {
return None;
}
Some(node)
}
pub fn update_expiry(&mut self, key: &K, expires_at: Instant) {
let node = if let Some(node) = self.entry_map.get_mut(key) {
node
} else {
return;
};
let old_expiry = node.expires_at;
node.expires_at = expires_at;
self.remove_expiry(key, old_expiry);
self.add_expiry(key.clone(), expires_at);
}
pub fn entry(&mut self, key: K) -> EphemeralEntry<'_, K, V> {
EphemeralEntry { map: self, key }
}
}
pub struct EphemeralNode<V> {
expires_at: Instant,
value: V,
}
impl<V: Debug> Debug for EphemeralNode<V> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("EphemeralNode")
.field("expires_at", &self.expires_at)
.field("value", &self.value)
.finish()
}
}
impl<V> EphemeralNode<V> {
pub fn is_expired(&self) -> bool {
self.expires_at <= Instant::now()
}
pub fn expires_at(&self) -> Instant {
self.expires_at
}
pub fn value(&self) -> &V {
&self.value
}
pub fn value_mut(&mut self) -> &mut V {
&mut self.value
}
pub fn into_value(self) -> V {
self.value
}
}
pub struct EphemeralEntry<'map, K: Ord + Clone, V> {
map: &'map mut EphemeralMap<K, V>,
key: K,
}
impl<'map, K: Ord + Clone + Debug, V> EphemeralEntry<'map, K, V> {
pub fn or_insert_with<F: FnOnce() -> V>(self, expires_at: Instant, f: F) -> &'map mut V {
self.or_insert(expires_at, f())
}
pub fn or_insert(self, expires_at: Instant, default: V) -> &'map mut V {
self.map.insert(self.key.clone(), default, expires_at);
self.map.get_mut(&self.key).unwrap()
}
pub fn and_modify<F: FnOnce(&mut V)>(self, f: F) -> Self {
if let Some(item) = self.map.get_mut(&self.key) {
f(item)
}
self
}
}
impl<'map, K: Ord + Clone + Debug, V: Default> EphemeralEntry<'map, K, V> {
pub fn or_insert_default(self, expires_at: Instant) -> &'map mut V {
self.or_insert(expires_at, Default::default())
}
}
#[derive(Default)]
pub struct EphemeralSet<K: Clone + Ord>(EphemeralMap<K, ()>);
impl<K: Debug + Clone + Ord> Debug for EphemeralSet<K> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_map()
.entries(
self.0
.entry_map
.iter()
.map(|(key, node)| (key, node.expires_at)),
)
.finish()
}
}
impl<K: Clone + Ord + Debug> EphemeralSet<K> {
pub fn new() -> Self {
EphemeralSet(EphemeralMap::new())
}
pub fn clean(&mut self) {
self.0.clean()
}
pub fn insert(&mut self, key: K, expires_at: Instant) -> bool {
self.0.insert(key, (), expires_at).is_none()
}
pub fn remove(&mut self, key: &K) -> bool {
self.0.remove(key).is_some()
}
pub fn contains(&self, key: &K) -> bool {
self.0.get_node(key).is_some()
}
pub fn remove_first(&mut self) -> Option<K> {
loop {
let first_key = self.0.entry_map.keys().next().cloned();
let item = if let Some(key) = first_key {
key
} else {
break;
};
if self.contains(&item) {
self.remove(&item);
return Some(item);
}
}
None
}
}
#[derive(Default)]
pub struct EphemeralQueueMap<K: Clone + Ord, V>(EphemeralMap<K, Vec<EphemeralNode<V>>>);
impl<K: Debug + Clone + Ord, V: Debug> Debug for EphemeralQueueMap<K, V> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_map().entries(self.0.entry_map.iter()).finish()
}
}
impl<K: Clone + Ord + Debug, V> EphemeralQueueMap<K, V> {
pub fn clean(&mut self) {
self.0.clean();
}
pub fn insert(&mut self, key: K, value: V, expires_at: Instant) {
let update_expiry = if let Some(node) = self.0.get_node_mut(&key) {
let mut old = std::mem::replace(&mut node.value, vec![]);
old.push(EphemeralNode { value, expires_at });
while let Some(old_node) = old.pop() {
if old_node.is_expired() {
continue;
}
node.value.push(old_node)
}
node.expires_at < expires_at
} else {
self.0.insert(
key.clone(),
vec![EphemeralNode { value, expires_at }],
expires_at,
);
false
};
if update_expiry {
self.0.update_expiry(&key, expires_at);
}
}
pub fn remove(&mut self, key: &K) {
self.0.remove(key);
}
pub fn pop(&mut self, key: &K) -> Option<V> {
if let Some(node) = self.0.get_node_mut(key) {
let old_values = std::mem::replace(&mut node.value, vec![]);
let mut selected = None;
let mut expiry: Option<Instant> = None;
for value in old_values {
if value.is_expired() {
continue;
}
if selected.is_none() {
selected = Some(value);
continue;
}
if *expiry.get_or_insert(value.expires_at) > value.expires_at {
expiry = Some(value.expires_at)
}
node.value.push(value)
}
if !node.value.is_empty() {
self.0.update_expiry(key, expiry.unwrap());
} else {
self.0.remove(key);
}
selected.map(|x| x.value)
} else {
None
}
}
}

19
torment-dht-node/Cargo.toml
Unescape Escape View File

@ -0,0 +1,19 @@
[package]
name = "torment-dht-node"
version = "0.1.0"
authors = ["eater <=@eater.me>"]
edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
torment-core = { path = "../torment-core" }
torment-dht = { path = "../torment-dht" }
futures = "0.3.5"
rand = "0.7.3"
async-std = { version = "1.6.3", features = ["attributes"] }
chrono = "0.4.15"
bytes = "0.5.6"
bson = "1.1.0"
serde = "1.0.115"
serde_derive = "1.0.115"

81
torment-dht-node/src/main.rs
Unescape Escape View File

@ -0,0 +1,81 @@
#![recursion_limit = "512"]
use async_std::net::{IpAddr, UdpSocket};
use futures::select;
use futures::FutureExt;
use std::net::SocketAddr;
use std::pin::Pin;
use std::str::FromStr;
use std::time::{Duration, Instant};
use torment_dht::host_node::HostNode;
use torment_dht::krpc::{FromBencode, Message, ToBencode};
async fn recv_from(socket: &UdpSocket) -> std::io::Result<(Vec<u8>, SocketAddr)> {
let mut buffer = [0u8; 2048];
socket
.recv_from(&mut buffer)
.await
.map(|(len, addr)| (buffer[..len].to_vec(), addr))
}
#[async_std::main]
async fn main() {
let socket = UdpSocket::bind(SocketAddr::from_str("[::]:50002").unwrap())
.await
.unwrap();
let mut node = HostNode::new(Default::default(), Some(50002));
node.add_bootstrap(SocketAddr::from_str("67.215.246.10:6881").unwrap(), None);
node.add_bootstrap(
SocketAddr::from_str("[2001:41d0:c:5ac:5::1]:6881").unwrap(),
None,
);
let mut fut_holder = recv_from(&socket).fuse();
let mut socket_fut = unsafe { Pin::new_unchecked(&mut fut_holder) };
let mut fut_sleep_holder = async_std::task::sleep(Duration::from_secs(10)).fuse();
let mut event_fut = unsafe { Pin::new_unchecked(&mut fut_sleep_holder) };
loop {
while let Some((message, to)) = node.next() {
println!("| {} <= {:?}", to, message);
socket
.send_to(&message.to_bencode().unwrap(), to)
.await
.unwrap();
}
select! {
res = socket_fut => {
if let Ok((msg, from)) = res {
let from = if let IpAddr::V6(ipv6) = from.ip() {
if let Some(ipv4) = ipv6.to_ipv4() {
SocketAddr::new(IpAddr::V4(ipv4), from.port())
} else {
from
}
} else {
from
};
match Message::from_bencode(&msg) {
Ok(msg) => {
println!("| {} => {:?}", from, msg);
node.process(msg, from);
},
Err(err) => eprintln!("{} => Failed parsing UDP message: {}", from, err),
}
}
fut_holder = recv_from(&socket).fuse();
socket_fut = unsafe { Pin::new_unchecked(&mut fut_holder) };
},
_ = event_fut => {
println!("| Housekeeping");
let instant = Instant::now();
node.housekeeping();
println!("| Housekeeping took {:?} (seen nodes: {}, table[ipv4={}, ipv6={}])", Instant::now() - instant, node.num_tracking_nodes(), node.num_ipv4_table_nodes(), node.num_ipv6_table_nodes());
fut_sleep_holder = async_std::task::sleep(Duration::from_secs(10)).fuse();
event_fut = unsafe { Pin::new_unchecked(&mut fut_sleep_holder) };
},
};
}
}

14
torment-dht/Cargo.toml
Unescape Escape View File

@ -0,0 +1,14 @@
[package]
name = "torment-dht"
version = "0.1.0"
authors = ["eater <=@eater.me>"]
edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
torment-core = { path = "../torment-core" }
bendy = "0.3.2"
serde = "1.0.115"
serde_derive = "1.0.115"
rand = "0.7.3"

797
torment-dht/src/host_node.rs
Unescape Escape View File

@ -0,0 +1,797 @@
use crate::krpc::{Message, MessageBody, RequestBody, ResponseBody, Want};
use crate::Table;
use rand::random;
use serde_derive::{Deserialize, Serialize};
use std::collections::{HashSet, VecDeque};
use std::convert::TryInto;
use std::net::{IpAddr, SocketAddr};
use std::ops::Add;
use std::rc::Rc;
use std::sync::RwLock;
use std::time::{Duration, Instant};
use torment_core::infohash::v1::U160;
use torment_core::ip::IpAddrExt;
use torment_core::utils::{EphemeralMap, EphemeralSet};
use torment_core::{ContactInfo, LookupFilter, PeerStorage};
#[derive(Debug)]
enum OpenRequest {
FindNode { target: U160 },
GetPeers { info_hash: U160 },
Ping,
Announce { info_hash: U160 },
}
const BOOTSTRAP_INTERVAL: Duration = Duration::from_secs(60);
#[derive(Debug)]
pub struct HostNode {
id: U160,
port: Option<u16>,
bootstrap_nodes: HashSet<(SocketAddr, Option<Want>)>,
last_bootstrap: Option<Instant>,
ipv4_table: Table,
ipv6_table: Table,
nodes: EphemeralMap<(U160, IpAddr, u16), PeerNode>,
queue: VecDeque<(Message, SocketAddr)>,
peer_storage: Rc<RwLock<PeerStorage>>,
}
#[derive(Debug)]
struct PeerNode {
id: U160,
addr: SocketAddr,
transactions: EphemeralMap<u16, OpenRequest>,
pre_announced: EphemeralSet<U160>,
last_activity: Instant,
tokens: EphemeralSet<u64>,
received_tokens: EphemeralSet<Vec<u8>>,
transaction_id: u16,
}
#[derive(Serialize, Deserialize)]
pub struct PersistentState {
id: U160,
nodes: Vec<ContactInfo>,
nodes6: Vec<ContactInfo>,
}
fn validate_ip<T: Into<IpAddr>>(ip: T) -> bool {
IpAddrExt::ext_is_global(&ip.into())
}
impl PeerNode {
fn clean(&mut self) {
self.transactions.clean();
self.tokens.clean();
self.received_tokens.clean();
self.pre_announced.clean();
}
fn contact_info(&self) -> ContactInfo {
ContactInfo {
id: self.id,
contact: self.addr,
}
}
}
impl PeerNode {
fn create<T: Into<SocketAddr>>(id: U160, addr: T) -> PeerNode {
PeerNode {
id,
addr: addr.into(),
transactions: EphemeralMap::new(),
pre_announced: Default::default(),
last_activity: Instant::now(),
tokens: Default::default(),
received_tokens: Default::default(),
transaction_id: 0,
}
}
}
impl HostNode {
pub fn new(peer_storage: Rc<RwLock<PeerStorage>>, port: Option<u16>) -> HostNode {
let id = U160::random();
HostNode {
id,
port,
bootstrap_nodes: Default::default(),
last_bootstrap: None,
ipv4_table: Table::new_with_id(id),
ipv6_table: Table::new_with_id(id),
nodes: EphemeralMap::new(),
queue: Default::default(),
peer_storage,
}
}
pub fn from_persistent_state(
peer_storage: Rc<RwLock<PeerStorage>>,
port: Option<u16>,
state: PersistentState,
) -> HostNode {
let mut host_node = HostNode {
id: state.id,
port,
bootstrap_nodes: Default::default(),
last_bootstrap: None,
ipv4_table: Table::new_with_id(state.id),
ipv6_table: Table::new_with_id(state.id),
nodes: EphemeralMap::new(),
queue: Default::default(),
peer_storage,
};
for node in state.nodes {
host_node
.ipv4_table
.add_node(node.id, node.contact.ip(), node.contact.port());
}
for node in state.nodes6 {
host_node
.ipv6_table
.add_node(node.id, node.contact.ip(), node.contact.port());
}
host_node.update_tables();
host_node
}
fn get_next_t(&mut self, contact_info: &ContactInfo) -> [u8; 2] {
let peer_node = self.get_peer_node(contact_info.id, contact_info.contact);
let old_id = peer_node.transaction_id;
peer_node.transaction_id = peer_node.transaction_id.overflowing_add(1).0;
old_id.to_ne_bytes()
}
pub fn process(&mut self, message: Message, from: SocketAddr) {
match message.body {
MessageBody::Request(ref req) => self.handle_request(req, message.transaction_id, from),
MessageBody::Response(ref resp) => {
self.handle_response(resp, message.transaction_id, from)
}
_ => {}
}
}
pub fn add_bootstrap(&mut self, socket_addr: SocketAddr, want: Option<Want>) {
self.bootstrap_nodes.insert((socket_addr, want));
}
fn get_k_closest_nodes(&self, id: U160, want: Option<Want>) -> Vec<ContactInfo> {
let mut nodes = vec![];
if want.unwrap_or(Want::N4) == Want::N4 {
nodes.extend(self.ipv4_table.find_k_closest_nodes(id, None))
}
if want.unwrap_or(Want::N6) == Want::N6 {
nodes.extend(self.ipv6_table.find_k_closest_nodes(id, None))
}
nodes
}
fn find_node(&mut self, target: U160, want: Option<Want>) {
if want.unwrap_or(Want::N4) == Want::N4 && self.ipv4_table.find_node(target).is_some() {
// noop
return;
}
if want.unwrap_or(Want::N6) == Want::N6 && self.ipv6_table.find_node(target).is_some() {
// noop
return;
}
for node in self.get_k_closest_nodes(target, want) {
let t = self.get_next_t(&node);
self.find_node_at(t, node.contact, target, want);
self.get_peer_node(node.id, node.contact)
.transactions
.insert(
u16::from_ne_bytes(t),
OpenRequest::FindNode { target },
Instant::now().add(Duration::from_secs(5 * 60)),
);
self.notify_table_of_outgoing_activity(&node);
}
}
fn find_node_at(
&mut self,
t: [u8; 2],
socket_addr: SocketAddr,
target: U160,
want: Option<Want>,
) {
let message = Message::request_find_node(self.id, t, target, want);
self.queue.push_back((message, socket_addr))
}
fn get_peers(&mut self, info_hash: U160, want: Option<Want>) {
let nodes = self.get_k_closest_nodes(info_hash, want);
for node in nodes {
self.get_peers_from(&node, info_hash, want)
}
}
fn get_peers_from(&mut self, node: &ContactInfo, info_hash: U160, want: Option<Want>) {
let message = Message::request_get_peers(self.id, self.get_next_t(&node), info_hash, want);
self.get_peer_node(node.id, node.contact)
.transactions
.insert(
u16::from_ne_bytes(message.transaction_id),
OpenRequest::GetPeers { info_hash },
Instant::now().add(Duration::from_secs(5 * 60)),
);
self.notify_table_of_outgoing_activity(node);
self.queue.push_back((message, node.contact))
}
fn notify_table_of_outgoing_activity(&mut self, node: &ContactInfo) {
if node.contact.is_ipv4() {
self.ipv4_table.handle_outgoing_activity(
node.id,
node.contact.ip(),
node.contact.port(),
);
} else {
self.ipv6_table.handle_outgoing_activity(
node.id,
node.contact.ip(),
node.contact.port(),
);
}
}
fn announce_peer_to(&mut self, node: &ContactInfo, info_hash: U160, token: Vec<u8>) {
let message = Message::request_announce_peer(
self.id,
self.get_next_t(&node),
info_hash,
token,
self.port.unwrap_or(0),
self.port.is_none(),
);
self.get_peer_node(node.id, node.contact)
.transactions
.insert(
u16::from_ne_bytes(message.transaction_id),
OpenRequest::Announce { info_hash },
Instant::now().add(Duration::from_secs(5 * 60)),
);
self.notify_table_of_outgoing_activity(node);
self.queue.push_back((message, node.contact))
}
pub fn ping(&mut self, node: &ContactInfo) {
let msg = Message::request_ping(self.id, self.get_next_t(node));
self.get_peer_node(node.id, node.contact)
.transactions
.insert(
u16::from_ne_bytes(msg.transaction_id),
OpenRequest::Ping,
Instant::now().add(Duration::from_secs(5 * 60)),
);
self.notify_table_of_outgoing_activity(node);
self.queue.push_back((msg, node.contact))
}
fn announce_peer(&mut self, info_hash: U160) {
let mut announce_to = vec![];
let mut peers_from = vec![];
for (_, details) in &mut self.nodes.iter_mut() {
if let Some(token) = details.received_tokens.remove_first() {
announce_to.push((details.contact_info(), token))
} else {
details
.pre_announced
.insert(info_hash, Instant::now().add(Duration::from_secs(5 * 60)));
peers_from.push(details.contact_info());
}
}
for (node, token) in announce_to {
self.announce_peer_to(&node, info_hash, token)
}
for node in peers_from {
self.get_peers_from(&node, info_hash, None)
}
}
fn has_node(&self, node: &ContactInfo) -> bool {
self.ipv4_table.has_node(node) || self.ipv6_table.has_node(node)
}
fn knows_node(&self, node: &ContactInfo) -> bool {
self.nodes
.contains_key(&(node.id, node.contact.ip(), node.contact.port()))
}
fn handle_response(&mut self, response: &ResponseBody, t: [u8; 2], from: SocketAddr) {
let from_id = response.node_id();
if from_id == self.id {
// In this house we don't talk to imposters
return;
}
match response {
ResponseBody::GetPeers(get_peers) => {
let node = self.get_peer_node(from_id, from);
let info_hash = if let Some(OpenRequest::GetPeers { info_hash }) =
node.transactions.remove_value(&u16::from_ne_bytes(t))
{
info_hash
} else {
return;
};
let is_pre_announce = if node.pre_announced.remove(&info_hash) {
true
} else {
node.received_tokens.insert(
get_peers.token.clone(),
Instant::now().add(Duration::from_secs(5 * 60)),
);
false
};
if is_pre_announce {
let contact_info = node.contact_info();
self.announce_peer_to(&contact_info, info_hash, get_peers.token.clone());
return;
}
if get_peers.values.len() > 0 {
self.peer_storage
.write()
.unwrap()
.add_peers(info_hash, get_peers.values.clone());
}
for node in &get_peers.nodes.nodes {
if !validate_ip(node.contact.ip()) {
continue;
}
if self.has_node(node) {
continue;
}
self.ipv4_table
.add_node(node.id, node.contact.ip(), node.contact.port());
self.get_peers_from(node, info_hash, None);
}
for node in &get_peers.nodes.nodes6 {
if !validate_ip(node.contact.ip()) {
continue;
}
if self.has_node(node) {
continue;
}
self.ipv6_table
.add_node(node.id, node.contact.ip(), node.contact.port());
self.get_peers_from(node, info_hash, None);
}
}
ResponseBody::FindNode(find_node) => {
let node = self.get_peer_node(response.node_id(), from);
// We don't really care if we had a transaction or not, since at bootstrap we didn't know the nodes ID yet
let request = node.transactions.remove_value(&u16::from_ne_bytes(t));
let target = if let Some(OpenRequest::FindNode { target }) = request {
target
} else {
self.id
};
let poke_if_closer = |host: &mut HostNode, node: &ContactInfo| {
if !host.knows_node(node) && (node.id ^ host.id) < (from_id ^ host.id) {
let t = host.get_next_t(node);
host.find_node_at(t, node.contact, target, None);
}
};
for node in &find_node.nodes {
if !validate_ip(node.contact.ip()) {
continue;
}
if self
.ipv4_table
.add_node(node.id, node.contact.ip(), node.contact.port())
{
poke_if_closer(self, node);
}
}
for node in &find_node.nodes6 {
if !validate_ip(node.contact.ip()) {
continue;
}
if self
.ipv6_table
.add_node(node.id, node.contact.ip(), node.contact.port())
{
poke_if_closer(self, node);
}
}
}
ResponseBody::Empty(id) => {
let node = self.get_peer_node(response.node_id(), from);
node.transactions.remove_value(&u16::from_ne_bytes(t));
if node.addr.is_ipv4() {
self.ipv4_table.handle_activity(*id, from.ip(), from.port());
} else {
self.ipv6_table.handle_activity(*id, from.ip(), from.port());
}
}
}
}
fn get_peer_node<T: Into<SocketAddr>>(&mut self, id: U160, addr: T) -> &mut PeerNode {
let addr = addr.into();
self.nodes.update_expiry(
&(id, addr.ip(), addr.port()),
Instant::now().add(Duration::from_secs(60 * 15)),
);
self.nodes
.entry((id, addr.ip(), addr.port()))
.or_insert_with(Instant::now().add(Duration::from_secs(60 * 15)), || {
PeerNode::create(id, addr)
})
}
fn create_token(&mut self, id: U160, addr: SocketAddr, _info_hash: U160) -> Vec<u8> {
let token: u64 = random();
self.get_peer_node(id, addr).tokens.insert(
token.clone(),
Instant::now().add(Duration::from_secs(60 * 5)),
);
token.to_ne_bytes().to_vec()
}
fn handle_request(&mut self, request: &RequestBody, t: [u8; 2], from: SocketAddr) {
if request.node_id() == self.id {
// In this house we don't talk to imposters
return;
}
let contact_info = ContactInfo {
id: request.node_id(),
contact: from,
};
if !self.has_node(&contact_info) {
if from.is_ipv4() {
self.ipv4_table.add_incoming_node(
contact_info.id,
contact_info.contact.ip(),
contact_info.contact.port(),
);
} else {
self.ipv6_table.add_incoming_node(
contact_info.id,
contact_info.contact.ip(),
contact_info.contact.port(),
);
}
}
match request {
RequestBody::GetPeers(get_peers) => {
let table =
get_peers
.want
.unwrap_or_else(|| if from.is_ipv4() { Want::N4 } else { Want::N6 });
let peers = self.peer_storage.read().unwrap().get_peers(
get_peers.info_hash,
if table == Want::N4 {
LookupFilter::IPv4
} else {
LookupFilter::IPv6
},
);
let token = self.create_token(get_peers.id, from, get_peers.info_hash);
if peers.len() == 0 {
self.queue.push_back((
Message::response_get_peers_not_found(
self.id,
t,
token,
if table == Want::N4 {
self.ipv4_table.find_k_closest_nodes(
get_peers.info_hash,
Some((get_peers.id, from.ip(), from.port())),
)
} else {
vec![]
},
if table == Want::N6 {
self.ipv6_table.find_k_closest_nodes(
get_peers.info_hash,
Some((get_peers.id, from.ip(), from.port())),
)
} else {
vec![]
},
),
from,
))
} else {
self.queue.push_back((
Message::response_get_peers_found(self.id, t, token, peers),
from,
));
}
}
RequestBody::FindNode(find_node) => {
let table =
find_node
.want
.unwrap_or_else(|| if from.is_ipv4() { Want::N4 } else { Want::N6 });
let nodes = if table == Want::N4 {
if let Some(node) = self.ipv4_table.find_node(find_node.target) {
vec![node]
} else {
self.ipv4_table.find_k_closest_nodes(
find_node.target,
Some((find_node.id, from.ip(), from.port())),
)
}
} else {
vec![]
};
let nodes6 = if table == Want::N6 {
if let Some(node) = self.ipv6_table.find_node(find_node.target) {
vec![node]
} else {
self.ipv6_table.find_k_closest_nodes(
find_node.target,
Some((find_node.id, from.ip(), from.port())),
)
}
} else {
vec![]
};
self.queue
.push_back((Message::response_find_node(self.id, t, nodes, nodes6), from))
}
RequestBody::Ping(id) => {
if from.is_ipv4() {
self.ipv4_table.handle_activity(*id, from.ip(), from.port())
} else {
self.ipv6_table.handle_activity(*id, from.ip(), from.port())
}
self.queue
.push_back((Message::response_empty(self.id, t), from))
}
RequestBody::AnnouncePeer(announce) => {
let ok = if let Some(node) =
self.nodes.get_mut(&(announce.id, from.ip(), from.port()))
{
if announce.token.len() != 8 {
false
} else {
node.tokens.remove(&u64::from_ne_bytes(
announce.token[0..8].try_into().unwrap(),
))
}
} else {
false
};
if !ok {
self.queue
.push_back((Message::error(t, 203, "Bad token".to_string()), from))
} else {
self.peer_storage.write().unwrap().add_peers(
announce.info_hash,
vec![SocketAddr::new(
from.ip(),
if announce.implied_port {
from.port()
} else {
announce.port
},
)],
);
self.queue
.push_back((Message::response_empty(self.id, t), from))
}
}
}
}
pub fn next(&mut self) -> Option<(Message, SocketAddr)> {
self.queue.pop_front()
}
fn clean(&mut self) {
for (_, node) in self.nodes.iter_mut() {
node.clean();
}
}
pub fn housekeeping(&mut self) {
self.clean();
self.update_tables();
if self
.last_bootstrap
.map(|last| last.add(BOOTSTRAP_INTERVAL) < Instant::now())
.unwrap_or(true)
{
for (addr, want) in self.bootstrap_nodes.clone() {
self.find_node_at(random(), addr, self.id, want);
}
self.last_bootstrap = Some(Instant::now());
}
}
fn update_tables(&mut self) {
self.ipv4_table.update_nodes();
self.ipv6_table.update_nodes();
for item in self.ipv6_table.get_silent_nodes() {
self.ping(&item);
}
for item in self.ipv4_table.get_silent_nodes() {
self.ping(&item);
}
}
pub fn num_tracking_nodes(&self) -> usize {
self.nodes.len()
}
pub fn num_ipv4_table_nodes(&self) -> usize {
self.ipv4_table.count_nodes()
}
pub fn num_ipv6_table_nodes(&self) -> usize {
self.ipv6_table.count_nodes()
}
}
#[cfg(test)]
mod test {
use crate::host_node::HostNode;
use crate::krpc::Message;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::rc::Rc;
use std::sync::RwLock;
use torment_core::infohash::v1::U160;
use torment_core::{ContactInfo, PeerStorage};
fn get_host() -> HostNode {
let peer_storage = Rc::new(RwLock::new(PeerStorage::new()));
HostNode::new(peer_storage, None)
}
#[test]
fn ping_pong() {
let mut host = get_host();
let t = [0u8; 2];
let ip = Ipv4Addr::LOCALHOST;
let id = U160::random();
let from = SocketAddr::new(IpAddr::from(ip), 4);
host.process(Message::request_ping(id, t), from);
assert_eq!(1, host.ipv4_table.count_nodes());
assert_eq!(
host.next(),
Some((Message::response_empty(host.id, t), from))
)
}
#[test]
fn get_peers() {
let mut host = get_host();
let t = [0u8; 2];
let ip = Ipv4Addr::LOCALHOST;
let from = SocketAddr::new(IpAddr::from(ip), 4);
let id = U160::random();
host.process(Message::request_ping(id, t), from);
assert_eq!(1, host.ipv4_table.count_nodes());
host.next();
let other = U160::random();
host.process(Message::request_ping(other, t), from);
assert_eq!(2, host.ipv4_table.count_nodes());
host.next();
let info_hash = U160::random();
host.process(Message::request_get_peers(id, t, info_hash, None), from);
let nodes = host
.next()
.unwrap()
.0
.expect_response()
.1
.expect_get_peers()
.nodes
.nodes;
assert_eq!(1, nodes.len());
assert_eq!(
&ContactInfo {
id: other,
contact: SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 4),
},
nodes.get(0).unwrap()
);
}
#[test]
fn announce_peer() {
let mut host = get_host();
let t = [0u8; 2];
let ip = Ipv4Addr::LOCALHOST;
let from = SocketAddr::new(IpAddr::from(ip), 4);
let id = U160::random();
host.process(Message::request_ping(id, t), from);
assert_eq!(1, host.ipv4_table.count_nodes());
host.next();
let other = U160::random();
host.process(Message::request_ping(other, t), from);
assert_eq!(2, host.ipv4_table.count_nodes());
host.next();
let info_hash = U160::random();
host.process(Message::request_get_peers(other, t, info_hash, None), from);
let resp = host
.next()
.unwrap()
.0
.expect_response()
.1
.expect_get_peers();
host.process(
Message::request_announce_peer(other, t, info_hash, resp.token.clone(), 4, false),
from,
);
host.next().unwrap().0.expect_response().1.expect_empty();
host.process(
Message::request_announce_peer(other, t, info_hash, resp.token, 4, false),
from,
);
host.next().unwrap().0.expect_error();
}
}

889
torment-dht/src/krpc.rs
Unescape Escape View File

@ -0,0 +1,889 @@
use bendy::decoding::{Error as DecodingError, Object};
use bendy::encoding::{Error as EncodingError, SingleItemEncoder};
use std::convert::TryInto;
use std::fmt::{Display, Formatter};
use std::option::Option::Some;
use std::str::FromStr;
use torment_core::infohash::v1::U160;
use torment_core::infohash::InfoHashCapable;
use torment_core::{CompactContact, ContactInfo, ParsingError};
pub use bendy::decoding::FromBencode;
pub use bendy::encoding::ToBencode;
use std::net::SocketAddr;
const VERSION: &'static [u8] = &[b'e', b't', 0, 0];
fn ver() -> Option<String> {
// None
Some(String::from_utf8_lossy(VERSION).to_string())
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Message {
pub transaction_id: [u8; 2],
pub body: MessageBody,
pub version: Option<String>,
}
impl Message {
pub fn expect_response(self) -> ([u8; 2], ResponseBody, Option<String>) {
(
self.transaction_id,
match self.body {
MessageBody::Response(resp) => resp,
b => panic!("Expected response got {}", b.name()),
},
self.version,
)
}
pub fn expect_request(self) -> ([u8; 2], RequestBody, Option<String>) {
(
self.transaction_id,
match self.body {
MessageBody::Request(req) => req,
b => panic!("Expected request got {}", b.name()),
},
self.version,
)
}
pub fn expect_error(self) -> ([u8; 2], ErrorBody, Option<String>) {
(
self.transaction_id,
match self.body {
MessageBody::Error(err) => err,
b => panic!("Expected error got {}", b.name()),
},
self.version,
)
}
pub fn request_ping(id: U160, t: [u8; 2]) -> Message {
Message {
transaction_id: t,
body: MessageBody::Request(RequestBody::Ping(id)),
version: ver(),
}
}
pub fn request_get_peers(id: U160, t: [u8; 2], info_hash: U160, want: Option<Want>) -> Message {
Message {
transaction_id: t,
body: MessageBody::Request(RequestBody::GetPeers(GetPeersRequest {
id,
want,
info_hash,
})),
version: ver(),
}
}
pub fn request_find_node(id: U160, t: [u8; 2], target: U160, want: Option<Want>) -> Message {
Message {
transaction_id: t,
body: MessageBody::Request(RequestBody::FindNode(FindNodeRequest { id, want, target })),
version: ver(),
}
}
pub fn request_announce_peer(
id: U160,
t: [u8; 2],
info_hash: U160,
token: Vec<u8>,
port: u16,
implied_port: bool,
) -> Message {
Message {
transaction_id: t,
body: MessageBody::Request(RequestBody::AnnouncePeer(AnnouncePeerRequest {
id,
implied_port,
info_hash,
port,
token,
})),
version: ver(),
}
}
pub fn response_get_peers_found(
id: U160,
t: [u8; 2],
token: Vec<u8>,
peers: Vec<SocketAddr>,
) -> Message {
Message {
transaction_id: t,
body: MessageBody::Response(ResponseBody::GetPeers(GetPeersResponse {
id,
token,
values: peers,
nodes: ContactNodes::default(),
})),
version: ver(),
}
}
pub fn response_get_peers_not_found(
id: U160,
t: [u8; 2],
token: Vec<u8>,
nodes: Vec<ContactInfo>,
nodes6: Vec<ContactInfo>,
) -> Message {
Message {
transaction_id: t,
body: MessageBody::Response(ResponseBody::GetPeers(GetPeersResponse {
id,
token,
values: vec![],
nodes: ContactNodes { nodes, nodes6 },
})),
version: ver(),
}
}
pub fn response_find_node(
id: U160,
t: [u8; 2],
nodes: Vec<ContactInfo>,
nodes6: Vec<ContactInfo>,
) -> Message {
Message {
transaction_id: t,
body: MessageBody::Response(ResponseBody::FindNode(FindNodeResponse {
id,
nodes,
nodes6,
})),
version: ver(),
}
}
pub fn response_empty(id: U160, t: [u8; 2]) -> Message {
Message {
transaction_id: t,
body: MessageBody::Response(ResponseBody::Empty(id)),
version: ver(),
}
}
pub fn error(t: [u8; 2], code: u16, text: String) -> Message {
Message {
transaction_id: t,
body: MessageBody::Error(ErrorBody { code, text }),
version: ver(),
}
}
}
impl ToBencode for Message {
const MAX_DEPTH: usize = 3;
fn encode(&self, encoder: SingleItemEncoder) -> Result<(), EncodingError> {
encoder.emit_unsorted_dict(|d| {
d.emit_pair_with(b"t", |e| e.emit_bytes(&self.transaction_id[..]))?;
match self.body {
MessageBody::Error(ref err) => {
d.emit_pair(b"y", "e")?;
d.emit_pair_with(b"e", |d| {
d.emit_list(|l| {
l.emit_int(err.code)?;
l.emit_str(&err.text)
})
})?;
}
MessageBody::Request(ref req) => {
d.emit_pair(b"y", "q")?;
match req {
RequestBody::Ping(ref p) => {
d.emit_pair(b"q", "ping")?;
d.emit_pair_with(b"a", |d| d.emit_dict(|mut d| d.emit_pair(b"id", p)))?;
}
RequestBody::AnnouncePeer(ref ann) => {
d.emit_pair(b"q", "announce_peer")?;
d.emit_pair_with(b"a", |d| {
d.emit_unsorted_dict(|d| {
d.emit_pair(b"id", ann.id)?;
d.emit_pair(
b"implied_port",
if ann.implied_port { 1 } else { 0 },
)?;
d.emit_pair(b"info_hash", ann.info_hash)?;
d.emit_pair(b"port", ann.port)?;
d.emit_pair_with(b"token", |v| v.emit_bytes(&ann.token))
})
})?;
}
RequestBody::FindNode(ref fin) => {
d.emit_pair(b"q", "find_node")?;
d.emit_pair_with(b"a", |d| {
d.emit_unsorted_dict(|d| {
d.emit_pair(b"id", fin.id)?;
if let Some(ref want) = fin.want {
d.emit_pair(
b"want",
match want {
Want::N4 => "n4",
Want::N6 => "n6",
},
)?;
}
d.emit_pair(b"target", fin.target)
})
})?;
}
RequestBody::GetPeers(ref get) => {
d.emit_pair(b"q", "get_peers")?;
d.emit_pair_with(b"a", |d| {
d.emit_unsorted_dict(|d| {
d.emit_pair(b"id", get.id)?;
if let Some(ref want) = get.want {
d.emit_pair(
b"want",
match want {
Want::N4 => "n4",
Want::N6 => "n6",
},
)?;
}
d.emit_pair(b"info_hash", get.info_hash)
})
})?;
}
};
}
MessageBody::Response(ref resp) => {
d.emit_pair(b"y", "r")?;
match resp {
ResponseBody::GetPeers(ref get) => {
d.emit_pair_with(b"r", |d| {
d.emit_unsorted_dict(|d| {
d.emit_pair(b"id", get.id)?;
d.emit_pair_with(b"token", |v| v.emit_bytes(&get.token))?;
if get.values.len() > 0 {
d.emit_pair_with(b"values", |e| {
e.emit_list(|l| {
for item in &get.values {
l.emit_bytes(&item.to_compact_contact())?;
}
Ok(())
})
})?;
}
if get.nodes.nodes.len() > 0 {
d.emit_pair_with(b"nodes", |v| {
v.emit_bytes(&compact_nodes(&get.nodes.nodes))
})?;
}
if get.nodes.nodes6.len() > 0 {
d.emit_pair_with(b"nodes6", |v| {
v.emit_bytes(&compact_nodes(&get.nodes.nodes6))
})?;
}
Ok(())
})
})?;
}
ResponseBody::FindNode(ref find) => {
d.emit_pair_with(b"r", |d| {
d.emit_unsorted_dict(|d| {
d.emit_pair(b"id", find.id)?;
if find.nodes.len() > 0 {
d.emit_pair_with(b"nodes", |v| {
v.emit_bytes(&compact_nodes(&find.nodes))
})?;
}
if find.nodes6.len() > 0 {
d.emit_pair_with(b"nodes6", |v| {
v.emit_bytes(&compact_nodes(&find.nodes6))
})?;
}
Ok(())
})
})?;
}
ResponseBody::Empty(ref id) => {
d.emit_pair_with(b"r", |d| {
d.emit_dict(|mut d| d.emit_pair(b"id", id))
})?;
}
}
}
}
if let Some(ref version) = self.version {
d.emit_pair(b"v", version)?;
}
Ok(())
})
}
}
#[derive(Debug)]
enum KRPCError {
Generic(String),
}
impl std::error::Error for KRPCError {}
impl Display for KRPCError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
KRPCError::Generic(ref error) => write!(f, "Generic error: {}", error),
}
}
}
impl FromBencode for Message {
fn decode_bencode_object(object: Object) -> Result<Self, DecodingError>
where
Self: Sized,
{
let mut dict = object.try_into_dictionary()?;
let mut t = None;
let mut y = None;
let mut v = None;
let mut q = None;
let mut info_hash = None;
let mut target = None;
let mut id = None;
let mut token = None;
let mut implied_port = false;
let mut want = None;
let mut port = 0;
let mut nodes = vec![];
let mut nodes6 = vec![];
let mut values = vec![];
let mut code = 0;
let mut text = String::new();
while let Some((name, obj)) = dict.next_pair()? {
match name {
b"y" => y = obj.try_into_bytes()?.first(),
b"t" => t = Some(obj.try_into_bytes()?.to_vec()),
b"v" => v = Some(obj.try_into_bytes()?.to_vec()),
b"q" => q = Some(obj.try_into_bytes()?.to_vec()),
b"a" => {
let mut args = obj.try_into_dictionary()?;
while let Some((arg_name, arg_obj)) = args.next_pair()? {
match arg_name {
b"info_hash" => {
info_hash = Some(U160::from_bytes(arg_obj.try_into_bytes()?)?)
}
b"id" => id = Some(U160::from_bytes(arg_obj.try_into_bytes()?)?),
b"target" => {
target = Some(U160::from_bytes(arg_obj.try_into_bytes()?)?)
}
b"token" => token = Some(arg_obj.try_into_bytes()?.to_vec()),
b"implied_port" => {
implied_port = usize::from_str(arg_obj.try_into_integer()?)? > 0
}
b"port" => port = u16::from_str(arg_obj.try_into_integer()?)?,
b"want" => {
want = match arg_obj.try_into_bytes() {
Ok(b"n4") => Some(Want::N4),
Ok(b"n6") => Some(Want::N6),
_ => None,
}
}
_ => {}
}
}
}
b"r" => {
let mut ret = obj.try_into_dictionary()?;
while let Some((ret_name, ret_obj)) = ret.next_pair()? {
match ret_name {
b"id" => id = Some(U160::from_bytes(ret_obj.try_into_bytes()?)?),
b"token" => token = Some(ret_obj.try_into_bytes()?.to_vec()),
b"values" => {
let mut list = ret_obj.try_into_list()?;
while let Some(value) = list.next_object()? {
values.push(SocketAddr::from_compact_contact(
value.try_into_bytes()?,
)?);
}
}
b"nodes" => {
nodes = ret_obj
.try_into_bytes()?
.chunks(26)
.filter_map(|x| ContactInfo::from_bytes(x).ok())
.collect()
}
b"nodes6" => {
nodes6 = ret_obj
.try_into_bytes()?
.chunks(38)
.filter_map(|x| ContactInfo::from_bytes(x).ok())
.collect()
}
_ => {}
}
}
}
b"e" => {
let mut err = obj.try_into_list()?;
code =
u16::from_str(err.next_object()?.ok_or(ParsingError)?.try_into_integer()?)?;
text = String::from_utf8_lossy(
err.next_object()?.ok_or(ParsingError)?.try_into_bytes()?,
)
.to_string();
}
_ => {}
}
}
let t = t.ok_or(ParsingError)?;
let t = if t.len() == 0 {
[0, 0]
} else if t.len() == 1 {
[0, t[0]]
} else {
t[..2].try_into()?
};
Ok(Message {
transaction_id: t,
body: match y {
Some(b'q') => {
let q = if let Some(q) = q {
q
} else {
return Err(KRPCError::Generic("DHT Query without type".to_string()).into());
};
let id = id.ok_or(ParsingError)?;
MessageBody::Request(match &q[..] {
b"ping" => RequestBody::Ping(id),
b"find_node" => RequestBody::FindNode(FindNodeRequest {
id,
want,
target: target.ok_or(ParsingError)?,
}),
b"get_peers" => RequestBody::GetPeers(GetPeersRequest {
id,
want,
info_hash: info_hash.ok_or(ParsingError)?,
}),
b"announce_peer" => RequestBody::AnnouncePeer(AnnouncePeerRequest {
id,
implied_port,
info_hash: info_hash.ok_or(ParsingError)?,
port,
token: token.ok_or(ParsingError)?,
}),
_ => Err(ParsingError)?,
})
}
Some(b'r') => {
let id = id.ok_or(ParsingError)?;
MessageBody::Response(if let Some(token) = token {
ResponseBody::GetPeers(GetPeersResponse {
id,
token,
values,
nodes: ContactNodes { nodes, nodes6 },
})
} else if nodes6.len() > 0 || nodes.len() > 0 {
ResponseBody::FindNode(FindNodeResponse { id, nodes, nodes6 })
} else {
ResponseBody::Empty(id)
})
}
Some(b'e') => MessageBody::Error(ErrorBody { code, text }),
Some(x) => return Err(KRPCError::Generic(format!("Type {} is unknown", x)).into()),
None => {
return Err(KRPCError::Generic("No type given for message".to_string()).into());
}
},
version: v.map(|x| String::from_utf8_lossy(&x).to_string()),
})
}
}
fn compact_nodes(input: &Vec<ContactInfo>) -> Vec<u8> {
let mut output = vec![];
for item in input {
output.extend(item.to_bytes())
}
output
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum MessageBody {
Request(RequestBody),
Response(ResponseBody),
Error(ErrorBody),
}
impl MessageBody {
fn name(&self) -> &str {
match self {
MessageBody::Request(_) => "request",
MessageBody::Response(_) => "response",
MessageBody::Error(_) => "error",
}
}
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum RequestBody {
Ping(U160),
FindNode(FindNodeRequest),
GetPeers(GetPeersRequest),
AnnouncePeer(AnnouncePeerRequest),
}
impl RequestBody {
pub fn node_id(&self) -> U160 {
match self {
RequestBody::Ping(id) => *id,
RequestBody::FindNode(find_node) => find_node.id,
RequestBody::GetPeers(get_peers) => get_peers.id,
RequestBody::AnnouncePeer(announce_peer) => announce_peer.id,
}
}
pub fn expect_ping(self) -> U160 {
match self {
RequestBody::Ping(id) => id,
_ => panic!("Expected ping request"),
}
}
pub fn expect_find_node(self) -> FindNodeRequest {
match self {
RequestBody::FindNode(find_node) => find_node,
_ => panic!("Expected find node request"),
}
}
pub fn expect_get_peers(self) -> GetPeersRequest {
match self {
RequestBody::GetPeers(get_peers) => get_peers,
_ => panic!("Expected get peers request"),
}
}
pub fn expect_announce_peer(self) -> AnnouncePeerRequest {
match self {
RequestBody::AnnouncePeer(announce_peer) => announce_peer,
_ => panic!("Expected announce peer request"),
}
}
}
#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq, Copy, Hash)]
pub enum Want {
N6,
N4,
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct FindNodeRequest {
pub id: U160,
pub want: Option<Want>,
pub target: U160,
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct GetPeersRequest {
pub id: U160,
pub want: Option<Want>,
pub info_hash: U160,
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct AnnouncePeerRequest {
pub id: U160,
pub implied_port: bool,
pub info_hash: U160,
pub port: u16,
pub token: Vec<u8>,
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum ResponseBody {
Empty(U160),
FindNode(FindNodeResponse),
GetPeers(GetPeersResponse),
}
impl ResponseBody {
pub fn node_id(&self) -> U160 {
match self {
ResponseBody::Empty(id) => *id,
ResponseBody::FindNode(find_node) => find_node.id,
ResponseBody::GetPeers(get_peers) => get_peers.id,
}
}
pub fn expect_empty(self) -> U160 {
match self {
ResponseBody::Empty(id) => id,
_ => panic!("Expected empty response"),
}
}
pub fn expect_find_node(self) -> FindNodeResponse {
match self {
ResponseBody::FindNode(find_node) => find_node,
_ => panic!("Expected find node response"),
}
}
pub fn expect_get_peers(self) -> GetPeersResponse {
match self {
ResponseBody::GetPeers(get_peers) => get_peers,
_ => panic!("Expected get peers response"),
}
}
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct FindNodeResponse {
pub id: U160,
pub nodes: Vec<ContactInfo>,
pub nodes6: Vec<ContactInfo>,
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct GetPeersResponse {
pub id: U160,
pub token: Vec<u8>,
pub values: Vec<SocketAddr>,
pub nodes: ContactNodes,
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct ContactNodes {
pub nodes: Vec<ContactInfo>,
pub nodes6: Vec<ContactInfo>,
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct ErrorBody {
pub code: u16,
pub text: String,
}
#[cfg(test)]
mod test {
use crate::krpc::Message;
use bendy::decoding::FromBencode;
use bendy::encoding::ToBencode;
use std::borrow::Borrow;
#[test]
fn error_response() {
let be: &[u8] = b"d1:eli201e23:A Generic Error Ocurrede1:t2:aa1:y1:ee";
let msg = Message::from_bencode(be).expect("de-bencoding");
let bencoded = msg.to_bencode().expect("Failed bencoding");
println!("{:?}", msg);
println!("{:?}", bencoded);
assert_eq!(
be,
bencoded.borrow() as &[u8],
"left [{}] != right [{}]",
String::from_utf8_lossy(be),
String::from_utf8_lossy(&bencoded)
)
}
#[test]
fn ping_request() {
let be: &[u8] = b"d1:ad2:id20:abcdefghij0123456789e1:q4:ping1:t2:aa1:y1:qe";
let msg = Message::from_bencode(be).expect("de-bencoding");
let bencoded = msg.to_bencode().expect("Failed bencoding");
println!("{:?}", msg);
println!("{:?}", bencoded);
assert_eq!(
be,
bencoded.borrow() as &[u8],
"left [{}] != right [{}]",
String::from_utf8_lossy(be),
String::from_utf8_lossy(&bencoded)
)
}
#[test]
fn ping_response() {
let be: &[u8] = b"d1:rd2:id20:mnopqrstuvwxyz123456e1:t2:aa1:y1:re";
let msg = Message::from_bencode(be).expect("de-bencoding");
let bencoded = msg.to_bencode().expect("Failed bencoding");
println!("{:?}", msg);
println!("{:?}", bencoded);
assert_eq!(
be,
bencoded.borrow() as &[u8],
"left [{}] != right [{}]",
String::from_utf8_lossy(be),
String::from_utf8_lossy(&bencoded)
)
}
#[test]
fn find_node_request() {
let be: &[u8] = b"d1:ad2:id20:abcdefghij01234567896:target20:mnopqrstuvwxyz123456e1:q9:find_node1:t2:aa1:y1:qe";
let msg = Message::from_bencode(be).expect("de-bencoding");
let bencoded = msg.to_bencode().expect("Failed bencoding");
println!("{:?}", msg);
println!("{:?}", bencoded);
assert_eq!(
be,
bencoded.borrow() as &[u8],
"left [{}] != right [{}]",
String::from_utf8_lossy(be),
String::from_utf8_lossy(&bencoded)
)
}
#[test]
fn find_node_response() {
let be: &[u8] =
b"d1:rd2:id20:0123456789abcdefghij5:nodes26:abcdefghij0123456789axje.ue1:t2:aa1:y1:re";
let msg = Message::from_bencode(be).expect("de-bencoding");
let bencoded = msg.to_bencode().expect("Failed bencoding");
println!("{:?}", msg);
println!("{:?}", bencoded);
assert_eq!(
be,
bencoded.borrow() as &[u8],
"left [{}] != right [{}]",
String::from_utf8_lossy(be),
String::from_utf8_lossy(&bencoded)
)
}
#[test]
fn get_peers_request() {
let be: &[u8] =
b"d1:ad2:id20:abcdefghij01234567899:info_hash20:mnopqrstuvwxyz123456e1:q9:get_peers1:t2:aa1:y1:qe";
let msg = Message::from_bencode(be).expect("de-bencoding");
let bencoded = msg.to_bencode().expect("Failed bencoding");
println!("{:?}", msg);
println!("{:?}", bencoded);
assert_eq!(
be,
bencoded.borrow() as &[u8],
"left [{}] != right [{}]",
String::from_utf8_lossy(be),
String::from_utf8_lossy(&bencoded)
)
}
#[test]
fn get_peers_response_a() {
let be: &[u8] =
b"d1:rd2:id20:abcdefghij01234567895:token8:aoeusnth6:valuesl6:axje.u6:idhtnmee1:t2:aa1:y1:re";
let msg = Message::from_bencode(be).expect("de-bencoding");
let bencoded = msg.to_bencode().expect("Failed bencoding");
println!("{:?}", msg);
println!("{:?}", bencoded);
assert_eq!(
be,
bencoded.borrow() as &[u8],
"left [{}] != right [{}]",
String::from_utf8_lossy(be),
String::from_utf8_lossy(&bencoded)
)
}
#[test]
fn get_peers_response_b() {
let be: &[u8] =
b"d1:rd2:id20:abcdefghij01234567895:nodes26:abcdefghij0123456789axje.u5:token8:aoeusnthe1:t2:aa1:y1:re";
let msg = Message::from_bencode(be).expect("de-bencoding");
let bencoded = msg.to_bencode().expect("Failed bencoding");
println!("{:?}", msg);
println!("{:?}", bencoded);
assert_eq!(
be,
bencoded.borrow() as &[u8],
"left [{}] != right [{}]",
String::from_utf8_lossy(be),
String::from_utf8_lossy(&bencoded)
)
}
#[test]
fn announce_peer_request() {
let be: &[u8] =
b"d1:ad2:id20:abcdefghij012345678912:implied_porti1e9:info_hash20:mnopqrstuvwxyz1234564:porti6881e5:token8:aoeusnthe1:q13:announce_peer1:t2:aa1:y1:qe";
let msg = Message::from_bencode(be).expect("de-bencoding");
let bencoded = msg.to_bencode().expect("Failed bencoding");
println!("{:?}", msg);
println!("{:?}", bencoded);
assert_eq!(
be,
bencoded.borrow() as &[u8],
"left [{}] != right [{}]",
String::from_utf8_lossy(be),
String::from_utf8_lossy(&bencoded)
)
}
#[test]
fn announce_peer_response() {
let be: &[u8] = b"d1:rd2:id20:mnopqrstuvwxyz123456e1:t2:aa1:y1:re";
let msg = Message::from_bencode(be).expect("de-bencoding");
let bencoded = msg.to_bencode().expect("Failed bencoding");
println!("{:?}", msg);
println!("{:?}", bencoded);
assert_eq!(
be,
bencoded.borrow() as &[u8],
"left [{}] != right [{}]",
String::from_utf8_lossy(be),
String::from_utf8_lossy(&bencoded)
)
}
}

647
torment-dht/src/lib.rs
Unescape Escape View File

@ -0,0 +1,647 @@
#![allow(dead_code)]
use std::cmp::min;
use std::collections::{BTreeMap, VecDeque};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::time::{Duration, Instant};
use torment_core::infohash::v1::U160;
use torment_core::ContactInfo;
use std::mem;
use std::ops::Add;
use std::ops::Bound::Included;
use std::option::Option::Some;
pub mod host_node;
pub mod krpc;
const DHT_K: usize = 8;
#[derive(Debug)]
pub struct Table {
id: U160,
buckets: BTreeMap<U160, Bucket>,
}
type NodeHandle = (U160, IpAddr, u16);
impl Default for Table {
fn default() -> Self {
Table::new()
}
}
impl Table {
pub fn new_with_id(id: U160) -> Table {
Table {
id,
buckets: {
let mut tree = BTreeMap::new();
tree.insert(U160::MAX, Bucket::new(U160::MIN, U160::MAX));
tree
},
}
}
pub fn new() -> Table {
Self::new_with_id(U160::random())
}
fn get_bucket_index(&self, id: U160) -> U160 {
*self.buckets.range(id..).next().unwrap().0
}
fn get_bucket(&self, id: U160) -> &Bucket {
self.buckets
.get(&self.get_bucket_index(id))
.expect("DHT corrupt")
}
fn get_bucket_mut(&mut self, id: U160) -> &mut Bucket {
self.buckets
.get_mut(&self.get_bucket_index(id))
.expect("DHT corrupt")
}
pub fn has_node(&self, contact_info: &ContactInfo) -> bool {
self.get_node(
contact_info.id,
contact_info.contact.ip(),
contact_info.contact.port(),
)
.is_some()
}
fn get_node(&self, id: U160, address: IpAddr, port: u16) -> Option<&Node> {
self.get_bucket(id).nodes.get(&(id, address, port))
}
pub fn find_node(&self, id: U160) -> Option<ContactInfo> {
self.get_bucket(id)
.nodes
.range((
Included(&(id, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)),
Included(&(
id,
IpAddr::V6(Ipv6Addr::new(
u16::MAX,
u16::MAX,
u16::MAX,
u16::MAX,
u16::MAX,
u16::MAX,
u16::MAX,
u16::MAX,
)),
u16::MAX,
)),
))
.next()
.map(|(_, node)| node.to_contact_info())
}
fn get_node_mut(&mut self, id: U160, address: IpAddr, port: u16) -> Option<&mut Node> {
self.get_bucket_mut(id).nodes.get_mut(&(id, address, port))
}
pub fn find_k_closest_nodes(
&self,
target: U160,
excluding: Option<(U160, IpAddr, u16)>,
) -> Vec<ContactInfo> {
let mut nodes = vec![];
let mut found_before = 0;
let mut found_after = 0;
let mut before_range = self.buckets.range(..target);
while let Some((_, bucket)) = before_range.next_back() {
let info = bucket.get_contact_info();
found_before += info.len();
nodes.extend(info);
if found_before >= DHT_K {
break;
}
}
let mut after_range = self.buckets.range(target..);
while let Some((_, bucket)) = after_range.next() {
let info = bucket.get_contact_info();
found_after += info.len();
nodes.extend(info);
if found_after >= DHT_K {
break;
}
}
nodes.sort_by_key(|x| x.id ^ target);
let nodes = if let Some(exclude) = excluding {
nodes
.iter()
.filter(|x| {
x.id != exclude.0
|| x.contact.ip() != exclude.1
|| x.contact.port() != exclude.2
})
.copied()
.collect()
} else {
nodes
};
nodes[..min(8, nodes.len())].to_vec()
}
pub fn handle_activity(&mut self, id: U160, address: IpAddr, port: u16) {
let bucket_has_activity = if let Some(node) = self.get_node_mut(id, address, port) {
node.handle_activity();
true
} else {
self.add_node(id, address, port)
};
if bucket_has_activity {
self.get_bucket_mut(id).last_activity = Some(Instant::now());
}
}
pub fn handle_outgoing_activity(&mut self, id: U160, address: IpAddr, port: u16) {
if let Some(node) = self.get_node_mut(id, address, port) {
node.handle_outgoing_activity();
}
}
pub fn add_incoming_node(&mut self, id: U160, address: IpAddr, port: u16) {
self.add_node(id, address, port);
self.handle_activity(id, address, port);
}
pub fn add_node(&mut self, id: U160, address: IpAddr, port: u16) -> bool {
let own_id = self.id;
if own_id == id {
return false;
}
let bucket = self.get_bucket_mut(id);
if bucket.add_node(id, address, port) {
return true;
}
if !(bucket.start <= own_id
&& own_id <= bucket.end
&& (bucket.end - bucket.start) > U160::ONE)
{
return false;
}
let low_end = bucket.start + (bucket.end - bucket.start).half();
fn split_node_map(
nodes: &mut BTreeMap<NodeHandle, Node>,
split_at: U160,
) -> (BTreeMap<NodeHandle, Node>, BTreeMap<NodeHandle, Node>) {
let mut high = BTreeMap::new();
let mut low = BTreeMap::new();
let keys = nodes.keys().copied().collect::<Vec<_>>();
for key in keys {
let node = nodes.remove(&key).unwrap();
if node.id <= split_at {
low.insert(key, node);
} else {
high.insert(key, node);
}
}
(low, high)
}
fn recount_nodes_per_id(nodes: &BTreeMap<NodeHandle, Node>) -> BTreeMap<U160, usize> {
nodes
.values()
.map(|x| x.id)
.fold(BTreeMap::new(), |mut c, i| {
c.entry(i).and_modify(|x| *x += 1).or_insert(1);
c
})
}
fn get_last_activity(nodes: &BTreeMap<NodeHandle, Node>) -> Option<Instant> {
nodes.values().filter_map(|x| x.last_activity).max()
}
let (low, high) = split_node_map(&mut bucket.nodes, low_end);
let (low_q, high_q) = split_node_map(&mut bucket.queue, low_end);
bucket.last_activity = get_last_activity(&high);
bucket.nodes_per_id = recount_nodes_per_id(&high);
bucket.nodes = high;
bucket.queue = high_q;
let low_start = bucket.start;
bucket.start = low_end + U160::ONE;
let mut new_bucket = Bucket::new(low_start, low_end);
new_bucket.last_activity = get_last_activity(&low);
new_bucket.nodes_per_id = recount_nodes_per_id(&low);
new_bucket.nodes = low;
new_bucket.queue = low_q;
self.buckets.insert(low_end, new_bucket);
self.add_node(id, address, port)
}
pub fn get_silent_nodes(&self) -> Vec<ContactInfo> {
self.iter()
.filter(|node| {
if let Some(last_activity) = node.last_activity {
if last_activity.add(Duration::from_secs(60 * 15)) < Instant::now() {
return true;
}
}
if let Some(last_outgoing_activity) = node.last_activity {
if node.last_activity.is_none()
&& last_outgoing_activity.add(Duration::from_secs(60)) < Instant::now()
{
return true;
}
}
false
})
.map(|node| node.to_contact_info())
.collect()
}
pub fn update_nodes(&mut self) {
for node in self.iter_mut() {
if node
.last_activity
.map(|x| x.add(Duration::from_secs(15 * 60)) < Instant::now())
.unwrap_or(false)
{
if let NodeState::Questioning(n) = node.state {
if n > 3 {
#[cfg(debug_assertions)]
{
if node.state != NodeState::Bad {
println!(
"Node[id={},addr={}] went from {:?} to {:?}",
node.id,
node.sock_addr(),
node.state,
NodeState::Bad
);
}
}
node.state = NodeState::Bad;
}
if node
.last_outgoing_activity
.map(|x| x.add(Duration::from_secs(3 * 60)) < Instant::now())
.unwrap_or(false)
{
#[cfg(debug_assertions)]
{
if node.state != NodeState::Bad {
println!(
"Node[id={},addr={}] went from {:?} to {:?}",
node.id,
node.sock_addr(),
node.state,
NodeState::Bad
);
}
}
node.state = NodeState::Bad;
}
} else {
#[cfg(debug_assertions)]
{
if node.state != NodeState::Questioning(0) {
println!(
"Node[id={},addr={}] went from {:?} to {:?}",
node.id,
node.sock_addr(),
node.state,
NodeState::Questioning(0)
);
}
}
node.state = NodeState::Questioning(0);
}
}
}
for (_, bucket) in &mut self.buckets {
if !bucket.queue.is_empty() {
let old_nodes = mem::replace(&mut bucket.nodes, BTreeMap::new());
let mut keys = bucket
.queue
.iter()
.map(|(key, node)| (key.clone(), node.first_activity))
.collect::<Vec<_>>();
keys.sort_by_key(|(_, first_activity)| *first_activity);
let mut keys = VecDeque::from(keys);
for (key, old_node) in old_nodes {
if bucket.queue.is_empty() {
bucket.nodes.insert(key, old_node);
continue;
}
if old_node.state == NodeState::Bad {
if let Some((key, _)) = keys.pop_front() {
if let Some(node) = bucket.queue.remove(&key) {
bucket.nodes.insert(key, node);
continue;
}
}
}
bucket.nodes.insert(key, old_node);
}
}
}
}
pub fn count_nodes(&self) -> usize {
self.buckets.values().map(|x| x.nodes.len()).sum()
}
fn iter(&self) -> impl Iterator<Item = &Node> {
self.buckets
.iter()
.flat_map(|(_, bucket)| bucket.nodes.iter().map(|(_, node)| node))
}
fn iter_mut(&mut self) -> impl Iterator<Item = &mut Node> {
self.buckets
.iter_mut()
.flat_map(|(_, bucket)| bucket.nodes.iter_mut().map(|(_, node)| node))
}
}
#[derive(Debug)]
struct Bucket {
start: U160,
end: U160,
last_activity: Option<Instant>,
nodes: BTreeMap<(U160, IpAddr, u16), Node>,
nodes_per_id: BTreeMap<U160, usize>,
queue: BTreeMap<(U160, IpAddr, u16), Node>,
}
impl Bucket {
fn new(start: U160, end: U160) -> Bucket {
Bucket {
start,
end,
last_activity: None,
nodes: BTreeMap::new(),
nodes_per_id: Default::default(),
queue: BTreeMap::new(),
}
}
fn get_contact_info(&self) -> Vec<ContactInfo> {
self.nodes
.iter()
.map(|(_, node)| ContactInfo {
id: node.id,
contact: SocketAddr::new(node.address, node.port),
})
.collect()
}
fn internal_add_node(&mut self, id: U160, address: IpAddr, port: u16) {
self.nodes_per_id
.entry(id)
.and_modify(|x| *x += 1)
.or_insert(1);
self.nodes
.insert((id, address, port), Node::new(id, address, port));
self.last_activity = None;
}
fn add_node(&mut self, id: U160, address: IpAddr, port: u16) -> bool {
// Discard after 4 nodes with the same id
if self.nodes_per_id.get(&id).unwrap_or(&0) >= &4 {
return true;
}
if self.nodes.len() < DHT_K {
self.internal_add_node(id, address, port);
true
} else {
let mut replace = None;
let mut has_questionable = false;
for (key, node) in &self.nodes {
if node.state == NodeState::Bad {
replace = Some(*key);
break;
}
if let NodeState::Questioning(_) = node.state {
has_questionable = true;
}
}
if let Some(key) = replace {
self.nodes.remove(&key);
self.nodes_per_id.entry(key.0).and_modify(|x| *x -= 1);
if Some(&0usize) == self.nodes_per_id.get(&key.0) {
self.nodes_per_id.remove(&key.0);
}
self.internal_add_node(id, address, port);
true
} else {
if has_questionable {
let len = self.queue.len();
let entry = self
.queue
.entry((id, address, port))
.and_modify(|node| node.handle_activity());
if len < DHT_K {
entry.or_insert_with(|| {
let mut queued_node = Node::new(id, address, port);
queued_node.set_queued();
queued_node
});
}
}
false
}
}
}
}
#[derive(Debug)]
struct Node {
id: U160,
address: IpAddr,
port: u16,
first_activity: Instant,
last_outgoing_activity: Option<Instant>,
last_activity: Option<Instant>,
is_queued: bool,
state: NodeState,
}
#[derive(Debug, Ord, PartialOrd, Eq, PartialEq, Copy, Clone)]
enum NodeState {
Good,
Bad,
Questioning(u16),
}
impl Node {
fn new(id: U160, address: IpAddr, port: u16) -> Self {
Node {
id,
address,
port,
first_activity: Instant::now(),
last_outgoing_activity: None,
last_activity: None,
is_queued: false,
state: NodeState::Questioning(0),
}
}
fn sock_addr(&self) -> SocketAddr {
SocketAddr::new(self.address, self.port)
}
fn handle_activity(&mut self) {
self.last_activity = Some(Instant::now());
#[cfg(debug_assertions)]
{
if self.state != NodeState::Good {
println!(
"Node[id={},addr={}] went from {:?} to {:?}",
self.id,
self.sock_addr(),
self.state,
NodeState::Good
);
}
}
self.state = NodeState::Good;
}
fn handle_outgoing_activity(&mut self) {
self.last_outgoing_activity = Some(Instant::now());
if let NodeState::Questioning(nr) = self.state {
#[cfg(debug_assertions)]
{
if self.state != NodeState::Questioning(nr + 1) {
println!(
"Node[id={},addr={}] went from {:?} to {:?}",
self.id,
self.sock_addr(),
self.state,
NodeState::Questioning(nr + 1)
);
}
}
self.state = NodeState::Questioning(nr + 1);
}
}
fn set_queued(&mut self) {
self.is_queued = true;
}
fn set_live(&mut self) {
self.is_queued = false;
}
fn is_queued(&self) -> bool {
self.is_queued
}
fn to_contact_info(&self) -> ContactInfo {
ContactInfo {
id: self.id,
contact: SocketAddr::new(self.address, self.port),
}
}
}
#[cfg(test)]
mod test {
use crate::Table;
use std::net::{IpAddr, Ipv4Addr};
use torment_core::infohash::v1::U160;
#[test]
fn test_table_fulfillment() {
let mut table = Table::new_with_id(U160::ZERO);
table.add_node(U160::ONE, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)), 1);
table.add_node(U160::ONE, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)), 2);
table.add_node(U160::ONE, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)), 3);
table.add_node(U160::ONE, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)), 4);
assert_eq!(4, table.count_nodes());
// Shouldn't be added anymore since we already have 4 nodes with the same id
table.add_node(U160::ONE, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)), 5);
assert_eq!(4, table.count_nodes());
// Don't add self
table.add_node(U160::ZERO, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)), 5);
assert_eq!(4, table.count_nodes());
let two = U160::ONE + U160::ONE;
table.add_node(two, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)), 5);
assert_eq!(5, table.count_nodes());
table.add_node(two, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)), 6);
table.add_node(two, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)), 7);
table.add_node(two, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)), 8);
// node_id: MAX
table.add_node(U160::MAX, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)), 9);
table.add_node(U160::MAX, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)), 10);
table.add_node(U160::MAX, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)), 11);
table.add_node(U160::MAX, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)), 12);
table.add_node(
U160::MAX - U160::ONE,
IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)),
13,
);
table.add_node(
U160::MAX - U160::ONE,
IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)),
14,
);
table.add_node(
U160::MAX - U160::ONE,
IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)),
15,
);
table.add_node(
U160::MAX - U160::ONE,
IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)),
16,
);
// This one however should too, since we already know 8 nodes with U160::MIN as id
table.add_node(
U160::MAX - (U160::ONE + U160::ONE),
IpAddr::V4(Ipv4Addr::new(0, 0, 0, 1)),
17,
);
assert_eq!(16, table.count_nodes());
}
}

9
torment-peer/Cargo.toml
Unescape Escape View File

@ -0,0 +1,9 @@
[package]
name = "torment-peer"
version = "0.1.0"
authors = ["eater <=@eater.me>"]
edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]

9
torment-peer/src/lib.rs
Unescape Escape View File

@ -0,0 +1,9 @@
struct Peer {}
#[cfg(test)]
mod tests {
#[test]
fn it_works() {
assert_eq!(2 + 2, 4);
}
}
Write Preview
Loading…
Cancel
Save