From aa35200c6fde8623ab705236243854d1484b68b6 Mon Sep 17 00:00:00 2001
From: jack <jackjackbits@users.noreply.github.com>
Date: Fri, 26 Sep 2025 12:30:31 +0200
Subject: [PATCH] Reset BLE assembler on stalled fragment trains

---
 bitchat/Services/BLEService.swift      | 63 ++++++++++++++++++--------
 bitchat/Services/TransportConfig.swift |  1 +
 2 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/bitchat/Services/BLEService.swift b/bitchat/Services/BLEService.swift
index fdbec914..da250a00 100644
--- a/bitchat/Services/BLEService.swift
+++ b/bitchat/Services/BLEService.swift
@@ -9,29 +9,39 @@ import UIKit
 
 struct NotificationStreamAssembler {
     private var buffer = Data()
+    private var pendingFrameStartedAt: DispatchTime?
+    private var pendingFrameExpectedLength: Int = 0
+
+    private mutating func resetState() {
+        buffer.removeAll(keepingCapacity: false)
+        pendingFrameStartedAt = nil
+        pendingFrameExpectedLength = 0
+    }
 
     mutating func append(_ chunk: Data) -> (frames: [Data], droppedPrefixes: [UInt8], reset: Bool) {
         guard !chunk.isEmpty else { return ([], [], false) }
 
         buffer.append(chunk)
 
-        if buffer.count > TransportConfig.bleNotificationAssemblerHardCapBytes {
-            SecureLogger.error("❌ Notification assembler overflow (\(buffer.count) bytes); dropping partial frame", category: .session)
-            buffer.removeAll(keepingCapacity: false)
-            return ([], [], true)
-        }
-
         var frames: [Data] = []
         var dropped: [UInt8] = []
-        var reset = false
+        var didReset = false
+        let now = DispatchTime.now()
         let maxFrameLength = TransportConfig.bleNotificationAssemblerHardCapBytes
-
         let minimumFramePrefix = BinaryProtocol.v1HeaderSize + BinaryProtocol.senderIDSize
 
+        if buffer.count > TransportConfig.bleNotificationAssemblerHardCapBytes {
+            SecureLogger.error("❌ Notification assembler overflow (\(buffer.count) bytes); dropping partial frame", category: .session)
+            resetState()
+            return ([], [], true)
+        }
+
         while buffer.count >= minimumFramePrefix {
             guard let version = buffer.first else { break }
             guard version == 1 || version == 2 else {
                 dropped.append(buffer.removeFirst())
+                pendingFrameStartedAt = nil
+                pendingFrameExpectedLength = 0
                 continue
             }
 
@@ -39,6 +49,8 @@ struct NotificationStreamAssembler {
             let framePrefix = headerSize + BinaryProtocol.senderIDSize
             guard headerSize > 0 else {
                 dropped.append(buffer.removeFirst())
+                pendingFrameStartedAt = nil
+                pendingFrameExpectedLength = 0
                 continue
             }
             guard buffer.count >= framePrefix else { break }
@@ -72,38 +84,51 @@ struct NotificationStreamAssembler {
                 let rawLengthFieldBytes = (version == 2) ? 4 : 2
                 if payloadLength < rawLengthFieldBytes {
                     SecureLogger.error("❌ Invalid compressed payload length (\(payloadLength))", category: .session)
-                    buffer.removeAll()
-                    reset = true
+                    resetState()
+                    didReset = true
                     break
                 }
             }
 
             guard frameLength > 0, frameLength <= maxFrameLength else {
                 SecureLogger.error("❌ Notification frame length \(frameLength) invalid (cap=\(maxFrameLength)); resetting stream", category: .session)
-                buffer.removeAll()
-                reset = true
+                resetState()
+                didReset = true
                 break
             }
 
             if buffer.count < frameLength {
-                SecureLogger.debug("⌛ Waiting for remaining \(frameLength - buffer.count)B to complete BLE frame", category: .session)
+                let remaining = frameLength - buffer.count
+                if pendingFrameStartedAt == nil || frameLength != pendingFrameExpectedLength {
+                    pendingFrameStartedAt = now
+                    pendingFrameExpectedLength = frameLength
+                } else if let started = pendingFrameStartedAt {
+                    let elapsed = now.uptimeNanoseconds - started.uptimeNanoseconds
+                    let threshold = UInt64(TransportConfig.bleAssemblerStallResetMs) * 1_000_000
+                    if elapsed >= threshold {
+                        SecureLogger.debug("📉 Resetting notification assembler after waiting \(remaining)B for \(TransportConfig.bleAssemblerStallResetMs)ms", category: .session)
+                        resetState()
+                        didReset = true
+                    } else {
+                        SecureLogger.debug("⌛ Waiting for remaining \(remaining)B to complete BLE frame", category: .session)
+                    }
+                }
                 break
             }
 
+            pendingFrameStartedAt = nil
+            pendingFrameExpectedLength = 0
+
             let frame = Data(buffer.prefix(frameLength))
             frames.append(frame)
             buffer.removeFirst(frameLength)
         }
 
         if !buffer.isEmpty, buffer.allSatisfy({ $0 == 0 }) {
-            buffer.removeAll(keepingCapacity: false)
+            resetState()
         }
 
-        return (frames, dropped, reset)
-    }
-
-    mutating func reset() {
-        buffer.removeAll(keepingCapacity: false)
+        return (frames, dropped, didReset)
     }
 }
 
diff --git a/bitchat/Services/TransportConfig.swift b/bitchat/Services/TransportConfig.swift
index 8b30af92..5702fd55 100644
--- a/bitchat/Services/TransportConfig.swift
+++ b/bitchat/Services/TransportConfig.swift
@@ -31,6 +31,7 @@ enum TransportConfig {
     static let bleConnectionCandidatesMax: Int = 100
     static let blePendingWriteBufferCapBytes: Int = 1_000_000
     static let bleNotificationAssemblerHardCapBytes: Int = 8 * 1024 * 1024
+    static let bleAssemblerStallResetMs: Int = 250
     static let blePendingNotificationsCapCount: Int = 128
     static let bleNotificationRetryDelayMs: Int = 25
     static let bleNotificationRetryMaxAttempts: Int = 80