diff --git a/Magic Switch/Manager/IncomingConnection.swift b/Magic Switch/Manager/IncomingConnection.swift index 48efcdf..f65acc0 100644 --- a/Magic Switch/Manager/IncomingConnection.swift +++ b/Magic Switch/Manager/IncomingConnection.swift @@ -68,6 +68,7 @@ final class IncomingConnection { private var selfRef: IncomingConnection? private var authenticated = false private var finished = false + private var pendingCounted = false /// Fingerprint of the accept-time PSK snapshot the handshake proved. private var provedFingerprint: String? @@ -108,6 +109,14 @@ final class IncomingConnection { return } + guard rateLimiter.beginPending(endpoint: endpoint) else { + print("Rejecting connection: too many concurrent unauthenticated handshakes") + connection.cancel() + release() + return + } + pendingCounted = true + let channel = SecureChannel( connection: connection, role: .server, psk: psk, queue: queue ) @@ -132,6 +141,7 @@ final class IncomingConnection { switch result { case .success: self.authenticated = true + self.endPendingIfNeeded() // The peer has proved possession of the pairing key this handshake // ran with — strictly stronger evidence than the fingerprint it // advertises over cleartext mDNS. If a registered device is stuck @@ -179,12 +189,22 @@ final class IncomingConnection { self.teardown() case .success(let data): self.resetIdleTimer() + guard self.currentlyAuthorized() else { + print("Dropping frame: pairing key changed or removed since handshake") + self.teardown() + return + } self.handleIncoming(data: data) self.readNext() } } } + private func currentlyAuthorized() -> Bool { + guard let key = pairingStore.currentKey() else { return false } + return PairingStore.fingerprint(forKey: key) == provedFingerprint + } + // MARK: - Command Handling private func handleIncoming(data: Data) { @@ -356,16 +376,24 @@ final class IncomingConnection { sendString(DeviceCommand.operationFailed.rawValue) break } + let store = bluetoothStore + let address = message // Read-only query: OP_SUCCESS only if we have a live connection to it, // so the peer's wake-time reclaim won't grab a peripheral we're using. // The BT check completes on the Bluetooth queue; hop back to the // connection queue so all sealed sends stay serialized there (the send // counter isn't synchronized across queues). - bluetoothStore.isHoldingPeripheral(address: message) { [weak self] held in + DispatchQueue.main.async { [weak self] in guard let self = self else { return } - self.queue.async { - self.sendString( - (held ? DeviceCommand.operationSuccess : DeviceCommand.operationFailed).rawValue) + guard store.peripherals.contains(where: { $0.id == address }) else { + self.queue.async { self.sendString(DeviceCommand.operationFailed.rawValue) } + return + } + store.isHoldingPeripheral(address: address) { held in + self.queue.async { + self.sendString( + (held ? DeviceCommand.operationSuccess : DeviceCommand.operationFailed).rawValue) + } } } case .adoptReleased: @@ -474,6 +502,7 @@ final class IncomingConnection { private func teardown() { guard !finished else { return } finished = true + endPendingIfNeeded() idleTimer?.cancel() totalTimer?.cancel() idleTimer = nil @@ -483,6 +512,12 @@ final class IncomingConnection { release() } + private func endPendingIfNeeded() { + guard pendingCounted else { return } + pendingCounted = false + rateLimiter.endPending(endpoint: endpoint) + } + private func release() { queue.async { [weak self] in self?.selfRef = nil diff --git a/Magic Switch/Manager/RateLimiter.swift b/Magic Switch/Manager/RateLimiter.swift index 5883030..ab1234d 100644 --- a/Magic Switch/Manager/RateLimiter.swift +++ b/Magic Switch/Manager/RateLimiter.swift @@ -17,6 +17,8 @@ final class RateLimiter { private static let failureThreshold = 5 private static let blockDuration: TimeInterval = 15 * 60 private static let blocksKey = "com.magicswitch.ratelimiter.blocks" + private static let maxPendingPerIP = 8 + private static let maxPendingTotal = 64 // MARK: - State @@ -25,6 +27,9 @@ final class RateLimiter { private var failuresByIP: [String: [CFTimeInterval]] = [:] /// Monotonic deadline (seconds since boot) at which the block lifts. private var blocksByIP: [String: CFTimeInterval] = [:] + /// In-flight pre-auth connection counts, per IP and in total. + private var pendingByIP: [String: Int] = [:] + private var pendingTotal = 0 // MARK: - Init @@ -46,6 +51,30 @@ final class RateLimiter { } } + /// Reserve a pre-auth connection slot; false if the per-IP or global cap is hit. + func beginPending(endpoint: NWEndpoint?) -> Bool { + let key = Self.bucket(for: endpoint) + return queue.sync { + guard pendingTotal < Self.maxPendingTotal, + pendingByIP[key, default: 0] < Self.maxPendingPerIP + else { return false } + pendingByIP[key, default: 0] += 1 + pendingTotal += 1 + return true + } + } + + /// Release a slot reserved by `beginPending`. + func endPending(endpoint: NWEndpoint?) { + let key = Self.bucket(for: endpoint) + queue.sync { + if let count = pendingByIP[key] { + if count <= 1 { pendingByIP.removeValue(forKey: key) } else { pendingByIP[key] = count - 1 } + } + if pendingTotal > 0 { pendingTotal -= 1 } + } + } + /// Record an authentication failure for `endpoint`. func recordFailure(endpoint: NWEndpoint?) { let key = Self.bucket(for: endpoint) diff --git a/Magic Switch/Model/Store/BluetoothPeripheralStore.swift b/Magic Switch/Model/Store/BluetoothPeripheralStore.swift index 3e0b80f..186777f 100644 --- a/Magic Switch/Model/Store/BluetoothPeripheralStore.swift +++ b/Magic Switch/Model/Store/BluetoothPeripheralStore.swift @@ -1558,8 +1558,14 @@ final class BluetoothPeripheralStore: NSObject, ObservableObject, BluetoothPerip else { return } - print("Accepting Bluetooth pairing confirmation for \(address): \(numericValue)") - pair.replyUserConfirmation(true) + DispatchQueue.main.async { [weak self] in + guard let self = self else { return } + let initiated = self.pendingPairs[address] === pair + print( + "Bluetooth pairing confirmation for \(address) (\(numericValue)): " + + (initiated ? "accept" : "refuse")) + pair.replyUserConfirmation(initiated) + } } @objc func devicePairingPINCodeRequest(_ sender: Any!) {