bitchat/bitchatTests/Fragmentation/FragmentationTests.swift

//
// FragmentationTests.swift
// bitchatTests
//
// This is free and unencumbered software released into the public domain.
// For more information, see <https://unlicense.org>
//

import Testing
import Foundation
import CoreBluetooth
import BitFoundation
@testable import bitchat

@Suite("Fragmentation Tests", .serialized)
struct FragmentationTests {

    @Test("Reassembly from fragments delivers a public message")
    func reassemblyFromFragmentsDeliversPublicMessage() async throws {
        let ble = makeBLEService()
        let capture = CaptureDelegate()
        ble.delegate = capture

        // Construct a big SIGNED public packet (3KB) from a remote sender. Public
        // messages must carry a valid signature, so the reassembled packet is
        // signed and the sender's signing key is preseeded into the registry.
        let signer = NoiseEncryptionService(keychain: MockKeychain())
        let signingKey = signer.getSigningPublicKeyData()
        let remoteShortID = PeerID(str: "1122334455667788")
        let original = try #require(
            signer.signPacket(makeLargePublicPacket(senderShortHex: remoteShortID, size: 3_000)),
            "Failed to sign public packet"
        )

        // Use a small fragment size to ensure multiple pieces
        let fragments = fragmentPacket(original, fragmentSize: 400)

        // Shuffle fragments to simulate out-of-order arrival
        let shuffled = fragments.shuffled()

        // Send fragments sequentially with small delays (no fire-and-forget Tasks)
        for (i, fragment) in shuffled.enumerated() {
            if i > 0 {
                try await Task.sleep(for: .milliseconds(5))
            }
            ble._test_handlePacket(fragment, fromPeerID: remoteShortID, signingPublicKey: signingKey)
        }

        // Wait for delegate callback with proper timeout
        try await capture.waitForPublicMessages(count: 1, timeout: .seconds(5))

        #expect(capture.publicMessages.count == 1)
        #expect(capture.publicMessages.first?.content.count == 3_000)
    }
    
    @Test("Duplicate fragment does not break reassembly")
    func duplicateFragmentDoesNotBreakReassembly() async throws {
        let ble = makeBLEService()
        let capture = CaptureDelegate()
        ble.delegate = capture

        let signer = NoiseEncryptionService(keychain: MockKeychain())
        let signingKey = signer.getSigningPublicKeyData()
        let remoteShortID = PeerID(str: "A1B2C3D4E5F60708")
        let original = try #require(
            signer.signPacket(makeLargePublicPacket(senderShortHex: remoteShortID, size: 2048)),
            "Failed to sign public packet"
        )
        var frags = fragmentPacket(original, fragmentSize: 300)

        // Duplicate one fragment
        if let dup = frags.first {
            frags.insert(dup, at: 1)
        }

        // Send fragments sequentially with small delays (no fire-and-forget Tasks)
        for (i, fragment) in frags.enumerated() {
            if i > 0 {
                try await Task.sleep(for: .milliseconds(5))
            }
            ble._test_handlePacket(fragment, fromPeerID: remoteShortID, signingPublicKey: signingKey)
        }

        // Wait for delegate callback with proper timeout
        try await capture.waitForPublicMessages(count: 1, timeout: .seconds(5))

        #expect(capture.publicMessages.count == 1)
        #expect(capture.publicMessages.first?.content.count == 2048)
    }

    @Test("Max-sized file transfer survives reassembly")
    func maxSizedFileTransferSurvivesReassembly() async throws {
        let ble = makeBLEService()
        let capture = CaptureDelegate()
        ble.delegate = capture

        let remoteID = PeerID(str: "CAFEBABECAFEBABE")
        let fileContent = Data(repeating: 0x42, count: FileTransferLimits.maxPayloadBytes)
        let filePacket = BitchatFilePacket(
            fileName: "limit.bin",
            fileSize: UInt64(fileContent.count),
            mimeType: "application/octet-stream",
            content: fileContent
        )
        let encoded = try #require(filePacket.encode(), "File packet encoding failed")

        let packet = BitchatPacket(
            type: MessageType.fileTransfer.rawValue,
            senderID: Data(hexString: remoteID.id) ?? Data(),
            recipientID: nil,
            timestamp: UInt64(Date().timeIntervalSince1970 * 1000),
            payload: encoded,
            signature: nil,
            ttl: 7,
            version: 2
        )

        let fragments = fragmentPacket(packet, fragmentSize: 4096, pad: false)
        #expect(!fragments.isEmpty)

        for (i, fragment) in fragments.enumerated() {
            if i > 0 {
                try await Task.sleep(for: .milliseconds(5))
            }
            ble._test_handlePacket(fragment, fromPeerID: remoteID)
        }

        try await capture.waitForReceivedMessages(count: 1, timeout: .seconds(5))

        let message = try #require(capture.receivedMessages.first, "Expected file transfer message")
        #expect(message.content.hasPrefix("[file]"))

        if let fileName = message.content.split(separator: " ").last {
            let base = try FileManager.default.url(for: .applicationSupportDirectory, in: .userDomainMask, appropriateFor: nil, create: true)
            let filesRoot = base.appendingPathComponent("files", isDirectory: true)
            let incoming = filesRoot.appendingPathComponent("files/incoming", isDirectory: true)
            let url = incoming.appendingPathComponent(String(fileName))
            try? FileManager.default.removeItem(at: url)
        }
    }
    
    @Test("Invalid fragment header is ignored")
    func invalidFragmentHeaderIsIgnored() async throws {
        let ble = makeBLEService()
        let capture = CaptureDelegate()
        ble.delegate = capture
        
        let remoteShortID = PeerID(str: "0011223344556677")
        let original = makeLargePublicPacket(senderShortHex: remoteShortID, size: 1000)
        let fragments = fragmentPacket(original, fragmentSize: 250)
        
        // Corrupt one fragment: make payload too short (header incomplete)
        var corrupted = fragments
        if !corrupted.isEmpty {
            var p = corrupted[0]
            p = BitchatPacket(
                type: p.type,
                senderID: p.senderID,
                recipientID: p.recipientID,
                timestamp: p.timestamp,
                payload: Data([0x00, 0x01, 0x02]), // invalid header
                signature: nil,
                ttl: p.ttl
            )
            corrupted[0] = p
        }
        
        for (i, fragment) in corrupted.enumerated() {
            if i > 0 {
                try await Task.sleep(for: .milliseconds(5))
            }
            ble._test_handlePacket(fragment, fromPeerID: remoteShortID)
        }
        
        // Allow async processing
        try await sleep(0.5)

        // Should not deliver since one fragment is invalid and reassembly can't complete
        #expect(capture.publicMessages.isEmpty)
    }
}

