From c1430eaeb91dd17eefb0ddb0d693f90451f230c8 Mon Sep 17 00:00:00 2001
From: jack <jackjackbits@users.noreply.github.com>
Date: Wed, 1 Oct 2025 02:40:45 +0200
Subject: [PATCH] Fix critical security issues in fragment reassembly and file
 cleanup

Fragment Reassembly Race Condition (CRITICAL):
- Wrap all incomingFragments/fragmentMetadata access in collectionsQueue.sync
- Prevents concurrent modification crashes from multi-threaded access
- Minimizes lock contention by doing heavy work (reassembly/decode) outside locks
- Add upper bound check: reject fragments with total > 10,000 (DoS prevention)
- Add cumulative size validation before storing fragments (memory DoS prevention)

File Cleanup Path Traversal (CRITICAL):
- Use NSString.lastPathComponent to extract filename safely
- Prevents directory traversal attacks via malicious filenames
- Add path prefix validation before file deletion
- Now checks both incoming and outgoing directories (fixes disk leak)

Additional Protections:
- Fragment assemblies now limited by both count (128) and cumulative bytes (1MB)
- Explicit checks for "." and ".." filenames in cleanup
- Defense-in-depth: multiple validation layers
---
 bitchat/Services/BLEService.swift      | 99 +++++++++++++++++---------
 bitchat/ViewModels/ChatViewModel.swift | 41 +++++++----
 2 files changed, 91 insertions(+), 49 deletions(-)

