fix: stop NoiseRateLimiter retaining every peer it has ever seen

bitchat/Noise/NoiseRateLimiter.swift

@@ -13,16 +13,43 @@ import Foundation
 final class NoiseRateLimiter {
     private var handshakeTimestamps: [PeerID: [Date]] = [:]
     private var messageTimestamps: [PeerID: [Date]] = [:]
+    private var lastPrune: Date = .distantPast
 
     // Global rate limiting
     private var globalHandshakeTimestamps: [Date] = []
     private var globalMessageTimestamps: [Date] = []
 
     private let queue = DispatchQueue(label: "chat.bitchat.noise.ratelimit", attributes: .concurrent)
+
+    /// Clock seam. Every sibling limiter takes `now` as a parameter
+    /// (`SyncResponseRateLimiter`, `BLESubscriptionAnnounceLimiter`,
+    /// `BLEAnnounceThrottle`); this one read `Date()` inline, which is why its
+    /// time-dependent behaviour had no coverage. Injected here instead of
+    /// threading a parameter through nine call sites.
+    private let currentDate: () -> Date
+
+    /// Sweeping every peer on every admission would put an O(peers) walk on the
+    /// message path, which runs up to `maxGlobalMessagesPerSecond` times a
+    /// second. Once a second is frequent enough to keep both maps to peers seen
+    /// inside their own windows.
+    private static let pruneInterval: TimeInterval = 1
+
+    init(currentDate: @escaping () -> Date = Date.init) {
+        self.currentDate = currentDate
+    }
+
+    /// Peers currently retained in either map. Mirrors
+    /// `BLESubscriptionAnnounceLimiter.trackedCentralCount`.
+    var trackedPeerCount: Int {
+        queue.sync {
+            Set(handshakeTimestamps.keys).union(messageTimestamps.keys).count
+        }
+    }
 
     func allowHandshake(from peerID: PeerID) -> Bool {
         return queue.sync(flags: .barrier) {
-            let now = Date()
+            let now = currentDate()
+            pruneStalePeersLocked(now: now)
             let oneMinuteAgo = now.addingTimeInterval(-60)
 
             // Check global rate limit first
@@ -51,7 +78,8 @@ final class NoiseRateLimiter {
 
     func allowMessage(from peerID: PeerID) -> Bool {
         return queue.sync(flags: .barrier) {
-            let now = Date()
+            let now = currentDate()
+            pruneStalePeersLocked(now: now)
             let oneSecondAgo = now.addingTimeInterval(-1)
 
             // Check global rate limit first
@@ -78,6 +106,30 @@ final class NoiseRateLimiter {
         }
     }
 
+    /// Drops peers whose timestamps have all aged out of their window.
+    ///
+    /// Without this the two maps only ever grew: a peer's array was filtered
+    /// when that same peer was next queried, but a peer that never came back
+    /// kept its entry for the lifetime of the process. `reset(for:)` below is
+    /// the per-peer counterpart and is never called from production code, so
+    /// nothing else reclaimed them. Must be called with the barrier held.
+    private func pruneStalePeersLocked(now: Date) {
+        guard now.timeIntervalSince(lastPrune) >= Self.pruneInterval else { return }
+        lastPrune = now
+
+        let handshakeCutoff = now.addingTimeInterval(-60)
+        handshakeTimestamps = handshakeTimestamps.compactMapValues { timestamps in
+            let recent = timestamps.filter { $0 > handshakeCutoff }
+            return recent.isEmpty ? nil : recent
+        }
+
+        let messageCutoff = now.addingTimeInterval(-1)
+        messageTimestamps = messageTimestamps.compactMapValues { timestamps in
+            let recent = timestamps.filter { $0 > messageCutoff }
+            return recent.isEmpty ? nil : recent
+        }
+    }
+
     func reset(for peerID: PeerID) {
         queue.async(flags: .barrier) {
             self.handshakeTimestamps.removeValue(forKey: peerID)

bitchatTests/Noise/NoiseRateLimiterTests.swift

@@ -67,4 +67,79 @@ final class NoiseRateLimiterTests: XCTestCase {
     private func makePeerID(_ value: Int) -> PeerID {
         PeerID(str: String(format: "%016x", value))
     }
+
+    // MARK: - Peer-map retention
+
+    /// Drives the limiter's clock so window expiry is observable.
+    private final class TestClock: @unchecked Sendable {
+        private let lock = NSLock()
+        private var _now: Date
+        init(_ start: Date) { _now = start }
+        var now: Date { lock.withLock { _now } }
+        func advance(_ seconds: TimeInterval) { lock.withLock { _now += seconds } }
+    }
+
+    func test_handshakeMap_doesNotRetainPeersPastTheirWindow() {
+        let clock = TestClock(Date(timeIntervalSince1970: 10_000))
+        let limiter = NoiseRateLimiter(currentDate: { clock.now })
+
+        // A burst of one-shot peers, each handshaking once and never returning.
+        // Stay inside the global per-minute budget so every one is admitted.
+        let peerCount = min(20, NoiseSecurityConstants.maxGlobalHandshakesPerMinute - 1)
+        for index in 0..<peerCount {
+            XCTAssertTrue(limiter.allowHandshake(from: makePeerID(index + 1)))
+        }
+        XCTAssertEqual(limiter.trackedPeerCount, peerCount)
+
+        // Past the one-minute handshake window, none of them is still relevant.
+        clock.advance(61)
+        _ = limiter.allowHandshake(from: makePeerID(200))
+
+        XCTAssertEqual(
+            limiter.trackedPeerCount,
+            1,
+            "departed peers must not be retained once their handshake window has passed"
+        )
+    }
+
+    func test_messageMap_doesNotRetainPeersPastTheirWindow() {
+        let clock = TestClock(Date(timeIntervalSince1970: 20_000))
+        let limiter = NoiseRateLimiter(currentDate: { clock.now })
+
+        let peerCount = min(10, NoiseSecurityConstants.maxGlobalMessagesPerSecond - 1)
+        for index in 0..<peerCount {
+            XCTAssertTrue(limiter.allowMessage(from: makePeerID(index + 1)))
+        }
+        XCTAssertEqual(limiter.trackedPeerCount, peerCount)
+
+        clock.advance(2)
+        _ = limiter.allowMessage(from: makePeerID(200))
+
+        XCTAssertEqual(
+            limiter.trackedPeerCount,
+            1,
+            "departed peers must not be retained once their message window has passed"
+        )
+    }
+
+    func test_pruningKeepsAPeerStillInsideItsWindow() {
+        // The counterpart: pruning must not discard a peer whose budget is still
+        // being enforced, or the limit becomes trivially bypassable by waiting.
+        let clock = TestClock(Date(timeIntervalSince1970: 30_000))
+        let limiter = NoiseRateLimiter(currentDate: { clock.now })
+        let peerID = makePeerID(7)
+
+        for _ in 0..<NoiseSecurityConstants.maxHandshakesPerMinute {
+            XCTAssertTrue(limiter.allowHandshake(from: peerID))
+        }
+        XCTAssertFalse(limiter.allowHandshake(from: peerID))
+
+        // Well past the prune interval, but well inside the one-minute window.
+        clock.advance(5)
+        XCTAssertFalse(
+            limiter.allowHandshake(from: peerID),
+            "a peer still inside its window must stay rate limited across a prune"
+        )
+        XCTAssertEqual(limiter.trackedPeerCount, 1)
+    }
 }