extension FragmentationTests {
    private func makeBLEService() -> BLEService {
        let mockKeychain = MockKeychain()
        let mockIdentityManager = MockIdentityManager(mockKeychain)
        let idBridge = NostrIdentityBridge(keychain: MockKeychainHelper())

        return BLEService(
            keychain: mockKeychain,
            idBridge: idBridge,
            identityManager: mockIdentityManager,
            initializeBluetoothManagers: false
        )
    }

    /// Thread-safe delegate that supports awaiting message delivery
    private final class CaptureDelegate: BitchatDelegate, @unchecked Sendable {
        private let lock = NSLock()
        private var _publicMessages: [(peerID: PeerID, nickname: String, content: String)] = []
        private var _receivedMessages: [BitchatMessage] = []
        private var publicMessageContinuation: CheckedContinuation<Void, Never>?
        private var receivedMessageContinuation: CheckedContinuation<Void, Never>?
        private var expectedPublicMessageCount: Int = 0
        private var expectedReceivedMessageCount: Int = 0

        private func withLock<T>(_ body: () -> T) -> T {
            lock.lock()
            defer { lock.unlock() }
            return body()
        }

        var publicMessages: [(peerID: PeerID, nickname: String, content: String)] {
            withLock { _publicMessages }
        }

        var receivedMessages: [BitchatMessage] {
            withLock { _receivedMessages }
        }

        func didReceiveMessage(_ message: BitchatMessage) {
            lock.lock()
            _receivedMessages.append(message)
            let count = _receivedMessages.count
            let expected = expectedReceivedMessageCount
            let continuation = receivedMessageContinuation
            lock.unlock()

            if count >= expected, let cont = continuation {
                lock.lock()
                receivedMessageContinuation = nil
                lock.unlock()
                cont.resume()
            }
        }

        func didReceivePublicMessage(from peerID: PeerID, nickname: String, content: String, timestamp: Date, messageID: String?) {
            lock.lock()
            _publicMessages.append((peerID, nickname, content))
            let count = _publicMessages.count
            let expected = expectedPublicMessageCount
            let continuation = publicMessageContinuation
            lock.unlock()

            if count >= expected, let cont = continuation {
                lock.lock()
                publicMessageContinuation = nil
                lock.unlock()
                cont.resume()
            }
        }

