diff --git a/bittorrent/src/connection_manager.rs b/bittorrent/src/connection_manager.rs index 6a5e69a..65cd590 100644 --- a/bittorrent/src/connection_manager.rs +++ b/bittorrent/src/connection_manager.rs @@ -124,6 +124,10 @@ impl ConnectionManager { } } + pub fn unban_peer(&mut self, peer_ip: IpAddr) { + self.ban_list.remove(&peer_ip); + } + /// 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(&mut self, peer: SockAddr, io: &mut Io) { diff --git a/bittorrent/src/event_loop.rs b/bittorrent/src/event_loop.rs index 144ef21..12cb2de 100644 --- a/bittorrent/src/event_loop.rs +++ b/bittorrent/src/event_loop.rs @@ -706,6 +706,10 @@ impl<'scope, 'state: 'scope> EventLoop { log::debug!("Banning peer(s): {peer_ip}"); connection_manager.ban_peer(peer_ip, io, state_ref); } + Command::UnbanPeer { peer_ip } => { + log::debug!("Unbanning peer(s): {peer_ip}"); + connection_manager.unban_peer(peer_ip); + } Command::Stop => { if !matches!(self.state, EventLoopState::ShuttingDown { .. }) { log::info!("Shutdown requested, closing all connections"); diff --git a/bittorrent/src/peer_comm/tests.rs b/bittorrent/src/peer_comm/tests.rs index 737ba87..0018244 100644 --- a/bittorrent/src/peer_comm/tests.rs +++ b/bittorrent/src/peer_comm/tests.rs @@ -4890,6 +4890,168 @@ fn banning_a_peer_records_a_banned_disconnect() { assert_eq!(*disconnects, 1); } +// Unbanning lifts the ban so the peer can be connected to again, but it does +// not undo the disconnect the ban caused: the connection only comes back once +// the caller asks for it again +#[test] +fn unbanned_peer_can_be_connected_to_again() { + 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 peer: SocketAddrV4 = "127.0.0.1:6881".parse().unwrap(); + let conn_id = insert_established_at(&mut connections, SocketAddr::V4(peer)); + connections.ban_peer(IpAddr::V4(*peer.ip()), &mut io, &mut state_ref); + connections.remove_closed(conn_id); + connections.maybe_connect_to_peer(peer.into(), &mut io); + assert!( + scheduled_connects(&io).is_empty(), + "the peer should still be banned" + ); + + connections.unban_peer(IpAddr::V4(*peer.ip())); + + // Lifting the ban doesn't reconnect on its own, the caller is the one + // that decides when the peer is worth another connection + assert!(connections.is_empty()); + assert!(scheduled_connects(&io).is_empty()); + + connections.maybe_connect_to_peer(peer.into(), &mut io); + assert_eq!( + scheduled_connects(&io).len(), + 1, + "an unbanned peer should be connectable again" + ); + // Reported on another port it is still the same, no longer banned, peer + 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_eq!(scheduled_connects(&io).len(), 2); + }); +} + +// The lifted ban also applies to the connections the peer initiates, which +// arrive from an ephemeral port rather than the address that was unbanned. The +// very same connection attempt is rejected before the unban and kept after it +#[test] +fn incoming_connection_from_an_unbanned_peer_is_accepted() { + 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(), + ); + + // Accepts a connection from the peer, which comes from an ephemeral port + // rather than the address the ban is keyed on + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let listening_addr: SocketAddrV4 = "127.0.0.1:6881".parse().unwrap(); + let connect_to_us = || { + 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(); + // The client end is kept alive for the duration of the test, a peer + // that hung up would be rejected for reasons that have nothing to + // do with the ban + (accepted, client) + }; + + connections.ban_peer(IpAddr::V4(*listening_addr.ip()), &mut io, &mut state_ref); + + let (accepted, _banned_client) = connect_to_us(); + 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(), + "the banned peer should not have been let in" + ); + assert_eq!(scheduled_closes(&io), vec![None]); + + connections.unban_peer(IpAddr::V4(*listening_addr.ip())); + + let (accepted, _unbanned_client) = connect_to_us(); + connections + .on_accepted(accepted.into(), info_hash, &mut io) + .unwrap(); + + assert_eq!(connections.total_connections(), 1); + assert_eq!( + scheduled_closes(&io), + vec![None], + "the connection should have been kept now that the ban is lifted, \ + leaving only the close of the one rejected while banned" + ); + // Kept and handshaked with like any other accepted connection + let [write_event_id] = io + .events + .iter() + .filter_map(|(event_id, event)| match event.typ { + EventType::Write { .. } => Some(event_id), + _ => None, + }) + .collect::>()[..] + else { + panic!("the accepted peer should have a handshake write scheduled"); + }; + // The write completion, which returns the buffer to the pool, is never + // handled against a mocked submission queue so do it here + let handshake_write = io.events.remove(write_event_id).unwrap(); + for buffer in handshake_write.buffers.into_iter().flatten() { + io.write_pool.return_buffer(buffer); + } + }); +} + +// Only the given peer is unbanned, everyone else keeps the ban they were given +#[test] +fn unbanning_a_peer_leaves_other_bans_in_place() { + 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 unbanned: SocketAddrV4 = "127.0.0.1:6881".parse().unwrap(); + let still_banned: SocketAddrV4 = "127.0.0.2:6881".parse().unwrap(); + connections.ban_peer(IpAddr::V4(*unbanned.ip()), &mut io, &mut state_ref); + connections.ban_peer(IpAddr::V4(*still_banned.ip()), &mut io, &mut state_ref); + + connections.unban_peer(IpAddr::V4(*unbanned.ip())); + // Unbanning a peer that was never banned changes nothing either + let never_banned: SocketAddrV4 = "127.0.0.3:6881".parse().unwrap(); + connections.unban_peer(IpAddr::V4(*never_banned.ip())); + + connections.maybe_connect_to_peer(still_banned.into(), &mut io); + assert!( + scheduled_connects(&io).is_empty(), + "the peer that wasn't unbanned should still be banned" + ); + connections.maybe_connect_to_peer(unbanned.into(), &mut io); + connections.maybe_connect_to_peer(never_banned.into(), &mut io); + assert_eq!(scheduled_connects(&io).len(), 2); + }); +} + // 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. diff --git a/bittorrent/src/torrent.rs b/bittorrent/src/torrent.rs index e516ef1..a8c5baa 100644 --- a/bittorrent/src/torrent.rs +++ b/bittorrent/src/torrent.rs @@ -175,9 +175,13 @@ 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. + /// Ban any peer(s) using the given ip. The current connection(s), if any, + /// are disconnected and no new peers from this ip will be connected to or accepted + /// again. NOTE: This includes ips from `Command::ConnectToPeers` messages. + /// Use `UnbanPeer` to reverse this decision. BanPeer { peer_ip: IpAddr }, + /// Unban any peer(s) using the given ip. + UnbanPeer { peer_ip: IpAddr }, } /// Metrics for a given peer