Contractless/src/rpc/server/structs.rs

use crate::records::memory::enums::ClientType;
use crate::records::memory::response_channels::Command;
use crate::rpc::server::flood_protection::{
    RPC_LONG_WINDOW_LIMIT, RPC_LONG_WINDOW_SECS, RPC_SHORT_WINDOW_SECS,
};
use crate::sled::Db;
use crate::Arc;
use crate::Mutex;
use crate::TcpStream;

pub struct IncomingCommand {
    pub command: u8,
    pub ip: String,
    pub client_type: ClientType,
}

// CombineAndSendDataParams carries the accepted handshake context into
// the response writer and RPC-loop spawner.
pub struct CombineAndSendDataParams {
    pub stream: Arc<Mutex<TcpStream>>,
    pub db: Db,
    pub connections_key: String,
    pub connection_type: String,
    pub wallet_key: String,
    pub map: Arc<Mutex<Command>>,
    pub returned_address: String,
}

// HandshakeTestParams groups the parsed handshake data and shared
// server context needed to validate an incoming peer connection.
pub struct HandshakeTestParams<'a> {
    pub map: Arc<Mutex<Command>>,
    pub stream: Arc<Mutex<TcpStream>>,
    pub count: usize,
    pub peer_time: u32,
    pub timestamp: u32,
    pub incoming_connections: u8,
    pub hash: &'a str,
    pub received_signed_message: &'a str,
    pub received_address: &'a str,
    pub received_ip: &'a str,
}

// RpcFloodState is stored as a compact sled value so flood counters can
// survive across commands without keeping a separate in-memory map.
pub struct RpcFloodState {
    pub short_window_start: i64,
    pub short_window_count: u32,
    pub long_window_start: i64,
    pub long_window_count: u32,
}

impl RpcFloodState {
    pub fn new(now: i64) -> Self {
        // New subjects start both flood windows at the current timestamp
        // with zero counted requests.
        Self {
            short_window_start: now,
            short_window_count: 0,
            long_window_start: now,
            long_window_count: 0,
        }
    }

    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
        // The serialized form is four little-endian fields:
        // i64, u32, i64, u32.
        if bytes.len() != 24 {
            return None;
        }

        Some(Self {
            short_window_start: i64::from_le_bytes(bytes[0..8].try_into().ok()?),
            short_window_count: u32::from_le_bytes(bytes[8..12].try_into().ok()?),
            long_window_start: i64::from_le_bytes(bytes[12..20].try_into().ok()?),
            long_window_count: u32::from_le_bytes(bytes[20..24].try_into().ok()?),
        })
    }

    pub fn to_bytes(&self) -> [u8; 24] {
        // Keep the byte layout stable because old flood tracker rows may
        // still exist while the node is running.
        let mut bytes = [0u8; 24];
        bytes[0..8].copy_from_slice(&self.short_window_start.to_le_bytes());
        bytes[8..12].copy_from_slice(&self.short_window_count.to_le_bytes());
        bytes[12..20].copy_from_slice(&self.long_window_start.to_le_bytes());
        bytes[20..24].copy_from_slice(&self.long_window_count.to_le_bytes());
        bytes
    }

    pub fn record_request(&mut self, now: i64) {
        // Expired windows reset before counting the current request.
        if now.saturating_sub(self.short_window_start) >= RPC_SHORT_WINDOW_SECS {
            self.short_window_start = now;
            self.short_window_count = 0;
        }

        if now.saturating_sub(self.long_window_start) >= RPC_LONG_WINDOW_SECS {
            self.long_window_start = now;
            self.long_window_count = 0;
        }

        self.short_window_count = self.short_window_count.saturating_add(1);
        self.long_window_count = self.long_window_count.saturating_add(1);
    }

    pub fn is_flooding(&self) -> bool {
        // Scoring uses the long window so normal short bursts from real
        // clients do not immediately penalize a peer.
        self.long_window_count > RPC_LONG_WINDOW_LIMIT
    }
}