diff --git a/src/records/memory/connections.rs b/src/records/memory/connections.rs index fee1415..8f368f0 100644 --- a/src/records/memory/connections.rs +++ b/src/records/memory/connections.rs @@ -55,6 +55,7 @@ struct NodeRuntimeContext { lazy_static! { static ref RECOVERY_IN_PROGRESS: StdMutex> = StdMutex::new(HashSet::new()); + static ref SETUP_RETRY_IN_PROGRESS: StdMutex> = StdMutex::new(HashSet::new()); static ref NODE_RUNTIME_CONTEXT: StdRwLock> = StdRwLock::new(None); } @@ -80,6 +81,26 @@ fn try_start_recovery(key: String) -> Option { Some(RecoveryGuard { key }) } +struct SetupRetryGuard { + key: String, +} + +impl Drop for SetupRetryGuard { + fn drop(&mut self) { + if let Ok(mut active) = SETUP_RETRY_IN_PROGRESS.lock() { + active.remove(&self.key); + } + } +} + +fn try_start_setup_retry(key: String) -> Option { + let mut active = SETUP_RETRY_IN_PROGRESS.lock().ok()?; + if !active.insert(key.clone()) { + return None; + } + Some(SetupRetryGuard { key }) +} + pub fn initialize_node_runtime_context( db: Db, wallet: Arc, @@ -1022,18 +1043,22 @@ pub async fn mark_peer_network_map_synced(key: &str) -> bool { } pub async fn mark_peer_operational(key: &str, map: Arc>) -> bool { + if peer_is_operational(key).await { + return true; + } + let setup = CONNECTIONS .write() .await .as_mut() .and_then(|connection| connection.mark_local_setup_complete(key, map.clone())); - let Some((stream, already_acknowledged)) = setup else { + let Some((stream, _already_acknowledged)) = setup else { return false; }; - if already_acknowledged { - return peer_is_operational(key).await; - } + // Setup completion is idempotent. Keep sending it while the peer exists + // until both sides have received and acknowledged the other's setup state. + // A previous local acknowledgement alone does not make the peer operational. for attempt in 1..=3 { let (uid, _tx, rx) = reserve_entry_with_context( map.clone(), @@ -1056,12 +1081,16 @@ pub async fn mark_peer_operational(key: &str, map: Arc>) -> bool }; delete_entry(map.clone(), uid).await; if matches!(response, Ok(Some(bytes)) if bytes == b"setup_complete_ack") { - return CONNECTIONS + let operational = CONNECTIONS .write() .await .as_mut() .map(|connection| connection.mark_local_setup_acknowledged(key, map.clone())) .unwrap_or(false); + if operational { + return true; + } + return peer_is_operational(key).await; } warn!( "[handshake_state] setup-complete acknowledgement failed: peer={key} attempt={attempt}/3" @@ -1082,6 +1111,10 @@ pub async fn mark_peer_remote_setup_complete(key: &str, map: Arc> pub fn spawn_peer_setup_retry(key: String, map: Arc>) { tokio::spawn(async move { + let Some(_retry_guard) = try_start_setup_retry(key.clone()) else { + return; + }; + loop { sleep(Duration::from_secs(10)).await; if Connection::get_stream_from_memory(&key).await.is_none() @@ -1267,4 +1300,17 @@ mod tests { assert!(try_start_recovery("test-peer:first".to_string()).is_some()); drop(second); } + + #[test] + fn setup_retry_is_deduplicated_per_peer() { + let first = try_start_setup_retry("test-setup:first".to_string()) + .expect("first setup retry should start"); + assert!(try_start_setup_retry("test-setup:first".to_string()).is_none()); + + let second = try_start_setup_retry("test-setup:second".to_string()) + .expect("another peer setup retry must run independently"); + drop(first); + assert!(try_start_setup_retry("test-setup:first".to_string()).is_some()); + drop(second); + } } diff --git a/src/rpc/server/command_loop_state.rs b/src/rpc/server/command_loop_state.rs index a6b88c7..cf25868 100644 --- a/src/rpc/server/command_loop_state.rs +++ b/src/rpc/server/command_loop_state.rs @@ -115,7 +115,8 @@ pub async fn next_incoming_command( // maintenance paths. All other commands still count against flood // protection, including non-sync commands from miners. if !skip_flood_check(command, client_type) { - check_request_frequency_with_client_type(db, ip.clone(), client_type, wallet).await; + check_request_frequency_with_client_type(db, ip.clone(), client_type, command, wallet) + .await; } Ok(Some(IncomingCommand { diff --git a/src/rpc/server/flood_protection.rs b/src/rpc/server/flood_protection.rs index b96f417..ceeba53 100644 --- a/src/rpc/server/flood_protection.rs +++ b/src/rpc/server/flood_protection.rs @@ -1,6 +1,10 @@ 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::command_maps::{ + RPC_ADD_NETWORK_NODE, RPC_NETWORK_MONITOR_ADD, RPC_NETWORK_MONITOR_REMOVE, + RPC_NETWORK_MONITOR_STATE, RPC_SETUP_COMPLETE, +}; use crate::rpc::server::structs::RpcFloodState; use crate::sled::Db; use crate::wallets::structures::Wallet; @@ -12,20 +16,36 @@ 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 { +fn request_subject(ip: &str, client_type: ClientType, command: u8) -> 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() + return format!("client:{ip}"); } + + // Join, monitor, and setup traffic can arrive in legitimate bursts when + // peers reconnect. Keep each authenticated miner control command bounded, + // but do not let unrelated control commands accumulate into one false + // flood event for the peer. + if matches!( + command, + RPC_ADD_NETWORK_NODE + | RPC_NETWORK_MONITOR_ADD + | RPC_NETWORK_MONITOR_REMOVE + | RPC_NETWORK_MONITOR_STATE + | RPC_SETUP_COMPLETE + ) { + return format!("miner:{ip}:control:{command}"); + } + + format!("miner:{ip}:rpc") } pub async fn check_request_frequency_with_client_type( db: &Db, ip: String, client_type: ClientType, + command: u8, wallet: Arc, ) { // Keep one compact flood-tracker row per subject and decay the @@ -37,7 +57,7 @@ pub async fn check_request_frequency_with_client_type( }; let now = Utc::now().timestamp(); - let subject = request_subject(&ip, client_type); + let subject = request_subject(&ip, client_type, command); let key = subject.as_bytes(); // Missing or unreadable flood state starts a fresh window instead of @@ -79,3 +99,43 @@ pub async fn check_request_frequency_with_client_type( let _ = tree.insert(key, state.to_bytes().to_vec()); } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::rpc::command_maps::{RPC_TIME, RPC_VALIDATE_ADDRESS}; + + #[test] + fn miner_control_commands_have_independent_bounded_subjects() { + let ip = "203.0.113.10"; + + assert_ne!( + request_subject(ip, ClientType::Miner, RPC_NETWORK_MONITOR_ADD), + request_subject(ip, ClientType::Miner, RPC_NETWORK_MONITOR_REMOVE) + ); + assert_ne!( + request_subject(ip, ClientType::Miner, RPC_NETWORK_MONITOR_STATE), + request_subject(ip, ClientType::Miner, RPC_SETUP_COMPLETE) + ); + } + + #[test] + fn ordinary_miner_commands_still_share_one_flood_subject() { + let ip = "203.0.113.10"; + + assert_eq!( + request_subject(ip, ClientType::Miner, RPC_TIME), + request_subject(ip, ClientType::Miner, RPC_VALIDATE_ADDRESS) + ); + } + + #[test] + fn client_commands_still_share_one_flood_subject() { + let ip = "203.0.113.10"; + + assert_eq!( + request_subject(ip, ClientType::Client, RPC_TIME), + request_subject(ip, ClientType::Client, RPC_SETUP_COMPLETE) + ); + } +}