From 92ae764f91a007021bb1c65b3d84d5cd826de7ad Mon Sep 17 00:00:00 2001
From: jack <212554440+jackjackbits@users.noreply.github.com>
Date: Tue, 2 Jun 2026 22:19:18 +0200
Subject: [PATCH] Fix unbounded public rate limit buckets

---
 bitchat/Services/TransportConfig.swift      |  3 +
 bitchat/ViewModels/MessageRateLimiter.swift | 70 ++++++++++++++++++---
 bitchatTests/MessageRateLimiterTests.swift  | 56 +++++++++++++++++
 3 files changed, 121 insertions(+), 8 deletions(-)
 create mode 100644 bitchatTests/MessageRateLimiterTests.swift

diff --git a/bitchat/Services/TransportConfig.swift b/bitchat/Services/TransportConfig.swift
index 6153c266..a6a56e43 100644
--- a/bitchat/Services/TransportConfig.swift
+++ b/bitchat/Services/TransportConfig.swift
@@ -56,6 +56,9 @@ enum TransportConfig {
     static let uiSenderRateBucketRefillPerSec: Double = 1.0
     static let uiContentRateBucketCapacity: Double = 3
     static let uiContentRateBucketRefillPerSec: Double = 0.5
+    static let uiSenderRateBucketMaxEntries: Int = 2000
+    static let uiContentRateBucketMaxEntries: Int = 2000
+    static let uiRateBucketIdleTTL: TimeInterval = 10 * 60
 
     // UI sleeps/delays
     static let uiStartupInitialDelaySeconds: TimeInterval = 1.0
diff --git a/bitchat/ViewModels/MessageRateLimiter.swift b/bitchat/ViewModels/MessageRateLimiter.swift
index 5bca9a6d..fd0b03c9 100644
--- a/bitchat/ViewModels/MessageRateLimiter.swift
+++ b/bitchat/ViewModels/MessageRateLimiter.swift
@@ -26,6 +26,10 @@ struct MessageRateLimiter {
             }
             return false
         }
+
+        func isIdle(since now: Date, idleTTL: TimeInterval) -> Bool {
+            now.timeIntervalSince(lastRefill) >= idleTTL
+        }
     }
 
     private var senderBuckets: [String: TokenBucket] = [:]
