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

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