bitchat-android/app/src/test/java/com/bitchat/android/mesh/FragmentManagerTest.kt

package com.bitchat.android.mesh

import com.bitchat.android.protocol.BitchatPacket
import com.bitchat.android.protocol.MessageType
import com.bitchat.android.model.FragmentPayload
import org.junit.Assert.assertEquals
import org.junit.Assert.assertNotNull
import org.junit.Assert.assertNull
import org.junit.Assert.assertThrows
import org.junit.Assert.assertTrue
import org.junit.Before
import org.junit.Test
import org.junit.runner.RunWith
import org.robolectric.RobolectricTestRunner
import java.util.Random

@RunWith(RobolectricTestRunner::class)
class FragmentManagerTest {

    private lateinit var fragmentManager: FragmentManager
    private val senderID = "1122334455667788"
    private val recipientID = "8877665544332211"

    @Before
    fun setup() {
        fragmentManager = FragmentManager()
    }

    @Test
    fun `test fragmentation without route`() {
        // Create a large payload (e.g., 1000 bytes)
        val payload = ByteArray(1000)
        Random().nextBytes(payload)

        val packet = BitchatPacket(
            version = 1u,
            type = MessageType.MESSAGE.value,
            senderID = hexStringToByteArray(senderID),
            recipientID = hexStringToByteArray(recipientID),
            timestamp = System.currentTimeMillis().toULong(),
            payload = payload,
            ttl = 7u,
            route = null
        )

        val fragments = fragmentManager.createFragments(packet)

        assertTrue("Should create multiple fragments", fragments.size > 1)
        
        // Verify each fragment fits in MTU (512)
        for (fragment in fragments) {
            val encodedSize = fragment.toBinaryData()?.size ?: 0
            assertTrue("Fragment encoded size should be <= 512, was $encodedSize", encodedSize <= 512)
            
            // Inspect the payload data size
            val fragmentPayload = FragmentPayload.decode(fragment.payload)
            assertNotNull(fragmentPayload)
        }
    }

    @Test
    fun `test fragmentation with route`() {
        // Create a large payload
        val payload = ByteArray(1000)
        Random().nextBytes(payload)
        
        // Create a fake route (3 hops)
        val route = listOf(
            hexStringToByteArray("AABBCCDDEEFF0011"),
            hexStringToByteArray("1100FFEEDDCCBBAA"),
            hexStringToByteArray("1234567890ABCDEF")
        )

        val packet = BitchatPacket(
            version = 2u,
            type = MessageType.MESSAGE.value,
            senderID = hexStringToByteArray(senderID),
            recipientID = hexStringToByteArray(recipientID),
            timestamp = System.currentTimeMillis().toULong(),
            payload = payload,
            ttl = 7u,
            route = route
        )

        val fragments = fragmentManager.createFragments(packet)

        assertTrue("Should create multiple fragments", fragments.size > 1)

        // Verify fragments retain the route and version 2
        for (fragment in fragments) {
            assertEquals("Fragment version should be 2", 2u.toUByte(), fragment.version)
            assertEquals("Fragment should have the route", route.size, fragment.route?.size)
            
            val encodedSize = fragment.toBinaryData()?.size ?: 0
            assertTrue("Fragment encoded size should be <= 512, was $encodedSize", encodedSize <= 512)
        }
    }
    
    @Test
    fun `test fragmentation size difference with and without route`() {
        // This test specifically checks if the dynamic calculation logic works 
        // by observing that fragments with routes carry less data payload per fragment
        
        val payload = ByteArray(2000) // Large enough to ensure full fragments
        Random().nextBytes(payload)
        
        // 1. Without route
        val packetNoRoute = BitchatPacket(
            version = 1u,
            type = MessageType.MESSAGE.value,
            senderID = hexStringToByteArray(senderID),
            recipientID = hexStringToByteArray(recipientID),
            timestamp = System.currentTimeMillis().toULong(),
            payload = payload,
            ttl = 7u,
            route = null
        )
        val fragmentsNoRoute = fragmentManager.createFragments(packetNoRoute)
        val firstFragPayloadNoRoute = FragmentPayload.decode(fragmentsNoRoute[0].payload)
        val dataSizeNoRoute = firstFragPayloadNoRoute?.data?.size ?: 0
        
        // 2. With large route (e.g., 5 hops)
        val route = List(5) { hexStringToByteArray("000000000000000$it") }
        val packetWithRoute = BitchatPacket(
            version = 2u,
            type = MessageType.MESSAGE.value,
            senderID = hexStringToByteArray(senderID),
            recipientID = hexStringToByteArray(recipientID),
            timestamp = System.currentTimeMillis().toULong(),
            payload = payload,
            ttl = 7u,
            route = route
        )
        val fragmentsWithRoute = fragmentManager.createFragments(packetWithRoute)
        val firstFragPayloadWithRoute = FragmentPayload.decode(fragmentsWithRoute[0].payload)
        val dataSizeWithRoute = firstFragPayloadWithRoute?.data?.size ?: 0
        
        println("Data size without route: $dataSizeNoRoute")
        println("Data size with route: $dataSizeWithRoute")
        
        assertTrue("Data payload should be smaller with route", dataSizeWithRoute < dataSizeNoRoute)
        
        // Rough verification of the math:
        // 5 hops * 8 bytes = 40 bytes extra.
        // Plus v2 header overhead differences.
        // The difference should be roughly 40+ bytes.
        assertTrue("Difference should be significant", (dataSizeNoRoute - dataSizeWithRoute) >= 40)
    }

