Contractless/src/rpc/server/handshake_processing.rs

use crate::common::binary_conversions::binary_to_ip_port;
use crate::common::skein::skein_256_hash_data;
use crate::io::ErrorKind;
use crate::log::error;
use crate::rpc::handshake_constants::{
    HANDSHAKE_ADDRESS_OFFSET, HANDSHAKE_IP_OFFSET, HANDSHAKE_MESSAGE_BYTES,
    HANDSHAKE_REQUEST_BYTES, HANDSHAKE_RESPONSE_BYTES, HANDSHAKE_SIGNATURE_OFFSET,
    HANDSHAKE_TIME_OFFSET,
};
use crate::rpc::server::connection_memory_manager::remove_key_from_memory;
use crate::rpc::server::rpc_command_loop::start_loop;
use crate::rpc::server::structs::CombineAndSendDataParams;
use crate::wallets::structures::Wallet;
use crate::Arc;
use crate::AsyncReadExt;
use crate::AsyncWriteExt;
use crate::Mutex;
use crate::TcpStream;
use crate::{decode, encode};

pub async fn parse_received_data(
    stream: Arc<Mutex<TcpStream>>,
) -> Result<(String, String, String, String, String, u32), String> {
    // A server-side handshake request has a fixed binary size, so read
    // the full frame before slicing out individual fields.
    let mut buffer = vec![0u8; HANDSHAKE_REQUEST_BYTES];
    let read_result = stream.lock().await.read_exact(&mut buffer).await;
    match read_result {
        // compiler requires an okay message even though we don't need it
        Ok(_) => {}
        Err(ref err) if err.kind() == ErrorKind::UnexpectedEof => {
            if let Err(e) = stream.lock().await.shutdown().await {
                error!("Error shutting down stream: {e}");
            }
            return Err(err.to_string());
        }
        Err(_) => {
            if let Err(e) = stream.lock().await.shutdown().await {
                error!("Error shutting down stream: {e}");
                return Err(e.to_string());
            }
            return Err("error: Failed to read handshake request".to_string());
        }
    }

    // The first bytes are the fixed challenge message; the signature
    // proves the sender controls the wallet address that follows it.
    let received_message = encode(&buffer[..HANDSHAKE_MESSAGE_BYTES]);
    let received_signed_message =
        encode(&buffer[HANDSHAKE_SIGNATURE_OFFSET..HANDSHAKE_ADDRESS_OFFSET]);
    let received_public_key = encode(&buffer[HANDSHAKE_ADDRESS_OFFSET..HANDSHAKE_TIME_OFFSET]);

    // The timestamp and announced ip:port are kept as raw protocol bytes
    // until this point so the offsets remain explicit.
    let peer_time = u32::from_le_bytes(
        buffer[HANDSHAKE_TIME_OFFSET..HANDSHAKE_IP_OFFSET]
            .try_into()
            .unwrap(),
    );
    let received_ip = binary_to_ip_port(&buffer[HANDSHAKE_IP_OFFSET..HANDSHAKE_REQUEST_BYTES]);
    let hash = skein_256_hash_data(&received_message);

    Ok((
        received_message,
        received_signed_message,
        received_public_key,
        hash,
        received_ip,
        peer_time,
    ))
}

pub async fn generate_and_sign_message(
    connection_type: &str,
    wallet: &Wallet,
) -> Result<(String, String, String), String> {
    let public_key = wallet.saved.public_key.clone();

    // if miner face is the return message, used because its hex so compressed better
    // otherwise face spelledbackwards is used
    if connection_type == "miner" {
        // Miner responses use the opposite challenge from the incoming
        // request so both sides prove they can sign independently.
        let message = "face";
        let hashed = skein_256_hash_data(message);
        let signed_message = Wallet::sign_transaction(&hashed, &wallet.saved.private_key).await;
        Ok((
            public_key.to_string(),
            message.to_string(),
            signed_message.to_string(),
        ))
    } else {
        // Non-miner clients get a client challenge while still using the
        // same wallet-signature verification path.
        let message = "ecaf";
        let hashed = skein_256_hash_data(message);
        let signed_message = Wallet::sign_transaction(&hashed, &wallet.saved.private_key).await;
        Ok((
            public_key.to_string(),
            message.to_string(),
            signed_message.to_string(),
        ))
    }
}

pub async fn return_handshake(
    stream: Arc<Mutex<TcpStream>>,
    public_key: &str,
    message: &str,
    signed_message: &str,
    connections_key: &str,
) -> Result<(), String> {
    // convert to binary/bytes
    let public_key_bin = Wallet::normalize_saved_public_key_bytes(public_key)
        .ok_or_else(|| "error: Invalid public key".to_string())?;
    let message_bin = decode(message).expect("Failed to decode message");
    let signed_bin = decode(signed_message).expect("Failed to decode signed message");
    let mut data = Vec::with_capacity(HANDSHAKE_RESPONSE_BYTES);

    // Response frames mirror the request layout: challenge, signature,
    // then the public key that signed the challenge.
    data.extend_from_slice(&message_bin);
    data.extend_from_slice(&signed_bin);
    data.extend(&public_key_bin);

    // get stream lock
    let mut stream_guard = stream.lock().await;

    // write to stream
    if let Err(err) = stream_guard.write_all(&data).await {
        let _ = remove_key_from_memory(connections_key).await;
        error!("Error writing to stream: {err:?}");
        return Err("error: Error writing to stream".to_string());
    } else if let Err(err) = stream_guard.flush().await {
        error!("Error flushing stream: {err:?}");
        return Err("error: Error flushing stream".to_string());
    }

    // drop lock
    drop(stream_guard);
    Ok(())
}

pub async fn combine_and_send_data(params: CombineAndSendDataParams) -> Result<(), String> {
    let CombineAndSendDataParams {
        stream,
        db,
        connections_key,
        connection_type,
        wallet,
        map,
        returned_address: _,
    } = params;
    // generate the message to send
    let result = generate_and_sign_message(&connection_type, &wallet).await;
    if let Err(err) = result {
        error!("Failed: {err}");
        return Err(err);
    }
    let (public_key, message, signed_message) = result.unwrap();

    // The handshake response must complete before the command loop is
    // spawned, otherwise the peer may start reading command data early.
    if let Err(err) = return_handshake(
        stream.clone(),
        &public_key,
        &message,
        &signed_message,
        &connections_key,
    )
    .await
    {
        error!("Handshake failed: {err}");
        return Err(err);
    }

    // start the rpc loop
    tokio::spawn(start_loop(
        stream.clone(),
        db,
        connections_key.to_string(),
        wallet,
        map,
    ));

    Ok(())
}