diff --git a/PeerConnectivityTests/NetworkPeerCoordinatorTests.swift b/PeerConnectivityTests/NetworkPeerCoordinatorTests.swift index b5922a6..0c2cbd9 100644 --- a/PeerConnectivityTests/NetworkPeerCoordinatorTests.swift +++ b/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(.none) let browserObserver = Observable(.none) let advertiserObserver = Observable(.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) } diff --git a/Sources/NetworkPeerCoordinator.swift b/Sources/NetworkPeerCoordinator.swift index 40e10b8..c687f4c 100644 --- a/Sources/NetworkPeerCoordinator.swift +++ b/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 { fileprivate struct PendingConnection { internal let connection : Connection internal let direction : NetworkPeerConnectionDirection + internal let timeout : DispatchWorkItem } fileprivate let localPeer : Peer @@ -22,6 +40,7 @@ internal final class NetworkPeerCoordinator fileprivate let sessionObserver : Observable fileprivate let browserObserver : Observable fileprivate let advertiserObserver : Observable + fileprivate let policy : NetworkPeerConnectionPolicy fileprivate var pendingConnections : [ObjectIdentifier:PendingConnection] = [:] fileprivate var connectionIdentities : [ObjectIdentifier:PeerIdentity] = [:] fileprivate var discoveredPeers : [PeerIdentity:Peer] = [:] @@ -29,13 +48,15 @@ internal final class NetworkPeerCoordinator internal init(localPeer: Peer, sessionObserver: Observable, browserObserver: Observable, - advertiserObserver: Observable) { + advertiserObserver: Observable, + 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 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 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 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 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 } 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 } } diff --git a/Sources/NetworkPeerTransport.swift b/Sources/NetworkPeerTransport.swift index 29d8267..c105b11 100644 --- a/Sources/NetworkPeerTransport.swift +++ b/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) diff --git a/Sources/NetworkPeerTransportAdapters.swift b/Sources/NetworkPeerTransportAdapters.swift index 0124f30..aa1cfe4 100644 --- a/Sources/NetworkPeerTransportAdapters.swift +++ b/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, 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 {