Merge pull request #305 from permissionlesstech/robust-connection-state

Fix handshake deadlock and improve connection reliability

bitchat/Noise/NoiseHandshakeCoordinator.swift

@@ -57,10 +57,24 @@ class NoiseHandshakeCoordinator {
     }
     
     /// Check if we should initiate handshake with a peer
-    func shouldInitiateHandshake(myPeerID: String, remotePeerID: String) -> Bool {
+    func shouldInitiateHandshake(myPeerID: String, remotePeerID: String, forceIfStale: Bool = false) -> Bool {
         return handshakeQueue.sync {
             // Check if we're already in an active handshake
             if let state = handshakeStates[remotePeerID], state.isActive {
+                // Check if the handshake is stale and we should force a new one
+                if forceIfStale {
+                    switch state {
+                    case .initiating(_, let lastAttempt):
+                        if Date().timeIntervalSince(lastAttempt) > handshakeTimeout {
+                            SecureLogger.log("Forcing new handshake with \(remotePeerID) - previous stuck in initiating", 
+                                           category: SecureLogger.handshake, level: .warning)
+                            return true
+                        }
+                    default:
+                        break
+                    }
+                }
+                
                 SecureLogger.log("Already in active handshake with \(remotePeerID), state: \(state)", 
                                category: SecureLogger.handshake, level: .debug)
                 return false
@@ -255,6 +269,26 @@ class NoiseHandshakeCoordinator {
         }
     }
     
+    /// Get current retry count for a peer
+    func getRetryCount(for peerID: String) -> Int {
+        return handshakeQueue.sync {
+            switch handshakeStates[peerID] {
+            case .initiating(let attempt, _):
+                return attempt - 1  // Attempts start at 1, retries start at 0
+            default:
+                return 0
+            }
+        }
+    }
+    
+    /// Increment retry count for a peer
+    func incrementRetryCount(for peerID: String) {
+        handshakeQueue.async(flags: .barrier) {
+            let currentAttempt = self.getCurrentAttempt(for: peerID)
+            self.handshakeStates[peerID] = .initiating(attempt: currentAttempt + 1, lastAttempt: Date())
+        }
+    }
+    
     // MARK: - Private Helpers
     
     private func getCurrentAttempt(for peerID: String) -> Int {

bitchat/Noise/NoiseSession.swift

@@ -338,12 +338,14 @@ class NoiseSessionManager {
             var existingSession: NoiseSession? = nil
             
             if let existing = sessions[peerID] {
-                // If we have an established session, reject new handshake attempts
+                // If we have an established session, the peer must have cleared their session
+                // for a good reason (e.g., decryption failure, restart, etc.)
+                // We should accept the new handshake to re-establish encryption
                 if existing.isEstablished() {
-                    // Don't destroy our working session just because the other side is confused
-                    // They should detect the established session through successful message exchange
-                    SecureLogger.log("Rejecting handshake attempt - session already established with \(peerID)", category: SecureLogger.session, level: .debug)
-                    throw NoiseSessionError.alreadyEstablished
+                    SecureLogger.log("Accepting handshake from \(peerID) despite existing session - peer likely cleared their session", 
+                                   category: SecureLogger.session, level: .info)
+                    _ = sessions.removeValue(forKey: peerID)
+                    shouldCreateNew = true
                 } else {
                     // If we're in the middle of a handshake and receive a new initiation,
                     // reset and start fresh (the other side may have restarted)

bitchat/Protocols/BitchatProtocol.swift

@@ -96,6 +96,11 @@ enum MessageType: UInt8 {
     case versionHello = 0x20            // Initial version announcement
     case versionAck = 0x21              // Version acknowledgment
     
+    // Protocol-level acknowledgments
+    case protocolAck = 0x22             // Generic protocol acknowledgment
+    case protocolNack = 0x23            // Negative acknowledgment (failure)
+    case systemValidation = 0x24        // Session validation ping
+    
     var description: String {
         switch self {
         case .announce: return "announce"
@@ -113,6 +118,9 @@ enum MessageType: UInt8 {
         case .noiseIdentityAnnounce: return "noiseIdentityAnnounce"
         case .versionHello: return "versionHello"
         case .versionAck: return "versionAck"
+        case .protocolAck: return "protocolAck"
+        case .protocolNack: return "protocolNack"
+        case .systemValidation: return "systemValidation"
         }
     }
 }
@@ -352,6 +360,172 @@ struct ReadReceipt: Codable {
     }
 }
 
+// MARK: - Protocol Acknowledgments
+
+// Protocol-level acknowledgment for reliable delivery
+struct ProtocolAck: Codable {
+    let originalPacketID: String    // ID of the packet being acknowledged
+    let ackID: String              // Unique ID for this ACK
+    let senderID: String           // Who sent the original packet
+    let receiverID: String         // Who received and is acknowledging
+    let packetType: UInt8          // Type of packet being acknowledged
+    let timestamp: Date            // When ACK was generated
+    let hopCount: UInt8            // Hops taken to reach receiver
+    
+    init(originalPacketID: String, senderID: String, receiverID: String, packetType: UInt8, hopCount: UInt8) {
+        self.originalPacketID = originalPacketID
+        self.ackID = UUID().uuidString
+        self.senderID = senderID
+        self.receiverID = receiverID
+        self.packetType = packetType
+        self.timestamp = Date()
+        self.hopCount = hopCount
+    }
+    
+    // Private init for binary decoding
+    private init(originalPacketID: String, ackID: String, senderID: String, receiverID: String, 
+                 packetType: UInt8, timestamp: Date, hopCount: UInt8) {
+        self.originalPacketID = originalPacketID
+        self.ackID = ackID
+        self.senderID = senderID
+        self.receiverID = receiverID
+        self.packetType = packetType
+        self.timestamp = timestamp
+        self.hopCount = hopCount
+    }
+    
+    func toBinaryData() -> Data {
+        var data = Data()
+        data.appendUUID(originalPacketID)
+        data.appendUUID(ackID)
+        
+        // Sender and receiver IDs as 8-byte hex strings
+        data.append(Data(hexString: senderID) ?? Data(repeating: 0, count: 8))
+        data.append(Data(hexString: receiverID) ?? Data(repeating: 0, count: 8))
+        
+        data.appendUInt8(packetType)
+        data.appendUInt8(hopCount)
+        data.appendDate(timestamp)
+        return data
+    }
+    
+    static func fromBinaryData(_ data: Data) -> ProtocolAck? {
+        let dataCopy = Data(data)
+        guard dataCopy.count >= 50 else { return nil } // 2 UUIDs + 2 IDs + type + hop + timestamp
+        
+        var offset = 0
+        guard let originalPacketID = dataCopy.readUUID(at: &offset),
+              let ackID = dataCopy.readUUID(at: &offset),
+              let senderIDData = dataCopy.readFixedBytes(at: &offset, count: 8),
+              let receiverIDData = dataCopy.readFixedBytes(at: &offset, count: 8),
+              let packetType = dataCopy.readUInt8(at: &offset),
+              let hopCount = dataCopy.readUInt8(at: &offset),
+              let timestamp = dataCopy.readDate(at: &offset) else { return nil }
+        
+        let senderID = senderIDData.hexEncodedString()
+        let receiverID = receiverIDData.hexEncodedString()
+        
+        return ProtocolAck(originalPacketID: originalPacketID,
+                          ackID: ackID,
+                          senderID: senderID,
+                          receiverID: receiverID,
+                          packetType: packetType,
+                          timestamp: timestamp,
+                          hopCount: hopCount)
+    }
+}
+
+// Protocol-level negative acknowledgment
+struct ProtocolNack: Codable {
+    let originalPacketID: String    // ID of the packet that failed
+    let nackID: String             // Unique ID for this NACK
+    let senderID: String           // Who sent the original packet
+    let receiverID: String         // Who is reporting the failure
+    let packetType: UInt8          // Type of packet that failed
+    let timestamp: Date            // When NACK was generated
+    let reason: String             // Reason for failure
+    let errorCode: UInt8           // Numeric error code
+    
+    // Error codes
+    enum ErrorCode: UInt8 {
+        case unknown = 0
+        case checksumFailed = 1
+        case decryptionFailed = 2
+        case malformedPacket = 3
+        case unsupportedVersion = 4
+        case resourceExhausted = 5
+        case routingFailed = 6
+        case sessionExpired = 7
+    }
+    
+    init(originalPacketID: String, senderID: String, receiverID: String, 
+         packetType: UInt8, reason: String, errorCode: ErrorCode = .unknown) {
+        self.originalPacketID = originalPacketID
+        self.nackID = UUID().uuidString
+        self.senderID = senderID
+        self.receiverID = receiverID
+        self.packetType = packetType
+        self.timestamp = Date()
+        self.reason = reason
+        self.errorCode = errorCode.rawValue
+    }
+    
+    // Private init for binary decoding
+    private init(originalPacketID: String, nackID: String, senderID: String, receiverID: String,
+                 packetType: UInt8, timestamp: Date, reason: String, errorCode: UInt8) {
+        self.originalPacketID = originalPacketID
+        self.nackID = nackID
+        self.senderID = senderID
+        self.receiverID = receiverID
+        self.packetType = packetType
+        self.timestamp = timestamp
+        self.reason = reason
+        self.errorCode = errorCode
+    }
+    
+    func toBinaryData() -> Data {
+        var data = Data()
+        data.appendUUID(originalPacketID)
+        data.appendUUID(nackID)
+        
+        // Sender and receiver IDs as 8-byte hex strings
+        data.append(Data(hexString: senderID) ?? Data(repeating: 0, count: 8))
+        data.append(Data(hexString: receiverID) ?? Data(repeating: 0, count: 8))
+        
+        data.appendUInt8(packetType)
+        data.appendUInt8(errorCode)
+        data.appendDate(timestamp)
+        data.appendString(reason)
+        return data
+    }
+    
+    static func fromBinaryData(_ data: Data) -> ProtocolNack? {
+        let dataCopy = Data(data)
+        guard dataCopy.count >= 52 else { return nil } // Minimum size
+        
+        var offset = 0
+        guard let originalPacketID = dataCopy.readUUID(at: &offset),
+              let nackID = dataCopy.readUUID(at: &offset),
+              let senderIDData = dataCopy.readFixedBytes(at: &offset, count: 8),
+              let receiverIDData = dataCopy.readFixedBytes(at: &offset, count: 8),
+              let packetType = dataCopy.readUInt8(at: &offset),
+              let errorCode = dataCopy.readUInt8(at: &offset),
+              let timestamp = dataCopy.readDate(at: &offset),
+              let reason = dataCopy.readString(at: &offset) else { return nil }
+        
+        let senderID = senderIDData.hexEncodedString()
+        let receiverID = receiverIDData.hexEncodedString()
+        
+        return ProtocolNack(originalPacketID: originalPacketID,
+                           nackID: nackID,
+                           senderID: senderID,
+                           receiverID: receiverID,
+                           packetType: packetType,
+                           timestamp: timestamp,
+                           reason: reason,
+                           errorCode: errorCode)
+    }
+}
 
 // MARK: - Peer Identity Rotation
 
@@ -796,6 +970,9 @@ protocol BitchatDelegate: AnyObject {
     func didReceiveDeliveryAck(_ ack: DeliveryAck)
     func didReceiveReadReceipt(_ receipt: ReadReceipt)
     func didUpdateMessageDeliveryStatus(_ messageID: String, status: DeliveryStatus)
+    
+    // Peer availability tracking
+    func peerAvailabilityChanged(_ peerID: String, available: Bool)
 }
 
 // Provide default implementation to make it effectively optional
@@ -815,4 +992,8 @@ extension BitchatDelegate {
     func didUpdateMessageDeliveryStatus(_ messageID: String, status: DeliveryStatus) {
         // Default empty implementation
     }
+    
+    func peerAvailabilityChanged(_ peerID: String, available: Bool) {
+        // Default empty implementation
+    }
 }

bitchat/ViewModels/ChatViewModel.swift

@@ -859,9 +859,11 @@ class ChatViewModel: ObservableObject {
             return primaryColor
         } else if let peerID = message.senderPeerID ?? getPeerIDForNickname(message.sender),
                   let rssi = meshService.getPeerRSSI()[peerID] {
+            // Use actual RSSI value
             return getRSSIColor(rssi: rssi.intValue, colorScheme: colorScheme)
         } else {
-            return primaryColor.opacity(0.9)
+            // No RSSI data available - use a neutral color
+            return primaryColor.opacity(0.7)
         }
     }
     
@@ -960,9 +962,11 @@ class ChatViewModel: ObservableObject {
                 senderColor = primaryColor
             } else if let peerID = message.senderPeerID ?? getPeerIDForNickname(message.sender),
                       let rssi = meshService.getPeerRSSI()[peerID] {
+                // Use actual RSSI value
                 senderColor = getRSSIColor(rssi: rssi.intValue, colorScheme: colorScheme)
             } else {
-                senderColor = primaryColor.opacity(0.9)
+                // No RSSI data available - use a neutral color
+                senderColor = primaryColor.opacity(0.7)
             }
             
             senderStyle.foregroundColor = senderColor
@@ -1103,9 +1107,11 @@ class ChatViewModel: ObservableObject {
                 senderColor = primaryColor
             } else if let peerID = message.senderPeerID ?? getPeerIDForNickname(message.sender),
                       let rssi = meshService.getPeerRSSI()[peerID] {
+                // Use actual RSSI value
                 senderColor = getRSSIColor(rssi: rssi.intValue, colorScheme: colorScheme)
             } else {
-                senderColor = primaryColor.opacity(0.9)
+                // No RSSI data available - use a neutral color
+                senderColor = primaryColor.opacity(0.7)
             }
             
             senderStyle.foregroundColor = senderColor

bitchat/Views/ContentView.swift

@@ -495,7 +495,7 @@ struct ContentView: View {
                         
                         ForEach(sortedPeers, id: \.self) { peerID in
                             let displayName = peerID == myPeerID ? viewModel.nickname : (peerNicknames[peerID] ?? "anon\(peerID.prefix(4))")
-                            let rssi = peerRSSI[peerID]?.intValue ?? -100
+                            let rssi = peerRSSI[peerID]?.intValue
                             let isFavorite = viewModel.isFavorite(peerID: peerID)
                             let isMe = peerID == myPeerID
                             
@@ -511,11 +511,17 @@ struct ContentView: View {
                                         .font(.system(size: 12))
                                         .foregroundColor(Color.orange)
                                         .accessibilityLabel("Unread message from \(displayName)")
-                                } else {
+                                } else if let rssi = rssi {
                                     Image(systemName: "circle.fill")
                                         .font(.system(size: 8))
                                         .foregroundColor(viewModel.getRSSIColor(rssi: rssi, colorScheme: colorScheme))
                                         .accessibilityLabel("Signal strength: \(rssi > -60 ? "excellent" : rssi > -70 ? "good" : rssi > -80 ? "fair" : "poor")")
+                                } else {
+                                    // No RSSI data available
+                                    Image(systemName: "circle")
+                                        .font(.system(size: 8))
+                                        .foregroundColor(Color.secondary.opacity(0.5))
+                                        .accessibilityLabel("Signal strength: unknown")
                                 }
                                 
                                 // Peer name

bitchatTests/Integration/IntegrationTests.swift

@@ -511,6 +511,168 @@ final class IntegrationTests: XCTestCase {
     
     // MARK: - Security Integration Tests
     
+    func testHandshakeAfterNACKDecryptionFailure() throws {
+        // Test the specific scenario where decryption fails, NACK is sent, and handshake is re-established
+        connect("Alice", "Bob")
+        
+        // Establish initial Noise session
+        try establishNoiseSession("Alice", "Bob")
+        
+        let expectation = XCTestExpectation(description: "Handshake re-established after NACK")
+        var nackSent = false
+        var newHandshakeCompleted = false
+        
+        // Exchange some messages to establish nonce state
+        for i in 0..<5 {
+            let msg = try noiseManagers["Alice"]!.encrypt("Message \(i)".data(using: .utf8)!, for: TestConstants.testPeerID2)
+            _ = try noiseManagers["Bob"]!.decrypt(msg, from: TestConstants.testPeerID1)
+        }
+        
+        // Simulate nonce desynchronization - Alice sends messages Bob doesn't receive
+        for _ in 0..<3 {
+            _ = try noiseManagers["Alice"]!.encrypt("Lost message".data(using: .utf8)!, for: TestConstants.testPeerID2)
+        }
+        
+        // Setup Bob's handler to send NACK on decryption failure
+        nodes["Bob"]!.packetDeliveryHandler = { packet in
+            if packet.type == MessageType.noiseEncrypted.rawValue {
+                do {
+                    _ = try self.noiseManagers["Bob"]!.decrypt(packet.payload, from: TestConstants.testPeerID1)
+                } catch {
+                    // Decryption failed - send NACK
+                    nackSent = true
+                    let nack = ProtocolNack(
+                        originalPacketID: UUID().uuidString,
+                        senderID: TestConstants.testPeerID2,
+                        receiverID: TestConstants.testPeerID1,
+                        packetType: packet.type,
+                        reason: "Decryption failed - session out of sync",
+                        errorCode: .decryptionFailed
+                    )
+                    
+                    let nackData = nack.toBinaryData()
+                    let nackPacket = TestHelpers.createTestPacket(
+                        type: MessageType.protocolNack.rawValue,
+                        payload: nackData
+                    )
+                    self.nodes["Alice"]!.simulateIncomingPacket(nackPacket)
+                    
+                    // Bob clears session
+                    self.noiseManagers["Bob"]!.removeSession(for: TestConstants.testPeerID1)
+                }
+            } else if packet.type == MessageType.noiseHandshakeInit.rawValue {
+                // Bob receives handshake init from Alice after NACK
+                do {
+                    let response = try self.noiseManagers["Bob"]!.handleIncomingHandshake(
+                        from: TestConstants.testPeerID1,
+                        message: packet.payload
+                    )
+                    if let resp = response {
+                        let responsePacket = TestHelpers.createTestPacket(
+                            type: MessageType.noiseHandshakeResp.rawValue,
+                            payload: resp
+                        )
+                        self.nodes["Alice"]!.simulateIncomingPacket(responsePacket)
+                    }
+                } catch {
+                    XCTFail("Bob failed to handle handshake: \(error)")
+                }
+            }
+        }
+        
+        // Setup Alice's handler to clear session on NACK and initiate handshake
+        nodes["Alice"]!.packetDeliveryHandler = { packet in
+            if packet.type == MessageType.protocolNack.rawValue {
+                // Alice receives NACK - clear session
+                self.noiseManagers["Alice"]!.removeSession(for: TestConstants.testPeerID2)
+                
+                // Initiate new handshake
+                do {
+                    let handshakeInit = try self.noiseManagers["Alice"]!.initiateHandshake(with: TestConstants.testPeerID2)
+                    let handshakePacket = TestHelpers.createTestPacket(
+                        type: MessageType.noiseHandshakeInit.rawValue,
+                        payload: handshakeInit
+                    )
+                    self.nodes["Bob"]!.simulateIncomingPacket(handshakePacket)
+                } catch {
+                    XCTFail("Alice failed to initiate handshake: \(error)")
+                }
+            } else if packet.type == MessageType.noiseHandshakeResp.rawValue {
+                // Complete handshake
+                do {
+                    let final = try self.noiseManagers["Alice"]!.handleIncomingHandshake(
+                        from: TestConstants.testPeerID2,
+                        message: packet.payload
+                    )
+                    if let finalMsg = final {
+                        let finalPacket = TestHelpers.createTestPacket(
+                            type: MessageType.noiseHandshakeResp.rawValue,
+                            payload: finalMsg
+                        )
+                        self.nodes["Bob"]!.simulateIncomingPacket(finalPacket)
+                        
+                        // Try sending a message with new session
+                        let testMsg = try self.noiseManagers["Alice"]!.encrypt(
+                            "After re-handshake".data(using: .utf8)!,
+                            for: TestConstants.testPeerID2
+                        )
+                        let msgPacket = TestHelpers.createTestPacket(
+                            type: MessageType.noiseEncrypted.rawValue,
+                            payload: testMsg
+                        )
+                        self.nodes["Bob"]!.simulateIncomingPacket(msgPacket)
+                    }
+                } catch {
+                    XCTFail("Alice failed to complete handshake: \(error)")
+                }
+            }
+        }
+        
+        // Add final handler to verify message works
+        let originalHandler = nodes["Bob"]!.packetDeliveryHandler
+        nodes["Bob"]!.packetDeliveryHandler = { packet in
+            originalHandler?(packet)
+            
+            if packet.type == MessageType.noiseHandshakeResp.rawValue && packet.senderID.hexEncodedString() == TestConstants.testPeerID1 {
+                // Final handshake message received
+                do {
+                    _ = try self.noiseManagers["Bob"]!.handleIncomingHandshake(
+                        from: TestConstants.testPeerID1,
+                        message: packet.payload
+                    )
+                } catch {
+                    XCTFail("Bob failed to complete handshake: \(error)")
+                }
+            } else if packet.type == MessageType.noiseEncrypted.rawValue && nackSent {
+                // Try to decrypt with new session
+                do {
+                    let decrypted = try self.noiseManagers["Bob"]!.decrypt(packet.payload, from: TestConstants.testPeerID1)
+                    if let msg = String(data: decrypted, encoding: .utf8), msg == "After re-handshake" {
+                        newHandshakeCompleted = true
+                        expectation.fulfill()
+                    }
+                } catch {
+                    XCTFail("Bob failed to decrypt after re-handshake: \(error)")
+                }
+            }
+        }
+        
+        // Trigger the scenario - send desynchronized message
+        let desyncMsg = try noiseManagers["Alice"]!.encrypt(
+            "This will fail".data(using: .utf8)!,
+            for: TestConstants.testPeerID2
+        )
+        let desyncPacket = TestHelpers.createTestPacket(
+            type: MessageType.noiseEncrypted.rawValue,
+            payload: desyncMsg
+        )
+        nodes["Bob"]!.simulateIncomingPacket(desyncPacket)
+        
+        wait(for: [expectation], timeout: TestConstants.defaultTimeout)
+        XCTAssertTrue(nackSent, "NACK should have been sent")
+        XCTAssertTrue(newHandshakeCompleted, "New handshake should have completed")
+    }
+    
     func testEndToEndSecurityScenario() throws {
         connect("Alice", "Bob")
         connect("Bob", "Charlie") // Charlie will try to eavesdrop

bitchatTests/Noise/NoiseProtocolTests.swift

@@ -454,6 +454,89 @@ final class NoiseProtocolTests: XCTestCase {
         XCTAssertNotNil(bobManager.getSession(for: TestConstants.testPeerID1))
     }
     
+    func testHandshakeAlwaysAcceptedWithExistingSession() throws {
+        // Test that handshake is always accepted even with existing valid session
+        let aliceManager = NoiseSessionManager(localStaticKey: aliceKey)
+        let bobManager = NoiseSessionManager(localStaticKey: bobKey)
+        
+        // Establish sessions
+        try establishManagerSessions(aliceManager: aliceManager, bobManager: bobManager)
+        
+        // Verify sessions are established
+        XCTAssertTrue(aliceManager.getSession(for: TestConstants.testPeerID2)?.isEstablished() ?? false)
+        XCTAssertTrue(bobManager.getSession(for: TestConstants.testPeerID1)?.isEstablished() ?? false)
+        
+        // Exchange messages to verify sessions work
+        let testMessage = "Session works".data(using: .utf8)!
+        let encrypted = try aliceManager.encrypt(testMessage, for: TestConstants.testPeerID2)
+        let decrypted = try bobManager.decrypt(encrypted, from: TestConstants.testPeerID1)
+        XCTAssertEqual(decrypted, testMessage)
+        
+        // Alice clears her session (simulating decryption failure)
+        aliceManager.removeSession(for: TestConstants.testPeerID2)
+        
+        // Alice initiates new handshake despite Bob having valid session
+        let newHandshake1 = try aliceManager.initiateHandshake(with: TestConstants.testPeerID2)
+        
+        // Bob should accept the new handshake even though he has a valid session
+        let newHandshake2 = try bobManager.handleIncomingHandshake(from: TestConstants.testPeerID1, message: newHandshake1)
+        XCTAssertNotNil(newHandshake2, "Bob should accept handshake despite having valid session")
+        
+        // Complete the handshake
+        let newHandshake3 = try aliceManager.handleIncomingHandshake(from: TestConstants.testPeerID2, message: newHandshake2!)
+        XCTAssertNotNil(newHandshake3)
+        _ = try bobManager.handleIncomingHandshake(from: TestConstants.testPeerID1, message: newHandshake3!)
+        
+        // Verify new sessions work
+        let testMessage2 = "New session works".data(using: .utf8)!
+        let encrypted2 = try aliceManager.encrypt(testMessage2, for: TestConstants.testPeerID2)
+        let decrypted2 = try bobManager.decrypt(encrypted2, from: TestConstants.testPeerID1)
+        XCTAssertEqual(decrypted2, testMessage2)
+    }
+    
+    func testNonceDesynchronizationCausesRehandshake() throws {
+        // Test that nonce desynchronization leads to proper re-handshake
+        let aliceManager = NoiseSessionManager(localStaticKey: aliceKey)
+        let bobManager = NoiseSessionManager(localStaticKey: bobKey)
+        
+        // Establish sessions
+        try establishManagerSessions(aliceManager: aliceManager, bobManager: bobManager)
+        
+        // Exchange messages normally
+        for i in 0..<5 {
+            let msg = try aliceManager.encrypt("Message \(i)".data(using: .utf8)!, for: TestConstants.testPeerID2)
+            _ = try bobManager.decrypt(msg, from: TestConstants.testPeerID1)
+        }
+        
+        // Simulate desynchronization - Alice sends messages that Bob doesn't receive
+        for i in 0..<3 {
+            _ = try aliceManager.encrypt("Lost message \(i)".data(using: .utf8)!, for: TestConstants.testPeerID2)
+        }
+        
+        // Next message from Alice should fail to decrypt at Bob (nonce mismatch)
+        let desyncMessage = try aliceManager.encrypt("This will fail".data(using: .utf8)!, for: TestConstants.testPeerID2)
+        XCTAssertThrowsError(try bobManager.decrypt(desyncMessage, from: TestConstants.testPeerID1), "Should fail due to nonce mismatch")
+        
+        // Bob clears session and initiates new handshake
+        bobManager.removeSession(for: TestConstants.testPeerID1)
+        let rehandshake1 = try bobManager.initiateHandshake(with: TestConstants.testPeerID1)
+        
+        // Alice should accept despite having a "valid" (but desynced) session
+        let rehandshake2 = try aliceManager.handleIncomingHandshake(from: TestConstants.testPeerID2, message: rehandshake1)
+        XCTAssertNotNil(rehandshake2, "Alice should accept handshake to fix desync")
+        
+        // Complete handshake
+        let rehandshake3 = try bobManager.handleIncomingHandshake(from: TestConstants.testPeerID1, message: rehandshake2!)
+        XCTAssertNotNil(rehandshake3)
+        _ = try aliceManager.handleIncomingHandshake(from: TestConstants.testPeerID2, message: rehandshake3!)
+        
+        // Verify communication works again
+        let testResynced = "Resynced".data(using: .utf8)!
+        let encryptedResync = try aliceManager.encrypt(testResynced, for: TestConstants.testPeerID2)
+        let decryptedResync = try bobManager.decrypt(encryptedResync, from: TestConstants.testPeerID1)
+        XCTAssertEqual(decryptedResync, testResynced)
+    }
+    
     // MARK: - Performance Tests
     
     func testHandshakePerformance() throws {