@@ -35,43 +39,93 @@ struct MessageRateLimiter {
     private let senderRefill: Double
     private let contentCapacity: Double
     private let contentRefill: Double
+    private let maxSenderBuckets: Int
+    private let maxContentBuckets: Int
+    private let bucketIdleTTL: TimeInterval
 
     init(
         senderCapacity: Double,
         senderRefillPerSec: Double,
         contentCapacity: Double,
-        contentRefillPerSec: Double
+        contentRefillPerSec: Double,
+        maxSenderBuckets: Int = TransportConfig.uiSenderRateBucketMaxEntries,
+        maxContentBuckets: Int = TransportConfig.uiContentRateBucketMaxEntries,
+        bucketIdleTTL: TimeInterval = TransportConfig.uiRateBucketIdleTTL
     ) {
         self.senderCapacity = senderCapacity
         self.senderRefill = senderRefillPerSec
         self.contentCapacity = contentCapacity
         self.contentRefill = contentRefillPerSec
+        self.maxSenderBuckets = max(1, maxSenderBuckets)
+        self.maxContentBuckets = max(1, maxContentBuckets)
+        self.bucketIdleTTL = bucketIdleTTL
     }
 
     mutating func allow(senderKey: String, contentKey: String, now: Date = Date()) -> Bool {
-        var senderBucket = senderBuckets[senderKey] ?? TokenBucket(
+        var senderBucket = bucket(
+            for: senderKey,
+            in: &senderBuckets,
             capacity: senderCapacity,
-            tokens: senderCapacity,
             refillPerSec: senderRefill,
-            lastRefill: now
+            maxBuckets: maxSenderBuckets,
+            now: now
         )
         let senderAllowed = senderBucket.allow(now: now)
         senderBuckets[senderKey] = senderBucket
+        guard senderAllowed else { return false }
 
-        var contentBucket = contentBuckets[contentKey] ?? TokenBucket(
+        var contentBucket = bucket(
+            for: contentKey,
+            in: &contentBuckets,
             capacity: contentCapacity,
-            tokens: contentCapacity,
             refillPerSec: contentRefill,
-            lastRefill: now
+            maxBuckets: maxContentBuckets,
+            now: now
         )
         let contentAllowed = contentBucket.allow(now: now)
         contentBuckets[contentKey] = contentBucket
 
-        return senderAllowed && contentAllowed
+        return contentAllowed
     }
 
     mutating func reset() {
         senderBuckets.removeAll()
         contentBuckets.removeAll()
     }
+
+    var bucketCountsForTesting: (sender: Int, content: Int) {
+        (senderBuckets.count, contentBuckets.count)
+    }
+
+    private mutating func bucket(
+        for key: String,
+        in buckets: inout [String: TokenBucket],
+        capacity: Double,
+        refillPerSec: Double,
+        maxBuckets: Int,
+        now: Date
+    ) -> TokenBucket {
+        if let bucket = buckets[key] {
+            return bucket
+        }
+
+        evictIfNeeded(from: &buckets, maxBuckets: maxBuckets, now: now)
+        return TokenBucket(
+            capacity: capacity,
+            tokens: capacity,
+            refillPerSec: refillPerSec,
+            lastRefill: now
+        )
+    }
+
+    private func evictIfNeeded(from buckets: inout [String: TokenBucket], maxBuckets: Int, now: Date) {
+        guard buckets.count >= maxBuckets else { return }
+
+        buckets = buckets.filter { !$0.value.isIdle(since: now, idleTTL: bucketIdleTTL) }
+        guard buckets.count >= maxBuckets else { return }
+
+        if let oldestKey = buckets.min(by: { $0.value.lastRefill < $1.value.lastRefill })?.key {
+            buckets.removeValue(forKey: oldestKey)
+        }
+    }
 }
diff --git a/bitchatTests/MessageRateLimiterTests.swift b/bitchatTests/MessageRateLimiterTests.swift
new file mode 100644
index 00000000..4219ece3
--- /dev/null
+++ b/bitchatTests/MessageRateLimiterTests.swift
@@ -0,0 +1,56 @@
+//
+// MessageRateLimiterTests.swift
+// bitchatTests
+//
+// Ensures public-message rate limiter state remains bounded for attacker-derived keys.
+//
+
+import Foundation
+import Testing
+@testable import bitchat
+
+struct MessageRateLimiterTests {
+    @Test("Content buckets do not grow when sender is rate limited")
+    func contentBucketsDoNotGrowAfterSenderLimit() {
+        var limiter = MessageRateLimiter(
+            senderCapacity: 1,
+            senderRefillPerSec: 0,
+            contentCapacity: 1,
+            contentRefillPerSec: 0,
+            maxSenderBuckets: 10,
+            maxContentBuckets: 10,
+            bucketIdleTTL: 60
+        )
+        let now = Date()
+
+        #expect(limiter.allow(senderKey: "sender", contentKey: "content-0", now: now))
+        for index in 1...100 {
+            #expect(!limiter.allow(senderKey: "sender", contentKey: "content-\(index)", now: now))
+        }
+
+        #expect(limiter.bucketCountsForTesting.sender == 1)
+        #expect(limiter.bucketCountsForTesting.content == 1)
+    }
+
+    @Test("Bucket maps evict entries at configured caps")
+    func bucketMapsEvictAtConfiguredCaps() {
+        let maxEntries = 3
+        var limiter = MessageRateLimiter(
+            senderCapacity: 1,
+            senderRefillPerSec: 0,
+            contentCapacity: 1,
+            contentRefillPerSec: 0,
+            maxSenderBuckets: maxEntries,
+            maxContentBuckets: maxEntries,
+            bucketIdleTTL: 60
+        )
+        let now = Date()
+
+        for index in 0..<25 {
+            #expect(limiter.allow(senderKey: "sender-\(index)", contentKey: "content-\(index)", now: now.addingTimeInterval(TimeInterval(index))))
+        }
+
+        #expect(limiter.bucketCountsForTesting.sender == maxEntries)
+        #expect(limiter.bucketCountsForTesting.content == maxEntries)
+    }
+}