From 53826d5145e5351ce55a80833e75dbd9c91cf211 Mon Sep 17 00:00:00 2001 From: islam <2553451+qalandarov@users.noreply.github.com> Date: Wed, 15 Oct 2025 03:59:45 +0100 Subject: [PATCH] Revive lost `NotificationStreamAssembler` changes --- .../NotificationStreamAssembler.swift | 108 +++++++++++++----- 1 file changed, 80 insertions(+), 28 deletions(-) diff --git a/bitchat/Services/NotificationStreamAssembler.swift b/bitchat/Services/NotificationStreamAssembler.swift index 77270fdb..39948dc6 100644 --- a/bitchat/Services/NotificationStreamAssembler.swift +++ b/bitchat/Services/NotificationStreamAssembler.swift @@ -6,10 +6,19 @@ // For more information, see // +import BitLogger import Foundation 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) } @@ -18,64 +27,107 @@ struct NotificationStreamAssembler { var frames: [Data] = [] var dropped: [UInt8] = [] - var reset = false - let maxFrameLength = TransportConfig.blePendingWriteBufferCapBytes + var didReset = false + let now = DispatchTime.now() + let maxFrameLength = TransportConfig.bleNotificationAssemblerHardCapBytes + let minimumFramePrefix = BinaryProtocol.v1HeaderSize + BinaryProtocol.senderIDSize - let minHeaderBytes = 14 // version + type + ttl + timestamp(8) + flags + length(2) - let minFramePrefix = minHeaderBytes + 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 >= minFramePrefix { - guard let first = buffer.first else { break } - if first != 1 { + 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 } - guard buffer.count >= minHeaderBytes else { break } + guard let headerSize = BinaryProtocol.headerSize(for: version) else { + dropped.append(buffer.removeFirst()) + pendingFrameStartedAt = nil + pendingFrameExpectedLength = 0 + continue + } + let framePrefix = headerSize + BinaryProtocol.senderIDSize + guard buffer.count >= framePrefix else { break } - let headerBytes = Array(buffer.prefix(minFramePrefix)) - guard headerBytes.count == minFramePrefix else { break } - - let flags = headerBytes[11] + let flagsIndex = buffer.startIndex + BinaryProtocol.Offsets.flags + guard flagsIndex < buffer.endIndex else { break } + let flags = buffer[flagsIndex] let hasRecipient = (flags & BinaryProtocol.Flags.hasRecipient) != 0 let hasSignature = (flags & BinaryProtocol.Flags.hasSignature) != 0 - let payloadLen = (Int(headerBytes[12]) << 8) | Int(headerBytes[13]) + let isCompressed = (flags & BinaryProtocol.Flags.isCompressed) != 0 - var frameLength = minFramePrefix + payloadLen + let lengthOffset = 12 + let payloadLength: Int + if version == 2 { + let lengthIndex = buffer.startIndex + lengthOffset + payloadLength = + (Int(buffer[lengthIndex]) << 24) | + (Int(buffer[lengthIndex + 1]) << 16) | + (Int(buffer[lengthIndex + 2]) << 8) | + Int(buffer[lengthIndex + 3]) + } else { + let lengthIndex = buffer.startIndex + lengthOffset + payloadLength = (Int(buffer[lengthIndex]) << 8) | Int(buffer[lengthIndex + 1]) + } + + var frameLength = framePrefix + payloadLength if hasRecipient { frameLength += BinaryProtocol.recipientIDSize } if hasSignature { frameLength += BinaryProtocol.signatureSize } + if isCompressed { + let rawLengthFieldBytes = (version == 2) ? 4 : 2 + if payloadLength < rawLengthFieldBytes { + SecureLogger.error("❌ Invalid compressed payload length (\(payloadLength))", category: .session) + resetState() + didReset = true + break + } + } guard frameLength > 0, frameLength <= maxFrameLength else { - buffer.removeAll() - reset = true + SecureLogger.error("❌ Notification frame length \(frameLength) invalid (cap=\(maxFrameLength)); resetting stream", category: .session) + resetState() + didReset = true break } if buffer.count < frameLength { - // Check if a new frame start exists within the incomplete buffer; if so, drop leading partial bytes. - if let nextStart = buffer.dropFirst().firstIndex(of: 1) { - let dropCount = buffer.distance(from: buffer.startIndex, to: nextStart) - if dropCount > 0 { - buffer.removeFirst(dropCount) - dropped.append(1) // treat as dropped partial start + 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) } }