mtorrent-core 0.5.2

use crate::pe;
use bitvec::prelude::*;
use std::mem::MaybeUninit;
use std::net::SocketAddr;
use std::{io, mem};
use tokio::io::{AsyncReadExt, AsyncWriteExt, ReadBuf};

/// Representation of the reserved bits in a handshake.
pub type ReservedBits = BitArray<[u8; 8], Lsb0>;

/// Generate reserved bits.
pub fn reserved_bits(extended_protocol: bool) -> ReservedBits {
    let mut bits = ReservedBits::ZERO;
    bits.set(44, extended_protocol);
    bits
}

/// Check if the extension protocol is enabled.
pub fn is_extension_protocol_enabled(reserved: &ReservedBits) -> bool {
    reserved[44]
}

/// Parsed handshake data.
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct Handshake {
    pub peer_id: [u8; 20],
    pub info_hash: [u8; 20],
    pub reserved: ReservedBits,
}

impl Default for Handshake {
    fn default() -> Self {
        Handshake {
            peer_id: [0u8; 20],
            info_hash: [0u8; 20],
            reserved: BitArray::ZERO,
        }
    }
}

macro_rules! decrypt_if_needed {
    ($crypto:expr, $data:expr) => {
        if let Some(crypto) = $crypto.as_mut() {
            crypto.decryptor.decrypt($data);
        }
    };
}

pub(super) async fn do_handshake_incoming<S>(
    remote_ip: &SocketAddr,
    mut socket: S,
    local_handshake: &Handshake,
    mut crypto: Option<&mut pe::Crypto>,
) -> io::Result<(S, Handshake)>
where
    S: AsyncReadExt + AsyncWriteExt + Unpin,
{
    // Read remote handshake up until peer id,
    // then send entire local handshake (with either local info_hash or remote one),
    // then read remote peer id.
    log::debug!("Receiving incoming handshake from {remote_ip}");

    let mut remote_handshake = Handshake::default();

    read_pstr_and_reserved(&mut socket, &mut remote_handshake.reserved, crypto.as_deref_mut())
        .await?;

    socket.read_exact(&mut remote_handshake.info_hash).await?;
    decrypt_if_needed!(crypto, &mut remote_handshake.info_hash);
    if local_handshake.info_hash != remote_handshake.info_hash {
        return Err(io::Error::other("info_hash doesn't match"));
    }

    write_handshake(&mut socket, local_handshake, crypto.as_mut().map(|c| &mut c.encryptor))
        .await?;

    socket.read_exact(&mut remote_handshake.peer_id).await?;
    decrypt_if_needed!(crypto, &mut remote_handshake.peer_id);
    if remote_handshake.peer_id == local_handshake.peer_id {
        // possible because some trackers include our own external ip
        Err(io::Error::other("incoming connect from ourselves"))
    } else {
        log::trace!(
            "Incoming handshake with {} DONE. Peer id: {}",
            remote_ip,
            String::from_utf8_lossy(&remote_handshake.peer_id[0..8])
        );
        Ok((socket, remote_handshake))
    }
}

pub(super) async fn do_handshake_outgoing<S>(
    remote_ip: &SocketAddr,
    mut socket: S,
    local_handshake: &Handshake,
    expected_remote_peer_id: Option<&[u8; 20]>,
    mut crypto: Option<&mut pe::Crypto>,
) -> io::Result<(S, Handshake)>
where
    S: AsyncReadExt + AsyncWriteExt + Unpin,
{
    // Send entire local handshake,
    // then wait for the entire remote handshake,
    log::debug!("Starting outgoing handshake with {remote_ip}");

    write_handshake(&mut socket, local_handshake, crypto.as_mut().map(|c| &mut c.encryptor))
        .await?;

    let mut remote_handshake = Handshake::default();

    read_pstr_and_reserved(&mut socket, &mut remote_handshake.reserved, crypto.as_deref_mut())
        .await
        .map_err(io::Error::other)?; // convert to Other to avoid reconnect

    socket.read_exact(&mut remote_handshake.info_hash).await?;
    decrypt_if_needed!(crypto, &mut remote_handshake.info_hash);
    if local_handshake.info_hash != remote_handshake.info_hash {
        return Err(io::Error::other("info_hash doesn't match"));
    }

    socket.read_exact(&mut remote_handshake.peer_id).await?;
    decrypt_if_needed!(crypto, &mut remote_handshake.peer_id);
    if let Some(expected_pid) = expected_remote_peer_id
        && expected_pid != &remote_handshake.peer_id
    {
        return Err(io::Error::other("remote peer_id doesn't match"));
    }

    if remote_handshake.peer_id == local_handshake.peer_id {
        Err(io::Error::other("connecting to ourselves"))
    } else {
        log::trace!(
            "Outgoing handshake with {} DONE. Peer id: {}",
            remote_ip,
            String::from_utf8_lossy(&remote_handshake.peer_id[0..8])
        );
        Ok((socket, remote_handshake))
    }
}

