Contractless/src/rpc/server/handshake_verifications.rs

use crate::records::memory::connections::CONNECTIONS;
use crate::records::memory::response_channels::{generate_uid, Command};
use crate::rpc::responses::RpcResponse;
use crate::rpc::server::structs::HandshakeTestParams;
use crate::rpc::server::tests::{ip_test, is_within_one_second};
use crate::wallets::structures::Wallet;
use crate::Arc;
use crate::Mutex;
use crate::TcpStream;

pub async fn connection_count() -> usize {
    // get memory read lock
    let connection_storage = CONNECTIONS.read().await;

    // read number of connected clients or set to 0 if none
    let count = if let Some(connections) = connection_storage.as_ref() {
        connections.count_incoming_connections()
    } else {
        0
    };

    drop(connection_storage);
    count
}

pub async fn check_max_connections(
    _map: Arc<Mutex<Command>>,
    stream: Arc<Mutex<TcpStream>>,
    count: usize,
    incoming_connections: u8,
) -> bool {
    if count >= incoming_connections as usize {
        // Handshake failures still use the normal RPC response framing,
        // so generate a temporary UID for the error packet.
        let hashmap_key = generate_uid();
        let padded_bytes = [hashmap_key[0], hashmap_key[1], hashmap_key[2], 0];
        let uid = u32::from_le_bytes(padded_bytes);

        let response_bytes = RpcResponse::Binary({
            let msg = "error: Sorry all our connections are filled. Please try again later."
                .to_string()
                .as_bytes()
                .to_vec();
            msg
        });
        response_bytes.send(&stream, None, uid).await;
        true
    } else {
        false
    }
}

pub async fn verify_timestamp(
    _map: Arc<Mutex<Command>>,
    stream: Arc<Mutex<TcpStream>>,
    peer_time: u32,
    timestamp: u32,
) -> bool {
    if !is_within_one_second(timestamp, peer_time) {
        // Tight timestamp checks keep peers from replaying stale signed
        // handshakes after the connection attempt has passed.
        let hashmap_key = generate_uid();
        let padded_bytes = [hashmap_key[0], hashmap_key[1], hashmap_key[2], 0];
        let uid = u32::from_le_bytes(padded_bytes);

        let response_bytes = RpcResponse::Binary({
            let msg = format!(
                "error: Handshake Failed: The time on your computer must be within 1 second of the time of our server. Your local time: {peer_time}. Our server time: {timestamp}. Please consider installing NTP to ensure proper timestamps."
            )
            .as_bytes()
            .to_vec();
            msg
        });

        response_bytes.send(&stream, None, uid).await;
        false
    } else {
        true
    }
}

pub async fn verify_ip_address(
    _map: Arc<Mutex<Command>>,
    stream: Arc<Mutex<TcpStream>>,
    received_ip: &str,
) -> bool {
    if !ip_test(received_ip, &stream).await {
        // The peer must announce the same public endpoint it is using
        // for this socket, otherwise it cannot be trusted as a node.
        let hashmap_key = generate_uid();
        let padded_bytes = [hashmap_key[0], hashmap_key[1], hashmap_key[2], 0];
        let uid = u32::from_le_bytes(padded_bytes);

        let response_bytes = RpcResponse::Binary({
            let msg =  "error: Handshake Failed: You must use a public IP, and your computer must be accessible by others. Please check your configuration, firewall settings, or NAT settings.".to_string().as_bytes().to_vec();
            msg
        });

        response_bytes.send(&stream, None, uid).await;
        false
    } else {
        true
    }
}

pub async fn verify_handshake(
    stream: Arc<Mutex<TcpStream>>,
    _map: Arc<Mutex<Command>>,
    hash: &str,
    received_signed_message: &str,
    received_public_key: &str,
) -> bool {
    // Signature verification proves the peer controls the wallet address
    // embedded in the handshake request.
    let hashmap_key = generate_uid();
    let padded_bytes = [hashmap_key[0], hashmap_key[1], hashmap_key[2], 0];
    let uid = u32::from_le_bytes(padded_bytes);

    if !Wallet::verify_transaction_with_public_key(
        hash,
        received_signed_message,
        received_public_key,
    )
    .await
    {
        let response_bytes = RpcResponse::Binary({
            "error: Handshake failed: Invalid handshake"
                .to_string()
                .as_bytes()
                .to_vec()
        });
        response_bytes.send(&stream, None, uid).await;
        return false;
    }
    true
}

pub async fn perform_handshake_tests(params: HandshakeTestParams<'_>) -> bool {
    // Stop at the first failed gate because each check sends the exact
    // error response that explains why the handshake was rejected.
    //check connection limits
    if check_max_connections(
        params.map.clone(),
        params.stream.clone(),
        params.count,
        params.incoming_connections,
    )
    .await
    {
        return false;
    }
    // verify timestmap
    if !verify_timestamp(
        params.map.clone(),
        params.stream.clone(),
        params.peer_time,
        params.timestamp,
    )
    .await
    {
        return false;
    }
    // verify IP
    if !verify_ip_address(
        params.map.clone(),
        params.stream.clone(),
        params.received_ip,
    )
    .await
    {
        return false;
    }
    // verify handshake
    if !verify_handshake(
        params.stream.clone(),
        params.map.clone(),
        params.hash,
        params.received_signed_message,
        params.received_address,
    )
    .await
    {
        return false;
    }
    true
}