fixed flood protection reconnect bug

This commit is contained in:
viraladmin 2026-07-16 14:19:16 -06:00
parent 800729bad7
commit 966d3f167f
3 changed files with 118 additions and 11 deletions

View File

@ -55,6 +55,7 @@ struct NodeRuntimeContext {
lazy_static! {
static ref RECOVERY_IN_PROGRESS: StdMutex<HashSet<String>> = StdMutex::new(HashSet::new());
static ref SETUP_RETRY_IN_PROGRESS: StdMutex<HashSet<String>> = StdMutex::new(HashSet::new());
static ref NODE_RUNTIME_CONTEXT: StdRwLock<Option<NodeRuntimeContext>> = StdRwLock::new(None);
}
@ -80,6 +81,26 @@ fn try_start_recovery(key: String) -> Option<RecoveryGuard> {
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<SetupRetryGuard> {
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<Wallet>,
@ -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<Mutex<Command>>) -> 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<Mutex<Command>>) -> 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<Mutex<Command>>
pub fn spawn_peer_setup_retry(key: String, map: Arc<Mutex<Command>>) {
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);
}
}

View File

@ -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 {

View File

@ -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<Wallet>,
) {
// 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)
);
}
}