        /// Waits for the specified number of public messages to be received
        func waitForPublicMessages(count: Int, timeout: Duration = .seconds(2)) async throws {
            let isAlreadySatisfied = withLock { () -> Bool in
                if _publicMessages.count >= count {
                    return true
                }
                expectedPublicMessageCount = count
                return false
            }
            if isAlreadySatisfied {
                return
            }

            try await withThrowingTaskGroup(of: Void.self) { group in
                group.addTask {
                    await withCheckedContinuation { continuation in
                        let shouldResumeImmediately = self.withLock {
                            // Recheck count after acquiring lock to avoid race condition
                            // where message arrives between initial check and continuation install
                            if self._publicMessages.count >= count {
                                return true
                            }
                            self.publicMessageContinuation = continuation
                            return false
                        }
                        if shouldResumeImmediately {
                            continuation.resume()
                        }
                    }
                }
                group.addTask {
                    try await Task.sleep(for: timeout)
                    throw CancellationError()
                }
                try await group.next()
                group.cancelAll()
            }
        }

        /// Waits for the specified number of received messages
        func waitForReceivedMessages(count: Int, timeout: Duration = .seconds(2)) async throws {
            let isAlreadySatisfied = withLock { () -> Bool in
                if _receivedMessages.count >= count {
                    return true
                }
                expectedReceivedMessageCount = count
                return false
            }
            if isAlreadySatisfied {
                return
            }

            try await withThrowingTaskGroup(of: Void.self) { group in
                group.addTask {
                    await withCheckedContinuation { continuation in
                        let shouldResumeImmediately = self.withLock {
                            // Recheck count after acquiring lock to avoid race condition
                            // where message arrives between initial check and continuation install
                            if self._receivedMessages.count >= count {
                                return true
                            }
                            self.receivedMessageContinuation = continuation
                            return false
                        }
                        if shouldResumeImmediately {
                            continuation.resume()
                        }
                    }
                }
                group.addTask {
                    try await Task.sleep(for: timeout)
                    throw CancellationError()
                }
                try await group.next()
                group.cancelAll()
            }
        }

        func didConnectToPeer(_ peerID: PeerID) {}
        func didDisconnectFromPeer(_ peerID: PeerID) {}
        func didUpdatePeerList(_ peers: [PeerID]) {}
        func isFavorite(fingerprint: String) -> Bool { false }
        func didUpdateMessageDeliveryStatus(_ messageID: String, status: DeliveryStatus) {}
        func didReceiveNoisePayload(from peerID: PeerID, type: NoisePayloadType, payload: Data, timestamp: Date) {}
        func didUpdateBluetoothState(_ state: CBManagerState) {}
        func didReceiveRegionalPublicMessage(from peerID: PeerID, nickname: String, content: String, timestamp: Date) {}
    }

    // Helper: build a large message packet (unencrypted public message)
    private func makeLargePublicPacket(senderShortHex: PeerID, size: Int) -> BitchatPacket {
        let content = String(repeating: "A", count: size)
        let payload = Data(content.utf8)
        let pkt = BitchatPacket(
            type: MessageType.message.rawValue,
            senderID: Data(hexString: senderShortHex.id) ?? Data(),
            recipientID: nil,
            timestamp: UInt64(Date().timeIntervalSince1970 * 1000),
            payload: payload,
            signature: nil,
            ttl: 7
        )
        return pkt
    }

    // Helper: fragment a packet using the same header format BLEService expects
    private func fragmentPacket(_ packet: BitchatPacket, fragmentSize: Int, fragmentID: Data? = nil, pad: Bool = true) -> [BitchatPacket] {
        guard let fullData = packet.toBinaryData(padding: pad) else { return [] }
        let fid = fragmentID ?? Data((0..<8).map { _ in UInt8.random(in: 0...255) })
        let chunks: [Data] = stride(from: 0, to: fullData.count, by: fragmentSize).map { off in
            Data(fullData[off..<min(off + fragmentSize, fullData.count)])
        }
        let total = UInt16(chunks.count)
        var packets: [BitchatPacket] = []
        for (i, chunk) in chunks.enumerated() {
            var payload = Data()
            payload.append(fid)
            var idxBE = UInt16(i).bigEndian
            var totBE = total.bigEndian
            withUnsafeBytes(of: &idxBE) { payload.append(contentsOf: $0) }
            withUnsafeBytes(of: &totBE) { payload.append(contentsOf: $0) }
            payload.append(packet.type)
            payload.append(chunk)
            let fpkt = BitchatPacket(
                type: MessageType.fragment.rawValue,
                senderID: packet.senderID,
                recipientID: packet.recipientID,
                timestamp: packet.timestamp,
                payload: payload,
                signature: nil,
                ttl: packet.ttl
            )
            packets.append(fpkt)
        }
        return packets
    }
}