diff --git a/bitchat/Services/BLEService.swift b/bitchat/Services/BLEService.swift
index bd2000c8..3c243be0 100644
--- a/bitchat/Services/BLEService.swift
+++ b/bitchat/Services/BLEService.swift
@@ -1255,7 +1255,7 @@ final class BLEService: NSObject {
         if peerID == myPeerID {
             return
         }
-        
+
         // Minimum header: 8 bytes ID + 2 index + 2 total + 1 type
         guard packet.payload.count >= 13 else { return }
 
@@ -1274,47 +1274,76 @@ final class BLEService: NSObject {
         let originalType = packet.payload[12]
         let fragmentData = packet.payload.suffix(from: 13)
 
-        // Sanity checks
-        guard total > 0 && index >= 0 && index < total else { return }
+        // Sanity checks - add reasonable upper bound on total to prevent DoS
+        guard total > 0 && total <= 10000 && index >= 0 && index < total else { return }
 
-        // Store fragment
+        // Compute fragment key for this assembly
         let key = FragmentKey(sender: senderU64, id: fragU64)
-        if incomingFragments[key] == nil {
-            // Cap in-flight assemblies to prevent memory/battery blowups
-            if incomingFragments.count >= maxInFlightAssemblies {
-                // Evict the oldest assembly by timestamp
-                if let oldest = fragmentMetadata.min(by: { $0.value.timestamp < $1.value.timestamp })?.key {
-                    incomingFragments.removeValue(forKey: oldest)
-                    fragmentMetadata.removeValue(forKey: oldest)
-                }
-            }
-            incomingFragments[key] = [:]
-            fragmentMetadata[key] = (originalType, total, Date())
-            SecureLogger.debug("📦 Started fragment assembly id=\(String(format: "%016llx", fragU64)) total=\(total)", category: .session)
-        }
-        incomingFragments[key]?[index] = Data(fragmentData)
-        SecureLogger.debug("📦 Fragment \(index + 1)/\(total) (len=\(fragmentData.count)) for id=\(String(format: "%016llx", fragU64))", category: .session)
 
-        // Check if complete
-        if let fragments = incomingFragments[key],
-           fragments.count == total {
-            // Reassemble
-            var reassembled = Data()
-            for i in 0..<total {
-                if let fragment = fragments[i] {
-                    reassembled.append(fragment)
+        // Critical section: Store fragment and check completion status
+        var shouldReassemble: Bool = false
+        var fragmentsToReassemble: [Int: Data]? = nil
+
+        collectionsQueue.sync(flags: .barrier) {
+            if incomingFragments[key] == nil {
+                // Cap in-flight assemblies to prevent memory/battery blowups
+                if incomingFragments.count >= maxInFlightAssemblies {
+                    // Evict the oldest assembly by timestamp
+                    if let oldest = fragmentMetadata.min(by: { $0.value.timestamp < $1.value.timestamp })?.key {
+                        incomingFragments.removeValue(forKey: oldest)
+                        fragmentMetadata.removeValue(forKey: oldest)
+                    }
                 }
+                incomingFragments[key] = [:]
+                fragmentMetadata[key] = (originalType, total, Date())
+                SecureLogger.debug("📦 Started fragment assembly id=\(String(format: "%016llx", fragU64)) total=\(total)", category: .session)
             }
-            
-            // Decode the original packet bytes we reassembled, so flags/compression are preserved
-            if let originalPacket = BinaryProtocol.decode(reassembled) {
-                SecureLogger.debug("✅ Reassembled packet id=\(String(format: "%016llx", fragU64)) type=\(originalPacket.type) bytes=\(reassembled.count)", category: .session)
-                handleReceivedPacket(originalPacket, from: peerID)
+
+            // Check cumulative size before storing this fragment
+            let currentSize = incomingFragments[key]?.values.reduce(0) { $0 + $1.count } ?? 0
+            guard currentSize + fragmentData.count <= FileTransferLimits.maxPayloadBytes else {
+                // Exceeds size limit - evict this assembly
+                SecureLogger.warning("🚫 Fragment assembly exceeds size limit (\(currentSize + fragmentData.count) bytes), evicting", category: .security)
+                incomingFragments.removeValue(forKey: key)
+                fragmentMetadata.removeValue(forKey: key)
+                shouldReassemble = false
+                fragmentsToReassemble = nil
+                return
+            }
+
+            incomingFragments[key]?[index] = Data(fragmentData)
+            SecureLogger.debug("📦 Fragment \(index + 1)/\(total) (len=\(fragmentData.count)) for id=\(String(format: "%016llx", fragU64))", category: .session)
+
+            // Check if complete
+            if let fragments = incomingFragments[key], fragments.count == total {
+                shouldReassemble = true
+                fragmentsToReassemble = fragments
             } else {
-                SecureLogger.error("❌ Failed to decode reassembled packet (type=\(originalType), total=\(total))", category: .session)
+                shouldReassemble = false
+                fragmentsToReassemble = nil
             }
-            
-            // Cleanup
+        }
+
+        // Heavy work outside lock: reassemble and decode
+        guard shouldReassemble, let fragments = fragmentsToReassemble else { return }
+
+        var reassembled = Data()
+        for i in 0..<total {
+            if let fragment = fragments[i] {
+                reassembled.append(fragment)
+            }
+        }
+
+        // Decode the original packet bytes we reassembled, so flags/compression are preserved
+        if let originalPacket = BinaryProtocol.decode(reassembled) {
+            SecureLogger.debug("✅ Reassembled packet id=\(String(format: "%016llx", fragU64)) type=\(originalPacket.type) bytes=\(reassembled.count)", category: .session)
+            handleReceivedPacket(originalPacket, from: peerID)
+        } else {
+            SecureLogger.error("❌ Failed to decode reassembled packet (type=\(originalType), total=\(total))", category: .session)
+        }
+
+        // Critical section: Cleanup completed assembly
+        collectionsQueue.sync(flags: .barrier) {
             incomingFragments.removeValue(forKey: key)
             fragmentMetadata.removeValue(forKey: key)
         }
diff --git a/bitchat/ViewModels/ChatViewModel.swift b/bitchat/ViewModels/ChatViewModel.swift
index 0ccc88f0..0955378b 100644
--- a/bitchat/ViewModels/ChatViewModel.swift
+++ b/bitchat/ViewModels/ChatViewModel.swift
@@ -2712,21 +2712,34 @@ final class ChatViewModel: ObservableObject, BitchatDelegate {
     }
 
     private func cleanupLocalFile(forMessage message: BitchatMessage) {
-        let prefixes = ["[voice] ": "voicenotes/outgoing",
-                        "[image] ": "images/outgoing",
-                        "[file] ": "files/outgoing"]
-        guard let entry = prefixes.first(where: { message.content.hasPrefix($0.key) }) else { return }
-        let filename = String(message.content.dropFirst(entry.key.count)).trimmingCharacters(in: .whitespacesAndNewlines)
-        guard !filename.isEmpty, let base = try? applicationFilesDirectory() else { return }
-        let target = base.appendingPathComponent(entry.value, isDirectory: true).appendingPathComponent(filename)
+        // Check both outgoing and incoming directories for thorough cleanup
+        let prefixes = ["[voice] ", "[image] ", "[file] "]
+        let subdirs = ["voicenotes/outgoing", "voicenotes/incoming",
+                       "images/outgoing", "images/incoming",
+                       "files/outgoing", "files/incoming"]
 
-        // Security: Remove file directly (no TOCTOU race)
-        do {
-            try FileManager.default.removeItem(at: target)
-        } catch CocoaError.fileNoSuchFile {
-            // Expected
-        } catch {
-            SecureLogger.error("Failed to cleanup \(filename): \(error)", category: .session)
+        guard let prefix = prefixes.first(where: { message.content.hasPrefix($0) }) else { return }
+        let rawFilename = String(message.content.dropFirst(prefix.count)).trimmingCharacters(in: .whitespacesAndNewlines)
+        guard !rawFilename.isEmpty, let base = try? applicationFilesDirectory() else { return }
+
+        // Security: Extract only the last path component to prevent directory traversal
+        let safeFilename = (rawFilename as NSString).lastPathComponent
+        guard !safeFilename.isEmpty && safeFilename != "." && safeFilename != ".." else { return }
+
+        // Try all possible locations (outgoing and incoming)
+        for subdir in subdirs {
+            let target = base.appendingPathComponent(subdir, isDirectory: true).appendingPathComponent(safeFilename)
+
+            // Security: Verify target is within expected directory before deletion
+            guard target.path.hasPrefix(base.path) else { continue }
+
+            do {
+                try FileManager.default.removeItem(at: target)
+            } catch CocoaError.fileNoSuchFile {
+                // Expected - file not in this directory
+            } catch {
+                SecureLogger.error("Failed to cleanup \(safeFilename): \(error)", category: .session)
+            }
         }
     }