torment/torment-core/src/infohash/v1.rs

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))
    }
}