pub(crate) const PROTO_STR: &[u8] = b"\x13BitTorrent protocol";

async fn read_pstr_and_reserved<S: AsyncReadExt + Unpin>(
    mut source: S,
    reserved: &mut ReservedBits,
    mut crypto: Option<&mut pe::Crypto>,
) -> io::Result<()> {
    let mut pstr = [0u8; mem::size_of_val(PROTO_STR)];
    source.read_exact(&mut pstr).await?;
    decrypt_if_needed!(crypto, &mut pstr);
    if pstr != PROTO_STR {
        return Err(io::Error::other(format!(
            "Unknown protocol: '{}'",
            String::from_utf8_lossy(&pstr)
        )));
    }
    source.read_exact(&mut reserved.data).await?;
    decrypt_if_needed!(crypto, &mut reserved.data);
    Ok(())
}

const TOTAL_HANDSHAKE_LEN: usize = mem::size_of::<Handshake>() + mem::size_of_val(PROTO_STR);

async fn write_handshake<S: AsyncWriteExt + Unpin>(
    mut sink: S,
    handshake: &Handshake,
    encryptor: Option<&mut pe::Encryptor>,
) -> io::Result<()> {
    let mut buf = [MaybeUninit::uninit(); TOTAL_HANDSHAKE_LEN];
    let mut writer = ReadBuf::uninit(&mut buf);

    writer.put_slice(PROTO_STR);
    writer.put_slice(&handshake.reserved.data);
    writer.put_slice(&handshake.info_hash);
    writer.put_slice(&handshake.peer_id);

    if let Some(enc) = encryptor {
        enc.encrypt(writer.filled_mut())
    }

    sink.write_all(writer.filled()).await?;
    sink.flush().await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
    use tokio::io::duplex;
    use tokio::join;

    const IP: SocketAddr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0));

    #[tokio::test]
    async fn test_handshake_specified_server_info_hash() {
        let (server_stream, client_stream) = duplex(1024);

        let client_hs_data = Handshake {
            peer_id: [1u8; 20],
            info_hash: [7u8; 20],
            reserved: BitArray::from([0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x00, 0x01]),
        };
        let server_hs_data = Handshake {
            peer_id: [2u8; 20],
            info_hash: [7u8; 20],
            reserved: BitArray::ZERO,
        };

        let client_hs_fut = async {
            do_handshake_outgoing(
                &IP,
                client_stream,
                &client_hs_data,
                Some(&server_hs_data.peer_id),
                None,
            )
            .await
            .unwrap()
            .1
        };
        let server_hs_fut = async {
            do_handshake_incoming(&IP, server_stream, &server_hs_data, None)
                .await
                .unwrap()
                .1
        };

        let (received_server_hs, received_client_hs): (Handshake, Handshake) =
            join!(client_hs_fut, server_hs_fut);

        assert_eq!(server_hs_data, received_server_hs);
        assert_eq!(client_hs_data, received_client_hs);
        assert!(received_client_hs.reserved[44]);
        assert!(received_client_hs.reserved[56]);
    }

    #[tokio::test]
    async fn test_handshake_encrypted() {
        let (server_stream, client_stream) = duplex(1024);
        let info_hash = [7u8; 20];

        let (outbound_enc, inbound_dec) = pe::crypto_pair(&info_hash);
        let (inbound_enc, outbound_dec) = pe::crypto_pair(&info_hash);

        let mut outbound_crypto = pe::Crypto {
            encryptor: outbound_enc,
            decryptor: outbound_dec,
        };
        let mut inbound_crypto = pe::Crypto {
            encryptor: inbound_enc,
            decryptor: inbound_dec,
        };

        let client_hs_data = Handshake {
            peer_id: [1u8; 20],
            info_hash: [7u8; 20],
            reserved: BitArray::from([0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x00, 0x01]),
        };
        let server_hs_data = Handshake {
            peer_id: [2u8; 20],
            info_hash: [7u8; 20],
            reserved: BitArray::ZERO,
        };

        let client_hs_fut = async {
            do_handshake_outgoing(
                &IP,
                client_stream,
                &client_hs_data,
                Some(&server_hs_data.peer_id),
                Some(&mut outbound_crypto),
            )
            .await
            .unwrap()
            .1
        };
        let server_hs_fut = async {
            do_handshake_incoming(&IP, server_stream, &server_hs_data, Some(&mut inbound_crypto))
                .await
                .unwrap()
                .1
        };

        let (received_server_hs, received_client_hs): (Handshake, Handshake) =
            join!(client_hs_fut, server_hs_fut);

        assert_eq!(server_hs_data, received_server_hs);
        assert_eq!(client_hs_data, received_client_hs);
        assert!(received_client_hs.reserved[44]);
        assert!(received_client_hs.reserved[56]);
    }

    #[tokio::test]
    async fn test_handshake_peer_id_doesnt_match() {
        let (server_stream, client_stream) = duplex(1024);

        let client_hs_data = Handshake {
            peer_id: [1u8; 20],
            info_hash: [7u8; 20],
            reserved: BitArray::ZERO,
        };
        let server_hs_data = Handshake {
            peer_id: [2u8; 20],
            info_hash: [7u8; 20],
            reserved: BitArray::ZERO,
        };

        let client_hs_fut = async {
            let result =
                do_handshake_outgoing(&IP, client_stream, &client_hs_data, Some(&[0u8; 20]), None)
                    .await;
            let error: io::Error = result.err().unwrap();
            assert_eq!("remote peer_id doesn't match", error.to_string(),)
        };
        let server_hs_fut = async {
            _ = do_handshake_incoming(&IP, server_stream, &server_hs_data, None).await;
        };
        join!(client_hs_fut, server_hs_fut);
    }

    #[tokio::test]
    async fn test_handshake_with_ourselves_returns_io_error_other() {
        let (server_stream, client_stream) = duplex(1024);

        let hs_data = Handshake {
            peer_id: [1u8; 20],
            info_hash: [7u8; 20],
            reserved: BitArray::ZERO,
        };

        let client_hs_fut = async {
            let result = do_handshake_outgoing(&IP, client_stream, &hs_data, None, None).await;
            assert!(result.is_err());
            assert_eq!(result.unwrap_err().kind(), io::ErrorKind::Other);
        };
        let server_hs_fut = async {
            let result = do_handshake_incoming(&IP, server_stream, &hs_data, None).await;
            assert!(result.is_err());
            assert_eq!(result.unwrap_err().kind(), io::ErrorKind::Other);
        };
        join!(client_hs_fut, server_hs_fut);
    }

    #[tokio::test]
    async fn test_handshake_info_hash_doesnt_match() {
        let (server_stream, client_stream) = duplex(1024);

        let client_hs_data = Handshake {
            peer_id: [1u8; 20],
            info_hash: [7u8; 20],
            reserved: BitArray::ZERO,
        };
        let server_hs_data = Handshake {
            peer_id: [2u8; 20],
            info_hash: [8u8; 20],
            reserved: BitArray::ZERO,
        };

        let client_hs_fut = async {
            let result =
                do_handshake_outgoing(&IP, client_stream, &client_hs_data, None, None).await;
            assert!(result.is_err());
        };
        let server_hs_fut = async {
            let result = do_handshake_incoming(&IP, server_stream, &server_hs_data, None).await;
            let error: io::Error = result.err().unwrap();
            assert_eq!("info_hash doesn't match", error.to_string())
        };
        join!(client_hs_fut, server_hs_fut);
    }

    #[tokio::test]
    async fn test_handshake_parse_entire_real_hanshake_message() {
        let server_hs_msg = b"\x13\x42\x69\x74\x54\x6f\x72\x72\x65\x6e\x74\x20\x70\x72\x6f\x74\
            \x6f\x63\x6f\x6c\x00\x00\x00\x00\x00\x10\x00\x05\x74\x4f\x27\x27\
            \xce\x5d\x3c\x4d\x6b\xa4\xcf\x5b\xa7\xac\x08\x78\x46\x0a\x9e\xed\
            \x2d\x42\x54\x37\x61\x35\x57\x2d\x11\xb4\x8d\x05\x19\x2c\x3e\x33\
            \x88\x7c\x4b\xca";

        let (mut server_stream, client_stream) = duplex(1024);

        let client_hs_data = Handshake {
            peer_id: [1u8; 20],
            info_hash:
                *b"\x74\x4f\x27\x27\xce\x5d\x3c\x4d\x6b\xa4\xcf\x5b\xa7\xac\x08\x78\x46\x0a\x9e\xed",
            reserved: BitArray::ZERO,
        };

        let client_hs_fut = async {
            do_handshake_outgoing(&IP, client_stream, &client_hs_data, None, None)
                .await
                .unwrap()
                .1
        };

        let server_fut = async {
            server_stream.write_all(&server_hs_msg[..]).await.unwrap();
        };

        let (received_server_hs, _) = join!(client_hs_fut, server_fut);

        let received_server_hs: Handshake = received_server_hs;

        assert_eq!(*b"\x00\x00\x00\x00\x00\x10\x00\x05", received_server_hs.reserved.data);
        assert_eq!(client_hs_data.info_hash, received_server_hs.info_hash);

        assert_eq!(
            *b"\x2d\x42\x54\x37\x61\x35\x57\x2d\x11\xb4\x8d\x05\x19\x2c\x3e\x33\x88\x7c\x4b\xca",
            received_server_hs.peer_id
        );

        assert!(received_server_hs.reserved[44]); // Extension protocol
        assert!(received_server_hs.reserved[56]); // DHT
        assert!(received_server_hs.reserved[58]); // Fast Extension
    }
}