Contractless/src/rpc/server/handshake.rs

use crate::records::memory::response_channels::Command;
use crate::records::memory::response_channels::generate_uid;
use crate::rpc::responses::RpcResponse;
use crate::rpc::server::connection_memory_manager::write_to_memory;
use crate::rpc::server::handshake_processing::{combine_and_send_data, parse_received_data};
use crate::rpc::server::structs::{CombineAndSendDataParams, HandshakeTestParams};
use crate::rpc::server::handshake_verifications::{connection_count, perform_handshake_tests};
use crate::rpc::server::tests::{endpoint_port, is_port_open};
use crate::sled::Db;
use crate::Arc;
use crate::AsyncWriteExt;
use crate::Mutex;
use crate::Settings;
use crate::TcpStream;
use crate::Utc;

async fn drop_failed_handshake(stream: &Arc<Mutex<TcpStream>>) {
    // Failed handshakes are never stored in connection memory, but the
    // accepted TCP socket should still be closed immediately.
    let mut stream_guard = stream.lock().await;
    let _ = stream_guard.flush().await;
    let _ = stream_guard.shutdown().await;
}

async fn get_connection_counts() -> (u8, u8) {
    // Handshake limits come from settings so the node can change its
    // connection policy without recompiling.
    let settings = Settings::load().expect("Failed to load settings");
    let incoming = settings.incoming_connections;
    let outgoing = settings.outgoing_connections;
    (incoming, outgoing)
}

// this function validates incoming handshake and determined
// what type of connection was made
pub async fn handle_handshake(
    stream: Arc<Mutex<TcpStream>>,
    db: Db,
    wallet_key: String,
    map: Arc<Mutex<Command>>,
) {
    // read number of connected clients or set to 0 if none
    let count = connection_count().await;

    // Only incoming capacity matters here; outgoing is loaded with the
    // same settings call but enforced by the connection starter.
    let (incoming_connections, _outgoing_connection) = get_connection_counts().await;

    // get data from stream
    let Ok((
        received_message,
        received_signed_message,
        received_address,
        hash,
        received_ip,
        peer_time,
    )) = parse_received_data(stream.clone()).await
    else {
        return;
    };

    // get local timestamp
    let timestamp = Utc::now().timestamp() as u32;

    // received message should be "aced"
    // aced is used instead of ping and pong
    // as its HEX and compressed better
    if received_message == "aced" {
        //validate handshake tests
        if !perform_handshake_tests(HandshakeTestParams {
            map: map.clone(),
            stream: stream.clone(),
            count,
            peer_time,
            timestamp,
            incoming_connections,
            hash: &hash,
            received_signed_message: &received_signed_message,
            received_address: &received_address,
            received_ip: &received_ip,
        })
        .await
        {
            // Each failed test sends its own error response, so the
            // handshake can stop here without writing another message.
            drop_failed_handshake(&stream).await;
            return;
        }

        // Port 0 is the explicit client marker. A node advertising a
        // nonzero miner port must actually be reachable before it can be
        // stored in connection memory or the network map.
        let Some(advertised_port) = endpoint_port(&received_ip) else {
            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("error: Invalid advertised endpoint.".as_bytes().to_vec());
            response_bytes.send(&stream, None, uid).await;
            drop_failed_handshake(&stream).await;
            return;
        };

        let connection_type = if advertised_port == 0 {
            "client"
        } else if is_port_open(&received_ip).await.unwrap_or(false) {
            "miner"
        } else {
            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(
                "error: Handshake failed: advertised miner port is not reachable."
                    .as_bytes()
                    .to_vec(),
            );
            response_bytes.send(&stream, None, uid).await;
            drop_failed_handshake(&stream).await;
            return;
        };

        if connection_type == "miner" {
            let miner_reserved_limit = incoming_connections.saturating_sub(1) as usize;
            let current_count = connection_count().await;
            if current_count >= miner_reserved_limit {
                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(
                    "error: Miner connection slots are filled. Please try again later."
                        .as_bytes()
                        .to_vec(),
                );
                response_bytes.send(&stream, None, uid).await;
                drop_failed_handshake(&stream).await;
                return;
            }
        }

        // write to memory
        let connections_key = write_to_memory(
            &received_ip,
            stream.clone(),
            connection_type,
            &received_address,
            map.clone(),
        )
        .await;

        if connections_key != "false" {
            // Once the peer is accepted into memory, return our signed
            // handshake response and start the long-lived RPC loop.
            let params = CombineAndSendDataParams {
                stream,
                db: db.clone(),
                connections_key,
                connection_type: connection_type.to_string(),
                wallet_key: wallet_key.clone(),
                map,
                returned_address: received_address.clone(),
            };
            let _ = combine_and_send_data(params).await;
        } else {
            drop_failed_handshake(&stream).await;
        }
    } else {
        let response_bytes = RpcResponse::Binary({
            "error: Invalid Handshake: Signature Failed"
                .to_string()
                .as_bytes()
                .to_vec()
        });
        response_bytes.send(&stream, None, 0).await;
        drop_failed_handshake(&stream).await;
    }
}