fix: add input validation for protocol decoding and fragment reassembly

[0x0v1]

No description provided.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 85a0499ae3

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[callebtc]

Thank you for this PR! The OOM protections and fragment caps are exactly what is needed to prevent memory exhaustion and malicious payload attacks. The fix in BinaryProtocol.kt handling Throwable for OOM errors is a great touch.

However, there is a critical concurrency issue in FragmentManager.kt that must be addressed before this can be merged. handleFragment is a suspend function executed concurrently, and the current implementation introduces check-then-act race conditions and uses a non-thread-safe inner collection (mutableMapOf()), which can lead to ConcurrentModificationExceptions, dropped fragments, or bypassed size caps.

The minimal fix:

Replace the 29 lines of size tracking and insertion with an atomic compute block, and initialize the inner map as a ConcurrentHashMap.

            val fragmentIDString = fragmentPayload.getFragmentIDString()
            
            Log.d(TAG, "Received fragment ${fragmentPayload.index}/${fragmentPayload.total} for fragmentID: $fragmentIDString, originalType: ${fragmentPayload.originalType}")

            val maxFragments = com.bitchat.android.util.AppConstants.Fragmentation.MAX_FRAGMENTS_PER_ID
            if (fragmentPayload.total > maxFragments) {
                Log.w(TAG, "Rejecting fragment with excessive total count: ${fragmentPayload.total} > $maxFragments")
                return null
            }

            val maxTotalBytes = com.bitchat.android.util.AppConstants.Fragmentation.MAX_FRAGMENT_TOTAL_BYTES
            var rejected = false
            
            // Use compute to atomically update the maps for this specific fragmentID
            incomingFragments.compute(fragmentIDString) { _, currentMap ->
                val map = currentMap ?: java.util.concurrent.ConcurrentHashMap<Int, ByteArray>().also {
                    // Initialize metadata for new fragments
                    fragmentMetadata[fragmentIDString] = Triple(
                        fragmentPayload.originalType,
                        fragmentPayload.total,
                        System.currentTimeMillis()
                    )
                    fragmentCumulativeSize[fragmentIDString] = 0
                }

                val currentSize = fragmentCumulativeSize[fragmentIDString] ?: 0
                val oldEntrySize = map[fragmentPayload.index]?.size ?: 0
                val newSize = currentSize - oldEntrySize + fragmentPayload.data.size

                if (newSize > maxTotalBytes) {
                    Log.w(TAG, "Rejecting fragment for $fragmentIDString: cumulative size $newSize exceeds cap $maxTotalBytes")
                    fragmentMetadata.remove(fragmentIDString)
                    fragmentCumulativeSize.remove(fragmentIDString)
                    rejected = true
                    return@compute null // Returning null removes the entry from incomingFragments
                }

                fragmentCumulativeSize[fragmentIDString] = newSize
                map[fragmentPayload.index] = fragmentPayload.data
                map
            }

            if (rejected) return null

            val fragmentMap = incomingFragments[fragmentIDString]
            if (fragmentMap != null && fragmentMap.size == fragmentPayload.total) {
                Log.d(TAG, "All fragments received for $fragmentIDString, reassembling...")

This guarantees atomicity per fragmentIDString and eliminates the risk of ConcurrentModificationException while correctly applying the byte size caps under concurrent load. Please apply this change and I'll gladly approve!