Contractless/src/rpc/server/flood_protection.rs

use crate::records::ip_score::enums::InfractionType;
use crate::records::ip_score::score::update_ip_score;
use crate::records::memory::enums::ClientType;
use crate::rpc::server::structs::RpcFloodState;
use crate::sled::Db;
use crate::wallets::structures::Wallet;
use crate::Arc;
use crate::Utc;

pub const MAX_TORRENT_METADATA_BYTES: usize = 8192;
pub const RPC_SHORT_WINDOW_SECS: i64 = 2;
pub const RPC_LONG_WINDOW_SECS: i64 = 60;
pub const RPC_LONG_WINDOW_LIMIT: u32 = 250;

pub fn request_subject(ip: &str, client_type: ClientType) -> String {
    // Wallet-backed clients share an address across changing IPs, so
    // flood tracking treats them differently from long-lived node peers.
    if client_type == ClientType::Client {
        format!("client:{ip}")
    } else {
        ip.to_string()
    }
}

pub async fn check_request_frequency_with_client_type(
    db: &Db,
    ip: String,
    client_type: ClientType,
    wallet: Arc<Wallet>,
) {
    // Keep one compact flood-tracker row per subject and decay the
    // counters by elapsed time so stale request history expires
    // automatically without growing the tree unbounded.
    let tree = match db.open_tree("rpc_flood_tracker") {
        Ok(t) => t,
        Err(_) => return,
    };

    let now = Utc::now().timestamp();
    let subject = request_subject(&ip, client_type);
    let key = subject.as_bytes();

    // Missing or unreadable flood state starts a fresh window instead of
    // blocking the request path on bookkeeping.
    let existing_state = match tree.get(key) {
        Ok(Some(value)) => {
            RpcFloodState::from_bytes(value.as_ref()).unwrap_or_else(|| RpcFloodState::new(now))
        }
        Ok(None) => RpcFloodState::new(now),
        Err(_) => return,
    };

    let mut state = existing_state;
    // Recording the request also rolls expired short/long windows
    // forward before the new count is stored.
    state.record_request(now);

    if tree.insert(key, state.to_bytes().to_vec()).is_err() {
        return;
    }

    // Only sustained long-window flooding is scored; short-window state
    // exists so the serialized tracker can be extended without changing
    // the storage shape again.
    if state.is_flooding() {
        let _ = update_ip_score(
            &ip,
            client_type.as_str(),
            InfractionType::RpcFloodAttack,
            now as u32,
            db,
            wallet,
        )
        .await;

        // Reset the flood window after a penalty so repeated bursts
        // require a fresh full threshold crossing before scoring again.
        state = RpcFloodState::new(now);
        let _ = tree.insert(key, state.to_bytes().to_vec());
    }
}