    @Test
    fun `test reassembly`() {
        val originalPayload = ByteArray(1500)
        Random().nextBytes(originalPayload)
        
        val originalPacket = BitchatPacket(
            version = 1u,
            type = MessageType.FILE_TRANSFER.value,
            senderID = hexStringToByteArray(senderID),
            recipientID = hexStringToByteArray(recipientID),
            timestamp = System.currentTimeMillis().toULong(),
            payload = originalPayload,
            ttl = 7u
        )
        
        val fragments = fragmentManager.createFragments(originalPacket)
        
        var reassembledPacket: BitchatPacket? = null
        
        // Feed fragments back into FragmentManager
        // Note: FragmentManager stores state in incomingFragments
        
        for (fragment in fragments) {
            val result = fragmentManager.handleFragment(fragment)
            if (result != null) {
                reassembledPacket = result
            }
        }
        
        assertNotNull("Should have reassembled packet", reassembledPacket)
        assertEquals("Type should match", originalPacket.type, reassembledPacket!!.type)
        assertEquals("Payload size should match", originalPacket.payload.size, reassembledPacket.payload.size)
        assertTrue("Payload content should match", originalPacket.payload.contentEquals(reassembledPacket.payload))
    }

    @Test
    fun `inbound fragment set above 256 is rejected`() {
        val payload = FragmentPayload(
            fragmentID = ByteArray(8) { 1 },
            index = 0,
            total = 257,
            originalType = MessageType.NOISE_ENCRYPTED.value,
            data = byteArrayOf(1)
        ).encode()
        val packet = BitchatPacket(
            version = 1u,
            type = MessageType.FRAGMENT.value,
            senderID = hexStringToByteArray(senderID),
            recipientID = hexStringToByteArray(recipientID),
            timestamp = 1u,
            payload = payload,
            ttl = 7u
        )

        assertNull(fragmentManager.handleFragment(packet))
    }

    @Test
    fun `fragment payload refuses UInt16 truncation`() {
        assertThrows(IllegalArgumentException::class.java) {
            FragmentPayload(
                fragmentID = ByteArray(8) { 2 },
                index = 0,
                total = 65_536,
                originalType = MessageType.NOISE_ENCRYPTED.value,
                data = byteArrayOf(1)
            ).encode()
        }
    }

    @Test
    fun `generic public packet retains a 257 fragment outbound plan`() {
        val randomPayload = ByteArray(180 * 1024).also { Random(0xB17C4A7).nextBytes(it) }

        fun plan(contentSize: Int): List<BitchatPacket> = fragmentManager.createFragments(
            BitchatPacket(
                version = 2u,
                type = MessageType.FILE_TRANSFER.value,
                senderID = hexStringToByteArray(senderID),
                recipientID = com.bitchat.android.protocol.SpecialRecipients.BROADCAST,
                timestamp = 1u,
                payload = randomPayload.copyOf(contentSize),
                signature = ByteArray(64) { 7 },
                ttl = 7u
            )
        )

        var low = 1
        var high = randomPayload.size
        while (low < high) {
            val mid = low + (high - low) / 2
            if (plan(mid).size >= 257) high = mid else low = mid + 1
        }

        assertEquals(257, plan(low).size)
    }

    private fun hexStringToByteArray(hexString: String): ByteArray {
        val result = ByteArray(8)
        for (i in 0 until 8) {
            val byteStr = hexString.substring(i * 2, i * 2 + 2)
            result[i] = byteStr.toInt(16).toByte()
        }
        return result
    }

