torment/torment-core/src/lib.rs

#![allow(dead_code)]
use crate::infohash::v1::U160;
use crate::infohash::InfoHashCapable;
use bytes::BytesMut;
use serde::export::fmt::Debug;
use serde_derive::{Deserialize, Serialize};
use std::cmp::min;
use std::convert::TryInto;
use std::fmt::{Display, Formatter};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::str::FromStr;

mod consts;
pub mod infohash;
pub mod ip;
pub mod metainfo;
pub mod peer_storage;
pub mod utils;

pub use consts::*;

pub fn peer_id(input: [u8; 12]) -> [u8; 20] {
    let peer_id = format!(
        "eT{:0>2}{:0>2}{:0>2}AAAAAAAAAAAA",
        env!("CARGO_PKG_VERSION_MAJOR"),
        env!("CARGO_PKG_VERSION_MINOR"),
        env!("CARGO_PKG_VERSION_PATCH")
    );

    let mut peer_id: [u8; 20] = peer_id.as_bytes().try_into().unwrap();
    &peer_id[8..20].copy_from_slice(&input);
    peer_id
}

pub fn compact_peer_id() -> [u8; 4] {
    [
        b'e',
        b'T',
        u8::from_str(env!("CARGO_PKG_VERSION_MAJOR")).unwrap(),
        u8::from_str(env!("CARGO_PKG_VERSION_MINOR")).unwrap(),
    ]
}

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(Clone, Ord, PartialOrd, Eq, PartialEq)]
pub struct Bitfield(BytesMut, usize);

impl Bitfield {
    pub fn size(&self) -> usize {
        self.1
    }

    pub fn new<T: AsRef<[u8]>>(field: T, size: usize) -> Bitfield {
        Bitfield(BytesMut::from(field.as_ref()), size)
    }

    pub fn with_size(size: usize) -> Bitfield {
        Bitfield(
            {
                let mut field = BytesMut::new();
                field.resize((size / 8) + if size % 8 > 0 { 1 } else { 0 }, 0);
                field
            },
            size,
        )
    }

    pub fn set(&mut self, index: u32) {
        debug_assert!(index < self.1 as u32, "index out of bounds");
        let byte_index = index / 8;
        let bitmask = 1 << (7 - (index % 8));
        self.0[byte_index as usize] |= bitmask;
    }

    pub fn unset(&mut self, index: u32) {
        debug_assert!(index < self.1 as u32, "index out of bounds");
        let byte_index = index / 8;
        let bitmask = 1 << (7 - (index % 8));
        self.0[byte_index as usize] &= !bitmask;
    }

    pub fn get(&self, index: u32) -> bool {
        debug_assert!(index < self.1 as u32, "index out of bounds");
        let byte_index = index / 8;
        let bitmask = 1 << (7 - (index % 8));
        (self.0[byte_index as usize] & bitmask) > 0
    }

    pub fn add<T: AsRef<[u8]>>(&mut self, field: T) {
        let field = field.as_ref();
        let max = min(self.0.len(), field.len());
        for i in 0..max {
            self.0[i] |= field[i]
        }
    }

    pub fn all(&self) -> bool {
        for i in 0..self.0.len() {
            let byte = self.0[i];
            if i + 1 == self.0.len() {
                let left = self.1 % 8;
                if left != 0 {
                    return byte == 255 ^ (255 >> left);
                }
            }

            if byte != u8::MAX {
                return false;
            }
        }

        return true;
    }
}

impl Debug for Bitfield {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "Bitfield(")?;
        for i in 0..self.0.len() {
            let byte = self.0[i];
            if i + 1 == self.0.len() {
                write!(
                    f,
                    "{}",
                    &format!("{:0>8b}", byte)[0..(if self.1 % 8 == 0 { 8 } else { self.1 % 8 })]
                )?;
            } else {
                write!(f, "{:b}", byte)?;
            }
        }

        write!(f, ")")
    }
}

impl AsRef<[u8]> for Bitfield {
    fn as_ref(&self) -> &[u8] {
        self.0.as_ref()
    }
}

#[derive(Ord, PartialOrd, Eq, PartialEq, Copy, Clone)]
pub enum LookupFilter {
    IPv6,
    IPv4,
    All,
}

#[cfg(test)]
mod tests {
    use crate::{compact_peer_id, peer_id, Bitfield};
    use rand::random;

    #[test]
    fn test_peer_id() {
        let x = peer_id(random());
        assert_eq!(String::from_utf8_lossy(&x[0..8]), "eT000100");
        assert_eq!(&compact_peer_id(), b"eT\x00\x01");
    }

    #[test]
    fn test_bitfield() {
        let mut field = Bitfield::with_size(3);
        field.set(0);
        assert!(field.get(0));
        field.unset(0);
        assert_eq!(false, field.get(0));
        field.set(2);

        let other_field = Bitfield::with_size(3);

        field.add(other_field);
        assert_eq!(true, field.get(2));

        assert!(!field.all());
        field.set(1);
        field.set(0);

        println!("{} = {:?}", field.as_ref()[0], field);
        assert!(field.all());
    }
}