diff --git a/bitchat/Noise/NoiseSession.swift b/bitchat/Noise/NoiseSession.swift index 41ff8a7e..e083aa15 100644 --- a/bitchat/Noise/NoiseSession.swift +++ b/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) diff --git a/bitchatTests/Integration/IntegrationTests.swift b/bitchatTests/Integration/IntegrationTests.swift index 32a14c9f..078c8c95 100644 --- a/bitchatTests/Integration/IntegrationTests.swift +++ b/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 diff --git a/bitchatTests/Noise/NoiseProtocolTests.swift b/bitchatTests/Noise/NoiseProtocolTests.swift index 56895af9..f54275de 100644 --- a/bitchatTests/Noise/NoiseProtocolTests.swift +++ b/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 {