Add a maximum limit to the number of listeners (#809)

* add max_listeners to swarm

* add swarm builder

* swarm_builder's build takes ownership of self

* replace max listeners with incoming limit

* don't disconnect from node after incoming limit has been reached

* update code according to recent changes

* don't poll listeners at all if incoming connection limit is reached

core/src/nodes/raw_swarm.rs

@@ -66,6 +66,9 @@ where
     /// The reach attempts of the swarm.
     /// This needs to be a separate struct in order to handle multiple mutable borrows issues.
     reach_attempts: ReachAttempts,
+
+    /// Max numer of incoming connections.
+    incoming_limit: Option<u32>,
 }
 
 #[derive(Debug)]
@@ -637,6 +640,25 @@ where
                 other_reach_attempts: Vec::new(),
                 connected_points: Default::default(),
             },
+            incoming_limit: None,
+        }
+    }
+
+    /// Creates a new node event stream with incoming connections limit.
+    #[inline]
+    pub fn new_with_incoming_limit(transport: TTrans,
+        local_peer_id: PeerId, incoming_limit: Option<u32>) -> Self
+    {
+        RawSwarm {
+            incoming_limit,
+            listeners: ListenersStream::new(transport),
+            active_nodes: CollectionStream::new(),
+            reach_attempts: ReachAttempts {
+                local_peer_id,
+                out_reach_attempts: Default::default(),
+                other_reach_attempts: Vec::new(),
+                connected_points: Default::default(),
+            },
         }
     }
 
@@ -658,6 +680,12 @@ where
         self.listeners.listeners()
     }
 
+    /// Returns limit on incoming connections.
+    #[inline]
+    pub fn incoming_limit(&self) -> Option<u32> {
+        self.incoming_limit
+    }
+
     /// Call this function in order to know which address remotes should dial in order to access
     /// your local node.
     ///
@@ -863,26 +891,39 @@ where
         <THandler::Handler as NodeHandler>::OutboundOpenInfo: Send + 'static, // TODO: shouldn't be necessary
         THandlerErr: error::Error + Send + 'static,
     {
-        // Start by polling the listeners for events.
-        match self.listeners.poll() {
-            Async::NotReady => (),
-            Async::Ready(ListenersEvent::Incoming { upgrade, listen_addr, send_back_addr }) => {
-                let event = IncomingConnectionEvent {
-                    upgrade,
-                    local_peer_id: self.reach_attempts.local_peer_id.clone(),
-                    listen_addr,
-                    send_back_addr,
-                    active_nodes: &mut self.active_nodes,
-                    other_reach_attempts: &mut self.reach_attempts.other_reach_attempts,
-                };
-                return Async::Ready(RawSwarmEvent::IncomingConnection(event));
-            }
-            Async::Ready(ListenersEvent::Closed { listen_addr, listener, result }) => {
-                return Async::Ready(RawSwarmEvent::ListenerClosed {
-                    listen_addr,
-                    listener,
-                    result,
-                });
+        // Start by polling the listeners for events, but only
+        // if numer of incoming connection does not exceed the limit.
+        match self.incoming_limit {
+            Some(x) if self.incoming_negotiated().count() >= (x as usize)
+                => (),
+            _ => {
+                match self.listeners.poll() {
+                    Async::NotReady => (),
+                    Async::Ready(ListenersEvent::Incoming {
+                        upgrade, listen_addr, send_back_addr }) =>
+                    {
+                        let event = IncomingConnectionEvent {
+                            upgrade,
+                            local_peer_id:
+                                self.reach_attempts.local_peer_id.clone(),
+                            listen_addr,
+                            send_back_addr,
+                            active_nodes: &mut self.active_nodes,
+                            other_reach_attempts: &mut self.reach_attempts.other_reach_attempts,
+                        };
+                        return Async::Ready(RawSwarmEvent::IncomingConnection(event));
+                     },
+                    Async::Ready(ListenersEvent::Closed {
+                        listen_addr, listener, result }) =>
+                    {
+                        return Async::Ready(RawSwarmEvent::ListenerClosed {
+                            listen_addr,
+                            listener,
+                            result,
+                        });
+                    }
+
+                }
             }
         }
 

core/src/nodes/raw_swarm/tests.rs

@@ -454,3 +454,48 @@ fn local_prio_equivalence_relation() {
         assert_ne!(has_dial_prio(&a, &b), has_dial_prio(&b, &a));
     }
 }
+
+#[test]
+fn limit_incoming_connections() {
+    let mut transport = DummyTransport::new();
+    let peer_id = PeerId::random();
+    let muxer = DummyMuxer::new();
+    let limit = 1;
+    transport.set_initial_listener_state(ListenerState::Ok(Async::Ready(
+        Some((peer_id, muxer)))));
+    let mut swarm = RawSwarm::<_, _, _, Handler, _>::new_with_incoming_limit(
+        transport, PeerId::random(), Some(limit));
+    assert_eq!(swarm.incoming_limit(), Some(limit));
+    swarm.listen_on("/memory".parse().unwrap()).unwrap();
+    assert_eq!(swarm.incoming_negotiated().count(), 0);
+
+    let swarm = Arc::new(Mutex::new(swarm));
+    let mut rt = Runtime::new().unwrap();
+    for i in 1..10 {
+        let swarm_fut = swarm.clone();
+        let fut = future::poll_fn(move || -> Poll<_, ()> {
+            let mut swarm_fut = swarm_fut.lock();
+            if i <= limit {
+                assert_matches!(swarm_fut.poll(),
+                    Async::Ready(RawSwarmEvent::IncomingConnection(incoming)) => {
+                        incoming.accept(Handler::default());
+                });
+            } else {
+                match swarm_fut.poll() {
+                    Async::NotReady => (),
+                    Async::Ready(x) => {
+                        match x {
+                            RawSwarmEvent::IncomingConnection(_) => (),
+                            RawSwarmEvent::Connected { .. } => (),
+                            _ => { panic!("Not expected event") },
+                        }
+                    },
+                }
+             }
+            Ok(Async::Ready(()))
+        });
+        rt.block_on(fut).expect("tokio works");
+        let swarm = swarm.lock();
+        assert!(swarm.incoming_negotiated().count() <= (limit as usize));
+    }
+}