Skip to content

Commit 1ce9c9b

Browse files
sanityclaude
andauthored
fix: add self-connection guard in should_accept() (#2140)
Co-authored-by: Claude <noreply@anthropic.com>
1 parent 535d684 commit 1ce9c9b

File tree

1 file changed

+91
-0
lines changed

1 file changed

+91
-0
lines changed

crates/core/src/ring/connection_manager.rs

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,14 @@ impl ConnectionManager {
147147
/// # Panic
148148
/// Will panic if the node checking for this condition has no location assigned.
149149
pub fn should_accept(&self, location: Location, peer_id: &PeerId) -> bool {
150+
// Don't accept connections from ourselves
151+
if let Some(own_id) = self.get_peer_key() {
152+
if &own_id == peer_id {
153+
tracing::warn!(%peer_id, "should_accept: rejecting self-connection attempt");
154+
return false;
155+
}
156+
}
157+
150158
tracing::info!("Checking if should accept connection");
151159
let open = self.connection_count();
152160
let reserved_before = self.pending_reservations.read().len();
@@ -635,3 +643,86 @@ impl ConnectionManager {
635643
read.keys().cloned().collect::<Vec<_>>().into_iter()
636644
}
637645
}
646+
647+
#[cfg(test)]
648+
mod tests {
649+
use super::*;
650+
use crate::node::PeerId;
651+
use crate::topology::rate::Rate;
652+
use crate::transport::TransportKeypair;
653+
use std::net::SocketAddr;
654+
use std::sync::atomic::AtomicU64;
655+
use std::time::Duration;
656+
657+
#[test]
658+
fn rejects_self_connection() {
659+
// Create a peer id for our node
660+
let keypair = TransportKeypair::new();
661+
let addr: SocketAddr = "127.0.0.1:8000".parse().unwrap();
662+
let own_peer_id = PeerId::new(addr, keypair.public().clone());
663+
664+
// Create connection manager with this peer id
665+
let cm = ConnectionManager::init(
666+
Rate::new_per_second(1_000_000.0),
667+
Rate::new_per_second(1_000_000.0),
668+
1,
669+
10,
670+
7,
671+
(
672+
keypair.public().clone(),
673+
Some(own_peer_id.clone()),
674+
AtomicU64::new(u64::from_le_bytes(0.5f64.to_le_bytes())),
675+
),
676+
false,
677+
10,
678+
Duration::from_secs(60),
679+
);
680+
681+
// Verify the connection manager has our peer id set
682+
assert_eq!(cm.get_peer_key(), Some(own_peer_id.clone()));
683+
684+
// Attempt to accept a connection from ourselves - should be rejected
685+
let location = Location::new(0.5);
686+
let accepted = cm.should_accept(location, &own_peer_id);
687+
assert!(!accepted, "should_accept must reject self-connection");
688+
}
689+
690+
#[test]
691+
fn accepts_connection_from_different_peer() {
692+
// Create our node's peer id
693+
let own_keypair = TransportKeypair::new();
694+
let own_addr: SocketAddr = "127.0.0.1:8000".parse().unwrap();
695+
let own_peer_id = PeerId::new(own_addr, own_keypair.public().clone());
696+
697+
// Create connection manager with our peer id
698+
let cm = ConnectionManager::init(
699+
Rate::new_per_second(1_000_000.0),
700+
Rate::new_per_second(1_000_000.0),
701+
1,
702+
10,
703+
7,
704+
(
705+
own_keypair.public().clone(),
706+
Some(own_peer_id.clone()),
707+
AtomicU64::new(u64::from_le_bytes(0.5f64.to_le_bytes())),
708+
),
709+
false,
710+
10,
711+
Duration::from_secs(60),
712+
);
713+
714+
// Create a different peer
715+
let other_keypair = TransportKeypair::new();
716+
let other_addr: SocketAddr = "127.0.0.1:9000".parse().unwrap();
717+
let other_peer_id = PeerId::new(other_addr, other_keypair.public().clone());
718+
719+
// Attempt to accept a connection from a different peer - should succeed
720+
// (first connection with no other constraints)
721+
let location = Location::new(0.6);
722+
let accepted = cm.should_accept(location, &other_peer_id);
723+
assert!(
724+
accepted,
725+
"should_accept must accept connection from different peer"
726+
);
727+
}
728+
}

0 commit comments

Comments
 (0)