diff --git a/bitchat/Protocols/BinaryProtocol.swift b/bitchat/Protocols/BinaryProtocol.swift index 5574946b..78c80313 100644 --- a/bitchat/Protocols/BinaryProtocol.swift +++ b/bitchat/Protocols/BinaryProtocol.swift @@ -320,8 +320,8 @@ struct BinaryProtocol { guard let rawSize = read16() else { return nil } originalSize = Int(rawSize) } - // Guard to keep decompression bounded (1 MiB ceiling aligns with previous behaviour) - guard originalSize >= 0 && originalSize <= 1_048_576 else { return nil } + // Guard to keep decompression bounded to sane BLE payload limits + guard originalSize >= 0 && originalSize <= FileTransferLimits.maxPayloadBytes else { return nil } let compressedSize = payloadLength - lengthFieldBytes guard compressedSize >= 0, let compressed = readData(compressedSize) else { return nil } guard let decompressed = CompressionUtil.decompress(compressed, originalSize: originalSize), diff --git a/bitchat/Services/BLEService.swift b/bitchat/Services/BLEService.swift index 9910e5c6..ad5eb621 100644 --- a/bitchat/Services/BLEService.swift +++ b/bitchat/Services/BLEService.swift @@ -26,7 +26,9 @@ struct NotificationStreamAssembler { var reset = false let maxFrameLength = TransportConfig.bleNotificationAssemblerHardCapBytes - while buffer.count >= BinaryProtocol.v1HeaderSize + BinaryProtocol.senderIDSize { + let minimumFramePrefix = BinaryProtocol.v1HeaderSize + BinaryProtocol.senderIDSize + + while buffer.count >= minimumFramePrefix { guard let version = buffer.first else { break } guard version == 1 || version == 2 else { dropped.append(buffer.removeFirst()) @@ -41,10 +43,13 @@ struct NotificationStreamAssembler { } guard buffer.count >= framePrefix else { break } - let flagsIndex = buffer.startIndex + 11 + let flagsOffset = 1 + 1 + 1 + 8 // version + type + ttl + timestamp + let flagsIndex = buffer.startIndex + flagsOffset + 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 isCompressed = (flags & BinaryProtocol.Flags.isCompressed) != 0 let lengthOffset = 12 let payloadLength: Int @@ -63,6 +68,15 @@ struct NotificationStreamAssembler { 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) + buffer.removeAll() + reset = true + break + } + } guard frameLength > 0, frameLength <= maxFrameLength else { SecureLogger.error("❌ Notification frame length \(frameLength) invalid (cap=\(maxFrameLength)); resetting stream", category: .session) diff --git a/bitchatTests/NotificationStreamAssemblerTests.swift b/bitchatTests/NotificationStreamAssemblerTests.swift index bbed45ef..07d3b00c 100644 --- a/bitchatTests/NotificationStreamAssemblerTests.swift +++ b/bitchatTests/NotificationStreamAssemblerTests.swift @@ -95,4 +95,67 @@ struct NotificationStreamAssemblerTests { let decoded = try #require(BinaryProtocol.decode(result.frames[0]), "Failed to decode frame after drop") #expect(decoded.timestamp == packet.timestamp) } + + func testAssemblesCompressedLargeFrame() throws { + var assembler = NotificationStreamAssembler() + + let largeContent = Data(repeating: 0x41, count: 2_500_000) + let filePacket = BitchatFilePacket( + fileName: "large.bin", + fileSize: UInt64(largeContent.count), + mimeType: "application/octet-stream", + content: largeContent + ) + guard let tlvPayload = filePacket.encode() else { + return XCTFail("Failed to encode file packet") + } + + let senderID = Data(repeating: 0xAA, count: BinaryProtocol.senderIDSize) + let packet = BitchatPacket( + type: MessageType.fileTransfer.rawValue, + senderID: senderID, + recipientID: nil, + timestamp: 0x010203040506, + payload: tlvPayload, + signature: nil, + ttl: 3, + version: 2 + ) + + guard let frame = packet.toBinaryData(padding: false) else { + return XCTFail("Failed to encode packet frame") + } + + let flagsOffset = 1 + 1 + 1 + 8 + XCTAssertLessThan(flagsOffset, frame.count) + let flags = frame[frame.startIndex + flagsOffset] + XCTAssertNotEqual(flags & BinaryProtocol.Flags.isCompressed, 0, "Frame should be compressed for large payloads") + + let splitIndex = min(4096, frame.count / 2) + var result = assembler.append(frame.prefix(splitIndex)) + XCTAssertTrue(result.frames.isEmpty) + + result = assembler.append(frame.suffix(from: splitIndex)) + XCTAssertEqual(result.frames.count, 1) + XCTAssertTrue(result.droppedPrefixes.isEmpty) + XCTAssertFalse(result.reset) + + guard let assembled = result.frames.first else { + return XCTFail("Missing assembled frame") + } + XCTAssertEqual(assembled.count, frame.count) + + guard let decodedPacket = BinaryProtocol.decode(assembled) else { + return XCTFail("Failed to decode compressed frame") + } + XCTAssertEqual(decodedPacket.payload.count, tlvPayload.count) + + guard let decodedFile = BitchatFilePacket.decode(decodedPacket.payload) else { + return XCTFail("Failed to decode TLV payload") + } + XCTAssertEqual(decodedFile.fileName, filePacket.fileName) + XCTAssertEqual(decodedFile.mimeType, filePacket.mimeType) + XCTAssertEqual(decodedFile.content.count, largeContent.count) + XCTAssertEqual(decodedFile.content.prefix(32), largeContent.prefix(32)) + } } diff --git a/bitchatTests/Protocols/BitchatFilePacketTests.swift b/bitchatTests/Protocols/BitchatFilePacketTests.swift new file mode 100644 index 00000000..51801fe5 --- /dev/null +++ b/bitchatTests/Protocols/BitchatFilePacketTests.swift @@ -0,0 +1,47 @@ +import XCTest +@testable import bitchat + +final class BitchatFilePacketTests: XCTestCase { + + func testRoundTripPreservesFields() throws { + let content = Data((0..<4096).map { UInt8($0 % 251) }) + let packet = BitchatFilePacket( + fileName: "sample.jpg", + fileSize: UInt64(content.count), + mimeType: "image/jpeg", + content: content + ) + + guard let encoded = packet.encode() else { + return XCTFail("Failed to encode file packet") + } + guard let decoded = BitchatFilePacket.decode(encoded) else { + return XCTFail("Failed to decode file packet") + } + + XCTAssertEqual(decoded.fileName, packet.fileName) + XCTAssertEqual(decoded.fileSize, packet.fileSize) + XCTAssertEqual(decoded.mimeType, packet.mimeType) + XCTAssertEqual(decoded.content, packet.content) + } + + func testDecodeFallsBackToContentSizeWhenFileSizeMissing() throws { + let content = Data(repeating: 0x7F, count: 1024) + let packet = BitchatFilePacket( + fileName: nil, + fileSize: nil, + mimeType: nil, + content: content + ) + + guard let encoded = packet.encode() else { + return XCTFail("Failed to encode file packet") + } + guard let decoded = BitchatFilePacket.decode(encoded) else { + return XCTFail("Failed to decode file packet") + } + + XCTAssertEqual(decoded.fileSize, UInt64(content.count)) + XCTAssertEqual(decoded.content, content) + } +}