Harden Network backend connection lifecycle

PeerConnectivityTests/NetworkPeerCoordinatorTests.swift

@@ -138,6 +138,56 @@ final class NetworkPeerCoordinatorTests : XCTestCase {
         XCTAssertTrue(harness.coordinator.connectedPeers.isEmpty)
     }
 
+    internal func testHandshakeTimeoutCancelsPendingConnection() {
+        let harness = makeHarness(policy: NetworkPeerConnectionPolicy(handshakeTimeout: 0.01))
+        let connection = MockCoordinatorConnection()
+
+        harness.coordinator.addPendingConnection(connection, direction: .outbound)
+        RunLoop.current.run(until: Date().addingTimeInterval(0.1))
+
+        XCTAssertEqual(connection.cancelCallCount, 1)
+        XCTAssertTrue(harness.coordinator.connectedPeers.isEmpty)
+    }
+
+    internal func testPendingConnectionLimitCancelsExcessConnection() {
+        let harness = makeHarness(policy: NetworkPeerConnectionPolicy(maxPendingConnections: 1))
+        let first = MockCoordinatorConnection()
+        let second = MockCoordinatorConnection()
+
+        harness.coordinator.addPendingConnection(first, direction: .outbound)
+        harness.coordinator.addPendingConnection(second, direction: .outbound)
+
+        XCTAssertEqual(first.cancelCallCount, 0)
+        XCTAssertEqual(second.cancelCallCount, 1)
+    }
+
+    internal func testConnectedPeerLimitCancelsExcessConnection() {
+        let harness = makeHarness(policy: NetworkPeerConnectionPolicy(maxConnectedPeers: 1))
+        let first = MockCoordinatorConnection()
+        let second = MockCoordinatorConnection()
+
+        harness.coordinator.addPendingConnection(first, direction: .outbound)
+        harness.coordinator.receiveFrame(handshakeFrame(identity("first")), from: first)
+        harness.coordinator.addPendingConnection(second, direction: .outbound)
+
+        XCTAssertEqual(second.cancelCallCount, 1)
+        XCTAssertEqual(harness.coordinator.connectedPeers, [Peer(identity: identity("first"), status: .connected)])
+    }
+
+    internal func testConnectedPeerLimitRejectsHandshakeWhenLimitReached() {
+        let harness = makeHarness(policy: NetworkPeerConnectionPolicy(maxPendingConnections: 2, maxConnectedPeers: 1))
+        let first = MockCoordinatorConnection()
+        let second = MockCoordinatorConnection()
+
+        harness.coordinator.addPendingConnection(first, direction: .outbound)
+        harness.coordinator.addPendingConnection(second, direction: .outbound)
+        harness.coordinator.receiveFrame(handshakeFrame(identity("first")), from: first)
+        harness.coordinator.receiveFrame(handshakeFrame(identity("second")), from: second)
+
+        XCTAssertEqual(second.cancelCallCount, 1)
+        XCTAssertEqual(harness.coordinator.connectedPeers, [Peer(identity: identity("first"), status: .connected)])
+    }
+
     internal func testInvalidHandshakeCancelsPendingConnection() {
         let harness = makeHarness()
         let connection = MockCoordinatorConnection()
@@ -215,8 +265,9 @@ final class NetworkPeerCoordinatorTests : XCTestCase {
         XCTAssertEqual(lostPeer, Peer(identity: remoteIdentity, status: .notConnected))
     }
 
-    private func makeHarness(localIdentifier: String = "local") -> Harness {
-        return Harness(localPeer: Peer(identity: identity(localIdentifier), status: .currentUser))
+    private func makeHarness(localIdentifier: String = "local",
+        policy: NetworkPeerConnectionPolicy = NetworkPeerConnectionPolicy()) -> Harness {
+        return Harness(localPeer: Peer(identity: identity(localIdentifier), status: .currentUser), policy: policy)
     }
 
     private func handshakeFrame(_ identity: PeerIdentity,
@@ -266,7 +317,7 @@ private final class Harness {
     internal private(set) var browserEvents : [PeerBrowserEvent] = []
     internal private(set) var advertiserEvents : [PeerAdvertiserEvent] = []
 
-    internal init(localPeer: Peer) {
+    internal init(localPeer: Peer, policy: NetworkPeerConnectionPolicy = NetworkPeerConnectionPolicy()) {
         let sessionObserver = Observable<PeerSessionEvent>(.none)
         let browserObserver = Observable<PeerBrowserEvent>(.none)
         let advertiserObserver = Observable<PeerAdvertiserEvent>(.none)
@@ -275,7 +326,8 @@ private final class Harness {
             localPeer: localPeer,
             sessionObserver: sessionObserver,
             browserObserver: browserObserver,
-            advertiserObserver: advertiserObserver
+            advertiserObserver: advertiserObserver,
+            policy: policy
         )
 
         sessionObserver.addObserver { [weak self] in self?.sessionEvents.append($0) }

Sources/NetworkPeerCoordinator.swift

@@ -8,11 +8,29 @@
 
 import Foundation
 
+internal struct NetworkPeerConnectionPolicy : Equatable {
+    internal let handshakeTimeout : TimeInterval
+    internal let maxPendingConnections : Int
+    internal let maxConnectedPeers : Int
+
+    internal init(handshakeTimeout: TimeInterval = 10,
+        maxPendingConnections: Int = 16,
+        maxConnectedPeers: Int = 8) {
+        precondition(handshakeTimeout > 0, "PeerConnectivity: Network handshake timeout must be positive")
+        precondition(maxPendingConnections > 0, "PeerConnectivity: Network pending connection limit must be positive")
+        precondition(maxConnectedPeers > 0, "PeerConnectivity: Network connected peer limit must be positive")
+        self.handshakeTimeout = handshakeTimeout
+        self.maxPendingConnections = maxPendingConnections
+        self.maxConnectedPeers = maxConnectedPeers
+    }
+}
+
 internal final class NetworkPeerCoordinator<Connection: NetworkPeerFrameSending> {
 
     fileprivate struct PendingConnection {
         internal let connection : Connection
         internal let direction : NetworkPeerConnectionDirection
+        internal let timeout : DispatchWorkItem
     }
 
     fileprivate let localPeer : Peer
@@ -22,20 +40,23 @@ internal final class NetworkPeerCoordinator<Connection: NetworkPeerFrameSending>
     fileprivate let sessionObserver : Observable<PeerSessionEvent>
     fileprivate let browserObserver : Observable<PeerBrowserEvent>
     fileprivate let advertiserObserver : Observable<PeerAdvertiserEvent>
+    fileprivate let policy : NetworkPeerConnectionPolicy
     fileprivate var pendingConnections : [ObjectIdentifier:PendingConnection] = [:]
     fileprivate var connectionIdentities : [ObjectIdentifier:PeerIdentity] = [:]
     fileprivate var discoveredPeers : [PeerIdentity:Peer] = [:]
 
     internal init(localPeer: Peer,
         sessionObserver: Observable<PeerSessionEvent>,
         browserObserver: Observable<PeerBrowserEvent>,
-        advertiserObserver: Observable<PeerAdvertiserEvent>) {
+        advertiserObserver: Observable<PeerAdvertiserEvent>,
+        policy: NetworkPeerConnectionPolicy = NetworkPeerConnectionPolicy()) {
         self.localPeer = localPeer
         self.registry = NetworkPeerConnectionRegistry(localIdentity: localPeer.identity)
         self.dataSender = NetworkPeerDataSender(registry: registry)
         self.sessionObserver = sessionObserver
         self.browserObserver = browserObserver
         self.advertiserObserver = advertiserObserver
+        self.policy = policy
     }
 
     internal var connectedPeers : [Peer] {
@@ -46,7 +67,21 @@ internal final class NetworkPeerCoordinator<Connection: NetworkPeerFrameSending>
 
     internal func addPendingConnection(_ connection: Connection, direction: NetworkPeerConnectionDirection) {
         queue.sync {
-            pendingConnections[ObjectIdentifier(connection)] = PendingConnection(connection: connection, direction: direction)
+            guard registry.connectedPeerIdentities.count < policy.maxConnectedPeers,
+                pendingConnections.count < policy.maxPendingConnections else {
+                connection.cancel()
+                return
+            }
+            let identifier = ObjectIdentifier(connection)
+            pendingConnections[identifier]?.timeout.cancel()
+            let timeout = DispatchWorkItem { [weak self, weak connection] in
+                guard let connection = connection else { return }
+                self?.expirePendingConnection(connection)
+            }
+            pendingConnections[identifier] = PendingConnection(connection: connection,
+                direction: direction,
+                timeout: timeout)
+            queue.asyncAfter(deadline: .now() + policy.handshakeTimeout, execute: timeout)
             sendHandshake(on: connection)
         }
     }
@@ -73,7 +108,7 @@ internal final class NetworkPeerCoordinator<Connection: NetworkPeerFrameSending>
     internal func removeConnection(_ connection: Connection) {
         let event : PeerSessionEvent? = queue.sync {
             let identifier = ObjectIdentifier(connection)
-            pendingConnections.removeValue(forKey: identifier)
+            pendingConnections.removeValue(forKey: identifier)?.timeout.cancel()
             guard let identity = connectionIdentities.removeValue(forKey: identifier) else { return nil }
             guard registry.connection(for: identity) === connection else { return nil }
             registry.remove(identity: identity)
@@ -85,7 +120,10 @@ internal final class NetworkPeerCoordinator<Connection: NetworkPeerFrameSending>
 
     internal func cancelAllConnections() {
         queue.sync {
-            pendingConnections.values.forEach { $0.connection.cancel() }
+            pendingConnections.values.forEach {
+                $0.timeout.cancel()
+                $0.connection.cancel()
+            }
             pendingConnections.removeAll()
             connectionIdentities.removeAll()
             registry.cancelAll()
@@ -127,6 +165,11 @@ internal final class NetworkPeerCoordinator<Connection: NetworkPeerFrameSending>
 
         let identifier = ObjectIdentifier(connection)
         let pending = pendingConnections.removeValue(forKey: identifier)
+        pending?.timeout.cancel()
+        guard registry.connectedPeerIdentities.count < policy.maxConnectedPeers || registry.connection(for: handshake.identity) != nil else {
+            connection.cancel()
+            return nil
+        }
         let direction = pending?.direction ?? NetworkPeerConnectionDirection.inbound
         let wasConnected = registry.connection(for: handshake.identity) != nil
         let isRegistered = registry.register(connection, for: handshake.identity, direction: direction)
@@ -139,10 +182,17 @@ internal final class NetworkPeerCoordinator<Connection: NetworkPeerFrameSending>
     }
 
     fileprivate func rejectHandshake(from connection: Connection) {
-        pendingConnections.removeValue(forKey: ObjectIdentifier(connection))
+        pendingConnections.removeValue(forKey: ObjectIdentifier(connection))?.timeout.cancel()
         connection.cancel()
     }
 
+    fileprivate func expirePendingConnection(_ connection: Connection) {
+        let identifier = ObjectIdentifier(connection)
+        guard let pending = pendingConnections.removeValue(forKey: identifier) else { return }
+        pending.timeout.cancel()
+        pending.connection.cancel()
+    }
+
     fileprivate func removeConnectionIdentity(for identity: PeerIdentity) {
         connectionIdentities = connectionIdentities.filter { $0.value != identity }
     }

Sources/NetworkPeerTransport.swift

@@ -30,7 +30,7 @@ internal final class NetworkPeerConnection : NetworkPeerConnectionCancellable {
 
     fileprivate let connection : NWConnection
     fileprivate let queue : DispatchQueue
-    fileprivate let stateHandler : StateHandler?
+    fileprivate var stateHandler : StateHandler?
     fileprivate var dataHandler : DataHandler?
     fileprivate var frameDecoder = PeerNetworkFrameDecoder()
 
@@ -59,6 +59,10 @@ internal final class NetworkPeerConnection : NetworkPeerConnectionCancellable {
         self.dataHandler = dataHandler
     }
 
+    internal func setStateHandler(_ stateHandler: StateHandler?) {
+        self.stateHandler = stateHandler
+    }
+
     internal func start() {
         connection.stateUpdateHandler = { [weak self] state in
             self?.stateHandler?(state)

Sources/NetworkPeerTransportAdapters.swift

@@ -96,19 +96,13 @@ internal final class NetworkPeerSessionTransport : PeerSessionTransport {
     }
 
     internal func addInboundConnection(_ connection: NetworkPeerConnection) {
-        connection.setDataHandler { [weak self, weak connection] frame in
-            guard let connection = connection else { return }
-            self?.coordinator.receiveFrame(frame, from: connection)
-        }
+        configureConnection(connection)
         coordinator.addPendingConnection(connection, direction: .inbound)
     }
 
     internal func connect(to endpoint: NWEndpoint) {
         let connection = NetworkPeerConnection(endpoint: endpoint, security: security)
-        connection.setDataHandler { [weak self, weak connection] frame in
-            guard let connection = connection else { return }
-            self?.coordinator.receiveFrame(frame, from: connection)
-        }
+        configureConnection(connection)
         coordinator.addPendingConnection(connection, direction: .outbound)
         connection.start()
     }
@@ -121,6 +115,21 @@ internal final class NetworkPeerSessionTransport : PeerSessionTransport {
         coordinator.lostPeer(identity: identity)
     }
 
+    fileprivate func configureConnection(_ connection: NetworkPeerConnection) {
+        connection.setDataHandler { [weak self, weak connection] frame in
+            guard let connection = connection else { return }
+            self?.coordinator.receiveFrame(frame, from: connection)
+        }
+        connection.setStateHandler { [weak self, weak connection] state in
+            guard let connection = connection else { return }
+            switch state {
+            case .failed, .cancelled:
+                self?.coordinator.removeConnection(connection)
+            default: break
+            }
+        }
+    }
+
     fileprivate static func makeListener(peer: Peer,
         coordinator: NetworkPeerCoordinator<NetworkPeerConnection>,
         serviceType: ServiceType,
@@ -134,6 +143,14 @@ internal final class NetworkPeerSessionTransport : PeerSessionTransport {
                     guard let connection = connection else { return }
                     coordinator?.receiveFrame(frame, from: connection)
                 }
+                connection.setStateHandler { [weak coordinator, weak connection] state in
+                    guard let connection = connection else { return }
+                    switch state {
+                    case .failed, .cancelled:
+                        coordinator?.removeConnection(connection)
+                    default: break
+                    }
+                }
                 coordinator.addPendingConnection(connection, direction: .inbound)
             })
         } catch let error {