    // Injected-fragment hardening. Fragment packets are unauthenticated and
    // both halves of the reassembly key are attacker-choosable, so a stream in
    // flight has to survive a crafted packet aimed at it. Ports the storage
    // rules from permissionlesstech/bitchat#1515.

    /** A real packet, fragmented the way the sender would fragment it. */
    private fun honestStream(): Pair<BitchatPacket, List<BitchatPacket>> {
        val payload = ByteArray(1200) { (it % 251).toByte() }
        val packet = BitchatPacket(
            version = 1u,
            type = MessageType.FILE_TRANSFER.value,
            senderID = hexStringToByteArray(senderID),
            recipientID = hexStringToByteArray(recipientID),
            timestamp = 1u,
            payload = payload,
            ttl = 7u
        )
        val fragments = fragmentManager.createFragments(packet)
        assertTrue("Need a multi-fragment stream", fragments.size >= 3)
        return packet to fragments
    }

    /** Same stream key as [honest], but attacker-chosen contents. */
    private fun injected(
        honest: BitchatPacket,
        index: Int,
        total: Int,
        data: ByteArray,
        originalType: UByte
    ): BitchatPacket {
        val decoded = FragmentPayload.decode(honest.payload)!!
        return honest.copy(
            payload = FragmentPayload(
                fragmentID = decoded.fragmentID,
                index = index,
                total = total,
                originalType = originalType,
                data = data
            ).encode()
        )
    }

    private fun feedAllButLast(fragments: List<BitchatPacket>) {
        fragments.dropLast(1).forEach {
            assertNull("No fragment before the last should complete", fragmentManager.handleFragment(it))
        }
    }

    private fun assertCompletes(original: BitchatPacket, last: BitchatPacket) {
        val reassembled = fragmentManager.handleFragment(last)
        assertNotNull("The honest stream must still complete", reassembled)
        assertTrue(
            "Reassembled payload must be the honest bytes",
            original.payload.contentEquals(reassembled!!.payload)
        )
    }

    @Test
    fun `a conflicting total is rejected without destroying the stream`() {
        val (original, fragments) = honestStream()
        feedAllButLast(fragments)

        val decoded = FragmentPayload.decode(fragments[0].payload)!!
        assertNull(
            fragmentManager.handleFragment(
                injected(fragments[0], index = 0, total = decoded.total - 1,
                         data = ByteArray(4), originalType = decoded.originalType)
            )
        )

        assertCompletes(original, fragments.last())
    }

    @Test
    fun `a conflicting original type is rejected without destroying the stream`() {
        val (original, fragments) = honestStream()
        feedAllButLast(fragments)

        val decoded = FragmentPayload.decode(fragments[0].payload)!!
        assertNull(
            fragmentManager.handleFragment(
                injected(fragments[0], index = 0, total = decoded.total,
                         data = ByteArray(4), originalType = MessageType.ANNOUNCE.value)
            )
        )

        assertCompletes(original, fragments.last())
    }

    @Test
    fun `an index already held keeps its first bytes`() {
        val (original, fragments) = honestStream()
        feedAllButLast(fragments)

        // Same stream, same index, different bytes: last-wins would swap the
        // real sender's bytes out from under the reassembly.
        val decoded = FragmentPayload.decode(fragments[0].payload)!!
        assertNull(
            fragmentManager.handleFragment(
                injected(fragments[0], index = 0, total = decoded.total,
                         data = ByteArray(decoded.data.size) { 0x66 },
                         originalType = decoded.originalType)
            )
        )

        assertCompletes(original, fragments.last())
    }

    @Test
    fun `an identical duplicate is still accepted`() {
        // Duplicate delivery is normal in a mesh; only a *differing* duplicate
        // is an attack, so redelivery must not be turned into a rejection.
        val (original, fragments) = honestStream()
        feedAllButLast(fragments)

        assertNull(fragmentManager.handleFragment(fragments[0]))

        assertCompletes(original, fragments.last())
    }

    @Test
    fun `an oversized fragment cannot wipe an assembly it did not start`() {
        val (original, fragments) = honestStream()
        feedAllButLast(fragments)

        val decoded = FragmentPayload.decode(fragments[0].payload)!!
        val oversized = ByteArray(
            com.bitchat.android.util.AppConstants.Fragmentation.MAX_FRAGMENT_TOTAL_BYTES + 1
        )
        assertNull(
            fragmentManager.handleFragment(
                injected(fragments[0], index = decoded.total - 1, total = decoded.total,
                         data = oversized, originalType = decoded.originalType)
            )
        )

        assertCompletes(original, fragments.last())
    }
}