fix: add input validation for protocol decoding and fragment reassembly (#666)

* fix: add input validation for protocol decoding and fragment reassembly

* fix:subtract the old entry's size before adding the new one, so duplicate/retransmitted fragments don't inflate the counter

* fix: FragmentManager.handleFragment() can be entered concurrently (e.g., fragments for the same fragmentID arriving from multiple peers/relays). In order for this to happen multiple devices would be needed to connect to the mesh. even with a per-fragmentID byte cap, an attacker could open many fragment IDs at once and force the device to buffer lots of fragment data overall (risking memory pressure/OOM). In this fix we made fragments atomic and thread safe. When a fragment index is retransmitted, we compute the size delta (new - old) so duplicates don’t inflate counters and can’t be used to bypass limits. Under heavy load/attack the app will drop/reject fragment sets earlier instead of growing memory usage without bound to reduce risk of oom. tested on pixel 6 and pixel 8.

* fix: add fragmenttest

app/src/main/java/com/bitchat/android/mesh/FragmentManager.kt

@@ -32,6 +32,10 @@ class FragmentManager {
     private val incomingFragments = ConcurrentHashMap<String, MutableMap<Int, ByteArray>>()
     // iOS equivalent: fragmentMetadata: [String: (type: UInt8, total: Int, timestamp: Date)]
     private val fragmentMetadata = ConcurrentHashMap<String, Triple<UByte, Int, Long>>() // originalType, totalFragments, timestamp
+    private val fragmentCumulativeSize = ConcurrentHashMap<String, Int>()
+
+    private val fragmentStateLock = Any()
+    private var globalBufferedBytes: Long = 0L
 
     // Delegate for callbacks
     var delegate: FragmentManagerDelegate? = null
@@ -173,42 +177,98 @@ class FragmentManager {
             
             Log.d(TAG, "Received fragment ${fragmentPayload.index}/${fragmentPayload.total} for fragmentID: $fragmentIDString, originalType: ${fragmentPayload.originalType}")
 
-            // iOS: if incomingFragments[fragmentID] == nil
+            val maxFragments = com.bitchat.android.util.AppConstants.Fragmentation.MAX_FRAGMENTS_PER_ID
-            if (!incomingFragments.containsKey(fragmentIDString)) {
+            if (fragmentPayload.total > maxFragments) {
+                Log.w(TAG, "Rejecting fragment with excessive total count: ${fragmentPayload.total} > $maxFragments")
+                return null
+            }
+
+            synchronized(fragmentStateLock) {
+                fragmentMetadata[fragmentIDString]?.let { (expectedType, expectedTotal, _) ->
+                    if (expectedTotal != fragmentPayload.total || expectedType != fragmentPayload.originalType) {
+                        Log.w(
+                            TAG,
+                            "Rejecting fragment for $fragmentIDString: inconsistent metadata " +
+                                "(expected type=$expectedType total=$expectedTotal, got type=${fragmentPayload.originalType} total=${fragmentPayload.total})"
+                        )
+                        removeFragmentSetLocked(fragmentIDString)
+                        return null
+                    }
+                }
+
+                val isNewSet = !incomingFragments.containsKey(fragmentIDString)
+                if (isNewSet) {
+                    val maxActive = com.bitchat.android.util.AppConstants.Fragmentation.MAX_ACTIVE_FRAGMENT_SETS
+                    if (incomingFragments.size >= maxActive) {
+                        Log.w(TAG, "Rejecting new fragment set $fragmentIDString: active fragment sets ${incomingFragments.size} >= $maxActive")
+                        return null
+                    }
+
                     incomingFragments[fragmentIDString] = mutableMapOf()
                     fragmentMetadata[fragmentIDString] = Triple(
                         fragmentPayload.originalType,
                         fragmentPayload.total,
                         System.currentTimeMillis()
                     )
+                    fragmentCumulativeSize[fragmentIDString] = 0
                 }
 
-            // iOS: incomingFragments[fragmentID]?[index] = Data(fragmentData)
-            incomingFragments[fragmentIDString]?.put(fragmentPayload.index, fragmentPayload.data)
-            
-            // iOS: if let fragments = incomingFragments[fragmentID], fragments.count == total
                 val fragmentMap = incomingFragments[fragmentIDString]
-            if (fragmentMap != null && fragmentMap.size == fragmentPayload.total) {
+                if (fragmentMap == null) {
+                    Log.w(TAG, "Dropping fragment set $fragmentIDString due to missing fragment map")
+                    removeFragmentSetLocked(fragmentIDString)
+                    return null
+                }
+
+                val currentSize = fragmentCumulativeSize[fragmentIDString]
+                if (currentSize == null) {
+                    Log.w(TAG, "Dropping fragment set $fragmentIDString due to missing size tracker")
+                    removeFragmentSetLocked(fragmentIDString)
+                    return null
+                }
+
+                val oldEntrySize = fragmentMap[fragmentPayload.index]?.size ?: 0
+                val newSize = currentSize - oldEntrySize + fragmentPayload.data.size
+                val maxTotalBytes = com.bitchat.android.util.AppConstants.Fragmentation.MAX_FRAGMENT_TOTAL_BYTES
+                if (newSize > maxTotalBytes) {
+                    Log.w(TAG, "Rejecting fragment for $fragmentIDString: cumulative size $newSize exceeds cap $maxTotalBytes")
+                    removeFragmentSetLocked(fragmentIDString)
+                    return null
+                }
+
+                val delta = (fragmentPayload.data.size - oldEntrySize).toLong()
+                val maxGlobalBytes = com.bitchat.android.util.AppConstants.Fragmentation.MAX_GLOBAL_FRAGMENT_TOTAL_BYTES
+                if (globalBufferedBytes + delta > maxGlobalBytes) {
+                    Log.w(
+                        TAG,
+                        "Rejecting fragment for $fragmentIDString: global buffered bytes ${(globalBufferedBytes + delta)} exceeds cap $maxGlobalBytes"
+                    )
+                    if (isNewSet) {
+                        removeFragmentSetLocked(fragmentIDString)
+                    }
+                    return null
+                }
+
+                fragmentMap[fragmentPayload.index] = fragmentPayload.data
+                fragmentCumulativeSize[fragmentIDString] = newSize
+                globalBufferedBytes += delta
+
+                val expectedTotal = fragmentMetadata[fragmentIDString]?.second ?: fragmentPayload.total
+                if (fragmentMap.size == expectedTotal) {
                     Log.d(TAG, "All fragments received for $fragmentIDString, reassembling...")
 
                     // iOS reassembly logic: for i in 0..<total { if let fragment = fragments[i] { reassembled.append(fragment) } }
                     val reassembledData = mutableListOf<Byte>()
-                for (i in 0 until fragmentPayload.total) {
+                    for (i in 0 until expectedTotal) {
                         fragmentMap[i]?.let { data ->
                             reassembledData.addAll(data.asIterable())
                         }
                     }
 
-                // Decode the original packet bytes we reassembled, so flags/compression are preserved - iOS fix
                     val originalPacket = BitchatPacket.fromBinaryData(reassembledData.toByteArray())
                     if (originalPacket != null) {
-                    // iOS cleanup: incomingFragments.removeValue(forKey: fragmentID)
+                        removeFragmentSetLocked(fragmentIDString)
-                    incomingFragments.remove(fragmentIDString)
-                    fragmentMetadata.remove(fragmentIDString)
 
-                    // Suppress re-broadcast of the reassembled packet by zeroing TTL.
-                    // We already relayed the incoming fragments; setting TTL=0 ensures
-                    // PacketRelayManager will skip relaying this reconstructed packet.
                         val suppressedTtlPacket = originalPacket.copy(ttl = 0u.toUByte())
                         Log.d(TAG, "Successfully reassembled original (${reassembledData.size} bytes); set TTL=0 to suppress relay")
                         return suppressedTtlPacket
@@ -217,8 +277,9 @@ class FragmentManager {
                         Log.e(TAG, "Failed to decode reassembled packet (type=${metadata?.first}, total=${metadata?.second})")
                     }
                 } else {
-                val received = fragmentMap?.size ?: 0
+                    val received = fragmentMap.size
-                Log.d(TAG, "Fragment ${fragmentPayload.index} stored, have $received/${fragmentPayload.total} fragments for $fragmentIDString")
+                    Log.d(TAG, "Fragment ${fragmentPayload.index} stored, have $received/$expectedTotal fragments for $fragmentIDString")
+                }
             }
             
         } catch (e: Exception) {
@@ -228,6 +289,15 @@ class FragmentManager {
         return null
     }
 
+    private fun removeFragmentSetLocked(fragmentIDString: String) {
+        incomingFragments.remove(fragmentIDString)
+        fragmentMetadata.remove(fragmentIDString)
+        val bytes = fragmentCumulativeSize.remove(fragmentIDString)?.toLong() ?: 0L
+        if (bytes != 0L) {
+            globalBufferedBytes = (globalBufferedBytes - bytes).coerceAtLeast(0L)
+        }
+    }
+    
     /**
      * Helper function to match iOS stride functionality
      * stride(from: 0, to: fullData.count, by: maxFragmentSize)
@@ -247,38 +317,42 @@ class FragmentManager {
      * Clean old fragments (> 30 seconds old)
      */
     private fun cleanupOldFragments() {
+        synchronized(fragmentStateLock) {
             val now = System.currentTimeMillis()
             val cutoff = now - FRAGMENT_TIMEOUT
 
             // iOS: let oldFragments = fragmentMetadata.filter { $0.value.timestamp < cutoff }.map { $0.key }
             val oldFragments = fragmentMetadata.filter { it.value.third < cutoff }.map { it.key }
 
-        // iOS: for fragmentID in oldFragments { incomingFragments.removeValue(forKey: fragmentID) }
             for (fragmentID in oldFragments) {
-            incomingFragments.remove(fragmentID)
+                removeFragmentSetLocked(fragmentID)
-            fragmentMetadata.remove(fragmentID)
             }
 
             if (oldFragments.isNotEmpty()) {
                 Log.d(TAG, "Cleaned up ${oldFragments.size} old fragment sets (iOS compatible)")
             }
         }
+    }
     
     /**
      * Get debug information - matches iOS debugging
      */
     fun getDebugInfo(): String {
+        synchronized(fragmentStateLock) {
             return buildString {
                 appendLine("=== Fragment Manager Debug Info (iOS Compatible) ===")
                 appendLine("Active Fragment Sets: ${incomingFragments.size}")
                 appendLine("Fragment Size Threshold: $FRAGMENT_SIZE_THRESHOLD bytes")
                 appendLine("Max Fragment Size: $MAX_FRAGMENT_SIZE bytes")
+                appendLine("Global Buffered Bytes: $globalBufferedBytes")
 
                 fragmentMetadata.forEach { (fragmentID, metadata) ->
                     val (originalType, totalFragments, timestamp) = metadata
                     val received = incomingFragments[fragmentID]?.size ?: 0
                     val ageSeconds = (System.currentTimeMillis() - timestamp) / 1000
-                appendLine("  - $fragmentID: $received/$totalFragments fragments, type: $originalType, age: ${ageSeconds}s")
+                    val bytes = fragmentCumulativeSize[fragmentID] ?: 0
+                    appendLine("  - $fragmentID: $received/$totalFragments fragments, bytes=$bytes, type: $originalType, age: ${ageSeconds}s")
+                }
             }
         }
     }
@@ -299,8 +373,12 @@ class FragmentManager {
      * Clear all fragments
      */
     fun clearAllFragments() {
+        synchronized(fragmentStateLock) {
             incomingFragments.clear()
             fragmentMetadata.clear()
+            fragmentCumulativeSize.clear()
+            globalBufferedBytes = 0L
+        }
     }
     
     /**

app/src/main/java/com/bitchat/android/protocol/BinaryProtocol.kt

@@ -365,7 +365,11 @@ object BinaryProtocol {
                 buffer.getShort().toUShort().toUInt()  // 2 bytes for v1, convert to UInt
             }
 
-            // Calculate expected total size
+            if (payloadLength > com.bitchat.android.util.AppConstants.Protocol.MAX_PAYLOAD_LENGTH.toUInt()) {
+                Log.w("BinaryProtocol", "Payload length ${payloadLength} exceeds maximum allowed (${com.bitchat.android.util.AppConstants.Protocol.MAX_PAYLOAD_LENGTH})")
+                return null
+            }
+
             var expectedSize = headerSize + SENDER_ID_SIZE + payloadLength.toInt()
             if (hasRecipient) expectedSize += RECIPIENT_ID_SIZE
             var routeCount = 0
@@ -467,7 +471,7 @@ object BinaryProtocol {
                 route = route
             )
             
-        } catch (e: Exception) {
+        } catch (e: Throwable) {
             Log.e("BinaryProtocol", "Error decoding packet: ${e.message}")
             return null
         }

app/src/main/java/com/bitchat/android/util/AppConstants.kt

@@ -41,6 +41,10 @@ object AppConstants {
         const val MAX_FRAGMENT_SIZE: Int = 469
         const val FRAGMENT_TIMEOUT_MS: Long = 30_000L
         const val CLEANUP_INTERVAL_MS: Long = 10_000L
+        const val MAX_FRAGMENTS_PER_ID: Int = 256
+        const val MAX_FRAGMENT_TOTAL_BYTES: Int = 1_048_576
+        const val MAX_ACTIVE_FRAGMENT_SETS: Int = 64
+        const val MAX_GLOBAL_FRAGMENT_TOTAL_BYTES: Long = 4L * 1_048_576L
     }
 
     object Security {
@@ -64,6 +68,7 @@ object AppConstants {
 
     object Protocol {
         const val COMPRESSION_THRESHOLD_BYTES: Int = 100
+        const val MAX_PAYLOAD_LENGTH: Int = 10_485_760
     }
 
     object StoreForward {