feat: Support banning peers

bittorrent/src/connection_manager.rs

@@ -1,8 +1,9 @@
 use std::io;
-use std::net::SocketAddr;
+use std::net::{IpAddr, SocketAddr};
 use std::ops::{Index, IndexMut};
 use std::os::fd::{AsRawFd, RawFd};
 
+use ahash::HashSet;
 use bytes::BufMut;
 use heapless::spsc::Producer;
 use rayon::Scope;
@@ -68,6 +69,7 @@ impl ConnectionState {
 
 pub struct ConnectionManager {
     connections: SlotMap<ConnectionId, ConnectionState>,
+    ban_list: HashSet<IpAddr>,
     max_connections: usize,
     our_id: PeerId,
 }
@@ -77,13 +79,18 @@ impl ConnectionManager {
         Self {
             connections: SlotMap::with_capacity_and_key(max_connections),
             max_connections,
+            ban_list: Default::default(),
             our_id,
         }
     }
 
     // The address scan is linear but since it's bound by max_connections, which
     // is expected to be in the hundreds, it should be fast
     fn can_accept_new(&self, peer: SocketAddr) -> bool {
+        if self.ban_list.contains(&peer.ip()) {
+            log::trace!("Ignoring peer that has already been banned ({peer})",);
+            return false;
+        }
         if self.connections.len() >= self.max_connections {
             log::trace!(
                 "Ignoring peer, max connections ({}) reached",
@@ -94,6 +101,29 @@ impl ConnectionManager {
         !self.connections.values().any(|state| state.addr() == peer)
     }
 
+    /// Ban all peers that have the given IP, disconnects any existing peers.
+    /// ip is used instead of SocketAddr since incoming source ports are controlled
+    /// by the peer
+    pub fn ban_peer<'state, Q: SubmissionQueue>(
+        &mut self,
+        peer_ip: IpAddr,
+        io: &mut Io<Q>,
+        state: &mut StateRef<'state>,
+    ) {
+        self.ban_list.insert(peer_ip);
+        for (conn_id, entry) in self.connections.iter_mut() {
+            if entry.addr().ip() != peer_ip {
+                continue;
+            }
+            #[cfg(feature = "metrics")]
+            if let ConnectionState::Established { .. } = &*entry {
+                let counter = metrics::counter!("disconnects", "reason" => "banned");
+                counter.increment(1);
+            }
+            Self::disconnect_entry(entry, conn_id, io, state);
+        }
+    }
+
     /// Attempt to open an outgoing connection to `peer`. No-op if the peer is
     /// already tracked or the connection cap has been reached.
     pub fn maybe_connect_to_peer<Q: SubmissionQueue>(&mut self, peer: SockAddr, io: &mut Io<Q>) {
@@ -160,7 +190,7 @@ impl ConnectionManager {
         let addr = addr.as_socket().expect("must be AF_INET");
         if !self.can_accept_new(addr) {
             log::debug!(
-                "Rejecting incoming connection from already tracked or excess peer: {addr:?}"
+                "Rejecting incoming connection from already tracked, banned or excess peer: {addr:?}"
             );
             io.close_socket(socket, None);
             return Ok(());

bittorrent/src/event_loop.rs

@@ -702,6 +702,10 @@ impl<'scope, 'state: 'scope> EventLoop {
                         );
                     }
                 }
+                Command::BanPeer { peer_ip } => {
+                    log::debug!("Banning peer(s): {peer_ip}");
+                    connection_manager.ban_peer(peer_ip, io, state_ref);
+                }
                 Command::Stop => {
                     if !matches!(self.state, EventLoopState::ShuttingDown { .. }) {
                         log::info!("Shutdown requested, closing all connections");

bittorrent/src/peer_comm/tests.rs

@@ -1,5 +1,5 @@
 use std::{
-    net::{SocketAddr, SocketAddrV4, TcpListener, TcpStream},
+    net::{IpAddr, SocketAddr, SocketAddrV4, TcpListener, TcpStream},
     rc::Rc,
     time::{Duration, Instant},
 };
@@ -4371,14 +4371,20 @@ fn bad_peer_addrs(event_rx: &mut Consumer<'_, TorrentEvent>) -> Vec<SocketAddr>
     addrs
 }
 
+#[track_caller]
+fn insert_established_at(connections: &mut ConnectionManager, addr: SocketAddr) -> ConnectionId {
+    let conn_id = connections.insert_established_with_key(|k| generate_peer(true, k));
+    connections[conn_id].peer_addr = addr;
+    conn_id
+}
+
 #[track_caller]
 fn mark_for_disconnect(
     connections: &mut ConnectionManager,
     addr: &str,
     reason: DisconnectReason,
 ) -> ConnectionId {
-    let conn_id = connections.insert_established_with_key(|k| generate_peer(true, k));
-    connections[conn_id].peer_addr = addr.parse().unwrap();
+    let conn_id = insert_established_at(connections, addr.parse().unwrap());
     connections[conn_id].pending_disconnect = Some(reason);
     conn_id
 }
@@ -4704,6 +4710,186 @@ fn rejected_incoming_connection_is_closed_without_being_tracked() {
     assert_eq!(scheduled_closes(&io), vec![None]);
 }
 
+// Banning a peer tears down the connection it currently has and keeps it out
+// for good, even after the slot of the closed connection has been freed
+#[test]
+fn banned_peer_is_disconnected_and_kept_out() {
+    let mut download_state = setup_test();
+
+    rayon::in_place_scope(|_scope| {
+        let mut state_ref = download_state.as_ref();
+        let mut connections = ConnectionManager::new(PeerId::generate(), 128);
+        let mut io = Io::new(
+            BackloggedSubmissionQueue::new(MockSubmissionQueue),
+            &torrent::Config::default(),
+        );
+
+        let banned: SocketAddrV4 = "127.0.0.1:6881".parse().unwrap();
+        let conn_id = insert_established_at(&mut connections, SocketAddr::V4(banned));
+
+        connections.ban_peer(IpAddr::V4(*banned.ip()), &mut io, &mut state_ref);
+
+        // The established connection is torn down like any other disconnect
+        assert_eq!(scheduled_closes(&io), vec![Some(conn_id)]);
+        assert!(connections.established_mut(conn_id).is_none());
+        assert_eq!(connections.num_established(), 0);
+
+        // Freeing the slot doesn't make the peer connectable again
+        connections.remove_closed(conn_id);
+        assert!(connections.is_empty());
+        connections.maybe_connect_to_peer(banned.into(), &mut io);
+        // Neither does showing up on another port, the ban is on the peer and
+        // not on the address it happened to be reported on
+        let same_peer_other_port: SocketAddrV4 = "127.0.0.1:51413".parse().unwrap();
+        connections.maybe_connect_to_peer(same_peer_other_port.into(), &mut io);
+        assert!(
+            scheduled_connects(&io).is_empty(),
+            "a banned peer should never be connected to again"
+        );
+        assert!(connections.is_empty());
+    });
+}
+
+// The ban applies to connections the peer initiates as well, which is the case
+// that matters most since those are the ones that can't be filtered out before
+// a socket exists
+#[test]
+fn incoming_connection_from_a_banned_peer_is_rejected() {
+    let mut download_state = setup_test();
+
+    rayon::in_place_scope(|_scope| {
+        let mut state_ref = download_state.as_ref();
+        let info_hash = *state_ref.info_hash();
+        let mut connections = ConnectionManager::new(PeerId::generate(), 128);
+        let mut io = Io::new(
+            BackloggedSubmissionQueue::new(MockSubmissionQueue),
+            &torrent::Config::default(),
+        );
+
+        // The peer is banned by the address it listens on, while the connection
+        // it opens to us comes from an ephemeral port that is different every time
+        let listening_addr: SocketAddrV4 = "127.0.0.1:6881".parse().unwrap();
+        connections.ban_peer(IpAddr::V4(*listening_addr.ip()), &mut io, &mut state_ref);
+
+        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+        let client = TcpStream::connect(listener.local_addr().unwrap()).unwrap();
+        assert_ne!(
+            client.local_addr().unwrap().port(),
+            listening_addr.port(),
+            "the connection should come from an ephemeral port"
+        );
+        let (accepted, _) = listener.accept().unwrap();
+        connections
+            .on_accepted(accepted.into(), info_hash, &mut io)
+            .unwrap();
+
+        // Rejected before being tracked, so the close carries no connection id
+        assert!(connections.is_empty());
+        assert_eq!(scheduled_closes(&io), vec![None]);
+    });
+}
+
+// Only the banned peer is affected, connections to everyone else are untouched
+#[test]
+fn banning_a_peer_leaves_other_connections_alone() {
+    let mut download_state = setup_test();
+
+    rayon::in_place_scope(|_scope| {
+        let mut state_ref = download_state.as_ref();
+        let mut connections = ConnectionManager::new(PeerId::generate(), 128);
+        let mut io = Io::new(
+            BackloggedSubmissionQueue::new(MockSubmissionQueue),
+            &torrent::Config::default(),
+        );
+
+        // Both peers are connected to before either of them is banned
+        let banned: SocketAddrV4 = "127.0.0.1:6881".parse().unwrap();
+        let other: SocketAddrV4 = "127.0.0.2:6881".parse().unwrap();
+        connections.maybe_connect_to_peer(banned.into(), &mut io);
+        let [banned_conn] = scheduled_connects(&io)[..] else {
+            panic!("the peer about to be banned should have a connect scheduled");
+        };
+        connections.maybe_connect_to_peer(other.into(), &mut io);
+        let other_connects: Vec<_> = scheduled_connects(&io)
+            .into_iter()
+            .filter(|conn_id| *conn_id != banned_conn)
+            .collect();
+        let [other_conn] = other_connects[..] else {
+            panic!("the other peer should have a connect scheduled");
+        };
+        assert_eq!(connections.total_connections(), 2);
+
+        connections.ban_peer(IpAddr::V4(*banned.ip()), &mut io, &mut state_ref);
+
+        // Only the banned peer is torn down
+        assert_eq!(scheduled_closes(&io), vec![Some(banned_conn)]);
+        assert!(connections.fd(banned_conn).is_none());
+        // while the other one keeps its socket and is left to finish connecting
+        assert!(connections.fd(other_conn).is_some());
+        connections.remove_closed(banned_conn);
+        assert_eq!(connections.total_connections(), 1);
+        // and everyone that isn't banned can still be connected to. The banned
+        // peer's connect event is still counted here, it lives until its cancel
+        // completes, which never happens against a mocked submission queue
+        let not_banned: SocketAddrV4 = "127.0.0.3:6881".parse().unwrap();
+        let connects_before = scheduled_connects(&io).len();
+        connections.maybe_connect_to_peer(not_banned.into(), &mut io);
+        assert_eq!(
+            scheduled_connects(&io).len(),
+            connects_before + 1,
+            "a peer that isn't banned should still be connected to"
+        );
+    });
+}
+
+// Bans show up in the disconnect metrics like every other teardown, and a
+// connection that is already on its way out isn't counted a second time
+#[cfg(feature = "metrics")]
+#[test]
+fn banning_a_peer_records_a_banned_disconnect() {
+    use metrics::{Key, Label};
+    use metrics_util::debugging::{DebugValue, DebuggingRecorder};
+    use metrics_util::{CompositeKey, MetricKind};
+
+    let mut download_state = setup_test();
+    let recorder = DebuggingRecorder::new();
+    let snapshotter = recorder.snapshotter();
+
+    rayon::in_place_scope(|_scope| {
+        let mut state_ref = download_state.as_ref();
+        let mut connections = ConnectionManager::new(PeerId::generate(), 128);
+        let mut io = Io::new(
+            BackloggedSubmissionQueue::new(MockSubmissionQueue),
+            &torrent::Config::default(),
+        );
+
+        let banned: SocketAddrV4 = "127.0.0.1:6881".parse().unwrap();
+        let conn_id = insert_established_at(&mut connections, SocketAddr::V4(banned));
+
+        metrics::with_local_recorder(&recorder, || {
+            connections.ban_peer(IpAddr::V4(*banned.ip()), &mut io, &mut state_ref);
+            // The second ban finds the connection closing, its teardown has
+            // already been counted
+            connections.ban_peer(IpAddr::V4(*banned.ip()), &mut io, &mut state_ref);
+        });
+
+        assert!(connections.established_mut(conn_id).is_none());
+        assert_eq!(scheduled_closes(&io), vec![Some(conn_id)]);
+    });
+
+    let snapshot = snapshotter.snapshot();
+    #[allow(clippy::mutable_key_type)]
+    let metrics = snapshot.into_hashmap();
+    let banned_disconnects = CompositeKey::new(
+        MetricKind::Counter,
+        Key::from_parts("disconnects", vec![Label::new("reason", "banned")]),
+    );
+    let Some((_, _, DebugValue::Counter(disconnects))) = metrics.get(&banned_disconnects) else {
+        panic!("the ban should have been counted as a disconnect: {metrics:?}");
+    };
+    assert_eq!(*disconnects, 1);
+}
+
 // A handshake as it would arrive from a peer, advertising the fast extension
 // (BEP 6) or not. The extension protocol is deliberately left unadvertised so
 // that the only message queued up in response is the initial bitfield one.

bittorrent/src/torrent.rs

@@ -2,7 +2,7 @@ use std::{
     cell::OnceCell,
     collections::VecDeque,
     io::{self},
-    net::{SocketAddr, SocketAddrV4, TcpListener},
+    net::{IpAddr, SocketAddr, SocketAddrV4, TcpListener},
     path::{Path, PathBuf},
     sync::mpsc::{Receiver, Sender},
     time::{Duration, Instant},
@@ -175,6 +175,9 @@ pub enum Command {
     /// Use the `TorrentEvent::Paused` event to determine when the torrent successfully entered
     /// the paused state.
     Resume,
+    /// Ban any peer(s) for the remainder of this run: its current connection, if any,
+    /// is disconnected and it is neither connected to nor accepted again.
+    BanPeer { peer_ip: IpAddr },
 }
 
 /// Metrics for a given peer