Unban peers at the swarm level when purged (#2855)

## Issue Addressed
#2840

beacon_node/lighthouse_network/tests/common/behaviour.rs

@@ -0,0 +1,349 @@
+// NOTE: Taken from libp2p's swarm's testing utils.
+//
+// Copyright 2020 Parity Technologies (UK) Ltd.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+use std::collections::HashMap;
+use std::task::{Context, Poll};
+
+use libp2p::core::connection::{ConnectedPoint, ConnectionId, ListenerId};
+use libp2p::swarm::protocols_handler::{
+    DummyProtocolsHandler, IntoProtocolsHandler, ProtocolsHandler,
+};
+use libp2p::swarm::{DialError, NetworkBehaviour, NetworkBehaviourAction, PollParameters};
+use libp2p::{Multiaddr, PeerId};
+
+/// A `MockBehaviour` is a `NetworkBehaviour` that allows for
+/// the instrumentation of return values, without keeping
+/// any further state.
+pub struct MockBehaviour<
+    THandler = DummyProtocolsHandler,
+    TOutEvent = <DummyProtocolsHandler as ProtocolsHandler>::OutEvent,
+> where
+    THandler: ProtocolsHandler,
+{
+    /// The prototype protocols handler that is cloned for every
+    /// invocation of `new_handler`.
+    pub handler_proto: THandler,
+    /// The addresses to return from `addresses_of_peer`.
+    pub addresses: HashMap<PeerId, Vec<Multiaddr>>,
+    /// The next action to return from `poll`.
+    ///
+    /// An action is only returned once.
+    pub next_action: Option<NetworkBehaviourAction<TOutEvent, THandler>>,
+}
+
+impl<THandler, TOutEvent> MockBehaviour<THandler, TOutEvent>
+where
+    THandler: ProtocolsHandler,
+{
+    pub fn new(handler_proto: THandler) -> Self {
+        MockBehaviour {
+            handler_proto,
+            addresses: HashMap::new(),
+            next_action: None,
+        }
+    }
+}
+
+impl<THandler, TOutEvent> NetworkBehaviour for MockBehaviour<THandler, TOutEvent>
+where
+    THandler: ProtocolsHandler + Clone,
+    THandler::OutEvent: Clone,
+    TOutEvent: Send + 'static,
+{
+    type ProtocolsHandler = THandler;
+    type OutEvent = TOutEvent;
+
+    fn new_handler(&mut self) -> Self::ProtocolsHandler {
+        self.handler_proto.clone()
+    }
+
+    fn addresses_of_peer(&mut self, p: &PeerId) -> Vec<Multiaddr> {
+        self.addresses.get(p).map_or(Vec::new(), |v| v.clone())
+    }
+
+    fn inject_event(&mut self, _: PeerId, _: ConnectionId, _: THandler::OutEvent) {}
+
+    fn poll(
+        &mut self,
+        _: &mut Context,
+        _: &mut impl PollParameters,
+    ) -> Poll<NetworkBehaviourAction<Self::OutEvent, Self::ProtocolsHandler>> {
+        Option::take(&mut self.next_action).map_or(Poll::Pending, Poll::Ready)
+    }
+}
+
+/// A `CallTraceBehaviour` is a `NetworkBehaviour` that tracks invocations of callback methods and
+/// their arguments, wrapping around an inner behaviour. It ensures certain invariants are met.
+pub struct CallTraceBehaviour<TInner>
+where
+    TInner: NetworkBehaviour,
+{
+    inner: TInner,
+
+    pub addresses_of_peer: Vec<PeerId>,
+    pub inject_connected: Vec<PeerId>,
+    pub inject_disconnected: Vec<PeerId>,
+    pub inject_connection_established: Vec<(PeerId, ConnectionId, ConnectedPoint)>,
+    pub inject_connection_closed: Vec<(PeerId, ConnectionId, ConnectedPoint)>,
+    pub inject_event: Vec<(
+        PeerId,
+        ConnectionId,
+        <<TInner::ProtocolsHandler as IntoProtocolsHandler>::Handler as ProtocolsHandler>::OutEvent,
+    )>,
+    pub inject_dial_failure: Vec<Option<PeerId>>,
+    pub inject_new_listener: Vec<ListenerId>,
+    pub inject_new_listen_addr: Vec<(ListenerId, Multiaddr)>,
+    pub inject_new_external_addr: Vec<Multiaddr>,
+    pub inject_expired_listen_addr: Vec<(ListenerId, Multiaddr)>,
+    pub inject_expired_external_addr: Vec<Multiaddr>,
+    pub inject_listener_error: Vec<ListenerId>,
+    pub inject_listener_closed: Vec<(ListenerId, bool)>,
+    pub poll: usize,
+}
+
+impl<TInner> CallTraceBehaviour<TInner>
+where
+    TInner: NetworkBehaviour,
+{
+    pub fn new(inner: TInner) -> Self {
+        Self {
+            inner,
+            addresses_of_peer: Vec::new(),
+            inject_connected: Vec::new(),
+            inject_disconnected: Vec::new(),
+            inject_connection_established: Vec::new(),
+            inject_connection_closed: Vec::new(),
+            inject_event: Vec::new(),
+            inject_dial_failure: Vec::new(),
+            inject_new_listener: Vec::new(),
+            inject_new_listen_addr: Vec::new(),
+            inject_new_external_addr: Vec::new(),
+            inject_expired_listen_addr: Vec::new(),
+            inject_expired_external_addr: Vec::new(),
+            inject_listener_error: Vec::new(),
+            inject_listener_closed: Vec::new(),
+            poll: 0,
+        }
+    }
+
+    #[allow(dead_code)]
+    pub fn reset(&mut self) {
+        self.addresses_of_peer = Vec::new();
+        self.inject_connected = Vec::new();
+        self.inject_disconnected = Vec::new();
+        self.inject_connection_established = Vec::new();
+        self.inject_connection_closed = Vec::new();
+        self.inject_event = Vec::new();
+        self.inject_dial_failure = Vec::new();
+        self.inject_new_listen_addr = Vec::new();
+        self.inject_new_external_addr = Vec::new();
+        self.inject_expired_listen_addr = Vec::new();
+        self.inject_listener_error = Vec::new();
+        self.inject_listener_closed = Vec::new();
+        self.poll = 0;
+    }
+
+    pub fn inner(&mut self) -> &mut TInner {
+        &mut self.inner
+    }
+
+    /// Checks that when the expected number of closed connection notifications are received, a
+    /// given number of expected disconnections have been received as well.
+    ///
+    /// Returns if the first condition is met.
+    pub fn assert_disconnected(
+        &self,
+        expected_closed_connections: usize,
+        expected_disconnections: usize,
+    ) -> bool {
+        if self.inject_connection_closed.len() == expected_closed_connections {
+            assert_eq!(self.inject_disconnected.len(), expected_disconnections);
+            return true;
+        }
+
+        false
+    }
+
+    /// Checks that when the expected number of established connection notifications are received,
+    /// a given number of expected connections have been received as well.
+    ///
+    /// Returns if the first condition is met.
+    pub fn assert_connected(
+        &self,
+        expected_established_connections: usize,
+        expected_connections: usize,
+    ) -> bool {
+        if self.inject_connection_established.len() == expected_established_connections {
+            assert_eq!(self.inject_connected.len(), expected_connections);
+            return true;
+        }
+
+        false
+    }
+}
+
+impl<TInner> NetworkBehaviour for CallTraceBehaviour<TInner>
+where
+    TInner: NetworkBehaviour,
+    <<TInner::ProtocolsHandler as IntoProtocolsHandler>::Handler as ProtocolsHandler>::OutEvent:
+        Clone,
+{
+    type ProtocolsHandler = TInner::ProtocolsHandler;
+    type OutEvent = TInner::OutEvent;
+
+    fn new_handler(&mut self) -> Self::ProtocolsHandler {
+        self.inner.new_handler()
+    }
+
+    fn addresses_of_peer(&mut self, p: &PeerId) -> Vec<Multiaddr> {
+        self.addresses_of_peer.push(*p);
+        self.inner.addresses_of_peer(p)
+    }
+
+    fn inject_connected(&mut self, peer: &PeerId) {
+        assert!(
+            self.inject_connection_established
+                .iter()
+                .any(|(peer_id, _, _)| peer_id == peer),
+            "`inject_connected` is called after at least one `inject_connection_established`."
+        );
+        self.inject_connected.push(*peer);
+        self.inner.inject_connected(peer);
+    }
+
+    fn inject_connection_established(
+        &mut self,
+        p: &PeerId,
+        c: &ConnectionId,
+        e: &ConnectedPoint,
+        errors: Option<&Vec<Multiaddr>>,
+    ) {
+        self.inject_connection_established.push((*p, *c, e.clone()));
+        self.inner.inject_connection_established(p, c, e, errors);
+    }
+
+    fn inject_disconnected(&mut self, peer: &PeerId) {
+        assert!(
+            self.inject_connection_closed
+                .iter()
+                .any(|(peer_id, _, _)| peer_id == peer),
+            "`inject_disconnected` is called after at least one `inject_connection_closed`."
+        );
+        self.inject_disconnected.push(*peer);
+        self.inner.inject_disconnected(peer);
+    }
+
+    fn inject_connection_closed(
+        &mut self,
+        p: &PeerId,
+        c: &ConnectionId,
+        e: &ConnectedPoint,
+        handler: <Self::ProtocolsHandler as IntoProtocolsHandler>::Handler,
+    ) {
+        let connection = (*p, *c, e.clone());
+        assert!(
+            self.inject_connection_established.contains(&connection),
+            "`inject_connection_closed` is called only for connections for \
+            which `inject_connection_established` was called first."
+        );
+        self.inject_connection_closed.push(connection);
+        self.inner.inject_connection_closed(p, c, e, handler);
+    }
+
+    fn inject_event(
+        &mut self,
+        p: PeerId,
+        c: ConnectionId,
+        e: <<Self::ProtocolsHandler as IntoProtocolsHandler>::Handler as ProtocolsHandler>::OutEvent,
+    ) {
+        assert!(
+            self.inject_connection_established
+                .iter()
+                .any(|(peer_id, conn_id, _)| *peer_id == p && c == *conn_id),
+            "`inject_event` is called for reported connections."
+        );
+        assert!(
+            !self
+                .inject_connection_closed
+                .iter()
+                .any(|(peer_id, conn_id, _)| *peer_id == p && c == *conn_id),
+            "`inject_event` is never called for closed connections."
+        );
+
+        self.inject_event.push((p, c, e.clone()));
+        self.inner.inject_event(p, c, e);
+    }
+
+    fn inject_dial_failure(
+        &mut self,
+        p: Option<PeerId>,
+        handler: Self::ProtocolsHandler,
+        error: &DialError,
+    ) {
+        self.inject_dial_failure.push(p);
+        self.inner.inject_dial_failure(p, handler, error);
+    }
+
+    fn inject_new_listener(&mut self, id: ListenerId) {
+        self.inject_new_listener.push(id);
+        self.inner.inject_new_listener(id);
+    }
+
+    fn inject_new_listen_addr(&mut self, id: ListenerId, a: &Multiaddr) {
+        self.inject_new_listen_addr.push((id, a.clone()));
+        self.inner.inject_new_listen_addr(id, a);
+    }
+
+    fn inject_expired_listen_addr(&mut self, id: ListenerId, a: &Multiaddr) {
+        self.inject_expired_listen_addr.push((id, a.clone()));
+        self.inner.inject_expired_listen_addr(id, a);
+    }
+
+    fn inject_new_external_addr(&mut self, a: &Multiaddr) {
+        self.inject_new_external_addr.push(a.clone());
+        self.inner.inject_new_external_addr(a);
+    }
+
+    fn inject_expired_external_addr(&mut self, a: &Multiaddr) {
+        self.inject_expired_external_addr.push(a.clone());
+        self.inner.inject_expired_external_addr(a);
+    }
+
+    fn inject_listener_error(&mut self, l: ListenerId, e: &(dyn std::error::Error + 'static)) {
+        self.inject_listener_error.push(l);
+        self.inner.inject_listener_error(l, e);
+    }
+
+    fn inject_listener_closed(&mut self, l: ListenerId, r: Result<(), &std::io::Error>) {
+        self.inject_listener_closed.push((l, r.is_ok()));
+        self.inner.inject_listener_closed(l, r);
+    }
+
+    fn poll(
+        &mut self,
+        cx: &mut Context,
+        args: &mut impl PollParameters,
+    ) -> Poll<NetworkBehaviourAction<Self::OutEvent, Self::ProtocolsHandler>> {
+        self.poll += 1;
+        self.inner.poll(cx, args)
+    }
+}

beacon_node/lighthouse_network/tests/common/mod.rs

@@ -13,6 +13,13 @@ use std::time::Duration;
 use tokio::runtime::Runtime;
 use types::{ChainSpec, EnrForkId, EthSpec, ForkContext, Hash256, MinimalEthSpec};
 
+#[allow(clippy::type_complexity)]
+#[allow(unused)]
+pub mod behaviour;
+#[allow(clippy::type_complexity)]
+#[allow(unused)]
+pub mod swarm;
+
 type E = MinimalEthSpec;
 use tempfile::Builder as TempBuilder;
 
@@ -41,6 +48,7 @@ impl std::ops::DerefMut for Libp2pInstance {
     }
 }
 
+#[allow(unused)]
 pub fn build_log(level: slog::Level, enabled: bool) -> slog::Logger {
     let decorator = slog_term::TermDecorator::new().build();
     let drain = slog_term::FullFormat::new(decorator).build().fuse();

beacon_node/lighthouse_network/tests/common/swarm.rs

@@ -0,0 +1,99 @@
+use std::collections::HashMap;
+use std::pin::Pin;
+
+use super::behaviour::{CallTraceBehaviour, MockBehaviour};
+
+use futures::stream::Stream;
+use futures::task::{Context, Poll};
+use libp2p::swarm::protocols_handler::ProtocolsHandler;
+use libp2p::swarm::{IntoProtocolsHandler, NetworkBehaviour, Swarm, SwarmBuilder, SwarmEvent};
+use libp2p::{PeerId, Transport};
+
+use futures::StreamExt;
+
+pub fn new_test_swarm<B>(behaviour: B) -> Swarm<B>
+where
+    B: NetworkBehaviour,
+{
+    let id_keys = libp2p::identity::Keypair::generate_ed25519();
+    let local_public_key = id_keys.public();
+    let transport = libp2p::core::transport::MemoryTransport::default()
+        .upgrade(libp2p::core::upgrade::Version::V1)
+        .authenticate(libp2p::plaintext::PlainText2Config {
+            local_public_key: local_public_key.clone(),
+        })
+        .multiplex(libp2p::yamux::YamuxConfig::default())
+        .boxed();
+    SwarmBuilder::new(transport, behaviour, local_public_key.into()).build()
+}
+
+pub fn random_multiaddr() -> libp2p::multiaddr::Multiaddr {
+    libp2p::multiaddr::Protocol::Memory(rand::random::<u64>()).into()
+}
+
+/// Bind a memory multiaddr to a compatible swarm.
+pub async fn bind_listener<B: NetworkBehaviour>(
+    swarm: &mut Swarm<B>,
+) -> libp2p::multiaddr::Multiaddr {
+    swarm.listen_on(random_multiaddr()).unwrap();
+    match swarm.select_next_some().await {
+        SwarmEvent::NewListenAddr {
+            listener_id: _,
+            address,
+        } => address,
+        _ => panic!("Testing swarm's first event should be a new listener"),
+    }
+}
+
+#[derive(Default)]
+pub struct SwarmPool<B: NetworkBehaviour> {
+    swarms: HashMap<PeerId, Swarm<B>>,
+}
+
+impl<B: NetworkBehaviour> SwarmPool<B> {
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self {
+            swarms: HashMap::with_capacity(capacity),
+        }
+    }
+    pub fn insert(&mut self, swarm: Swarm<B>) -> PeerId {
+        let peer_id = *swarm.local_peer_id();
+        self.swarms.insert(peer_id, swarm);
+        peer_id
+    }
+
+    pub fn remove(&mut self, peer_id: &PeerId) {
+        self.swarms.remove(peer_id);
+    }
+
+    pub fn get_mut(&mut self, peer_id: &PeerId) -> Option<&mut Swarm<B>> {
+        self.swarms.get_mut(peer_id)
+    }
+
+    pub fn swarms(&self) -> &HashMap<PeerId, Swarm<B>> {
+        &self.swarms
+    }
+
+    pub fn swarms_mut(&mut self) -> &mut HashMap<PeerId, Swarm<B>> {
+        &mut self.swarms
+    }
+}
+
+impl<B> Stream for SwarmPool<B>
+where
+    B: NetworkBehaviour,
+    <B as NetworkBehaviour>::ProtocolsHandler: ProtocolsHandler,
+{
+    type Item = (PeerId,
+                 SwarmEvent<<B as NetworkBehaviour>::OutEvent, <<<B as NetworkBehaviour>::ProtocolsHandler as IntoProtocolsHandler>::Handler as ProtocolsHandler>::Error>);
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let mut polls = self
+            .get_mut()
+            .swarms
+            .iter_mut()
+            .map(|(&peer_id, swarm)| swarm.map(move |ev| (peer_id, ev)))
+            .collect::<futures::stream::SelectAll<_>>();
+        polls.poll_next_unpin(cx)
+    }
+}