diff --git a/bitchatTests/Integration/IntegrationTests.swift b/bitchatTests/Integration/IntegrationTests.swift index dd129963..21c213b6 100644 --- a/bitchatTests/Integration/IntegrationTests.swift +++ b/bitchatTests/Integration/IntegrationTests.swift @@ -270,6 +270,124 @@ final class IntegrationTests: XCTestCase { XCTAssertEqual(receivedMessages.count, totalMessages) } + func testPeerPresenceTrackingAndReconnection() { + // Test peer presence tracking and identity announcement on reconnection + connect("Alice", "Bob") + + // Establish Noise sessions + do { + try establishNoiseSession("Alice", "Bob") + } catch { + XCTFail("Failed to establish Noise session: \(error)") + } + + let expectation = XCTestExpectation(description: "Peer reconnection handled") + var bobReceivedIdentityAnnounce = false + var aliceReceivedIdentityAnnounce = false + + // Track identity announcements + nodes["Bob"]!.packetDeliveryHandler = { packet in + if packet.type == MessageType.noiseIdentityAnnounce.rawValue { + bobReceivedIdentityAnnounce = true + if aliceReceivedIdentityAnnounce { + expectation.fulfill() + } + } + } + + nodes["Alice"]!.packetDeliveryHandler = { packet in + if packet.type == MessageType.noiseIdentityAnnounce.rawValue { + aliceReceivedIdentityAnnounce = true + if bobReceivedIdentityAnnounce { + expectation.fulfill() + } + } + } + + // Simulate disconnect (out of range) + disconnect("Alice", "Bob") + + // Wait to simulate extended disconnect period + Thread.sleep(forTimeInterval: 0.5) + + // Reconnect + connect("Alice", "Bob") + + // Both should receive identity announcements after reconnection + wait(for: [expectation], timeout: TestConstants.defaultTimeout) + XCTAssertTrue(bobReceivedIdentityAnnounce) + XCTAssertTrue(aliceReceivedIdentityAnnounce) + } + + func testEncryptedMessageAfterPeerRestart() { + // Test that encrypted messages work after one peer restarts + connect("Alice", "Bob") + do { + try establishNoiseSession("Alice", "Bob") + } catch { + XCTFail("Failed to establish Noise session: \(error)") + } + + // Exchange an encrypted message + let firstExpectation = XCTestExpectation(description: "First message received") + nodes["Bob"]!.messageDeliveryHandler = { message in + if message.content == "Before restart" && message.isPrivate { + firstExpectation.fulfill() + } + } + + nodes["Alice"]!.sendPrivateMessage("Before restart", to: TestConstants.testPeerID2, recipientNickname: "Bob") + wait(for: [firstExpectation], timeout: TestConstants.defaultTimeout) + + // Simulate Bob restart by recreating his Noise manager + let bobKey = Curve25519.KeyAgreement.PrivateKey() + noiseManagers["Bob"] = NoiseSessionManager(localStaticKey: bobKey) + + // Bob should initiate new handshake + let handshakeExpectation = XCTestExpectation(description: "New handshake completed") + + nodes["Bob"]!.packetDeliveryHandler = { packet in + if packet.type == MessageType.noiseHandshakeInit.rawValue { + // Bob initiates new handshake after restart + do { + let response = try self.noiseManagers["Alice"]!.handleIncomingHandshake( + from: TestConstants.testPeerID2, + message: packet.payload + ) + if let resp = response { + // Send response back to Bob + let responsePacket = TestHelpers.createTestPacket( + type: MessageType.noiseHandshakeResp.rawValue, + payload: resp + ) + self.nodes["Bob"]!.simulateIncomingPacket(responsePacket) + } + } catch { + XCTFail("Handshake handling failed: \(error)") + } + } else if packet.type == MessageType.noiseHandshakeResp.rawValue && packet.senderID.hexEncodedString() == TestConstants.testPeerID1 { + // Final handshake message (message 3 in XX pattern) + handshakeExpectation.fulfill() + } + } + + // Trigger handshake by trying to send a message + nodes["Bob"]!.sendPrivateMessage("After restart", to: TestConstants.testPeerID1, recipientNickname: "Alice") + + wait(for: [handshakeExpectation], timeout: TestConstants.defaultTimeout) + + // Now messages should work again + let secondExpectation = XCTestExpectation(description: "Message after restart received") + nodes["Alice"]!.messageDeliveryHandler = { message in + if message.content == "After restart success" && message.isPrivate { + secondExpectation.fulfill() + } + } + + nodes["Bob"]!.sendPrivateMessage("After restart success", to: TestConstants.testPeerID1, recipientNickname: "Alice") + wait(for: [secondExpectation], timeout: TestConstants.defaultTimeout) + } + func testLargeScaleNetwork() { // Create larger network for i in 5...10 { diff --git a/bitchatTests/Noise/NoiseProtocolTests.swift b/bitchatTests/Noise/NoiseProtocolTests.swift index 8b6d3b68..56895af9 100644 --- a/bitchatTests/Noise/NoiseProtocolTests.swift +++ b/bitchatTests/Noise/NoiseProtocolTests.swift @@ -295,6 +295,165 @@ final class NoiseProtocolTests: XCTestCase { XCTAssertEqual(decrypted, plaintext) } + // MARK: - Session Recovery Tests + + func testPeerRestartDetection() throws { + // Establish initial sessions + let aliceManager = NoiseSessionManager(localStaticKey: aliceKey) + let bobManager = NoiseSessionManager(localStaticKey: bobKey) + + try establishManagerSessions(aliceManager: aliceManager, bobManager: bobManager) + + // Exchange some messages to establish nonce state + let message1 = try aliceManager.encrypt("Hello".data(using: .utf8)!, for: TestConstants.testPeerID2) + _ = try bobManager.decrypt(message1, from: TestConstants.testPeerID1) + + let message2 = try bobManager.encrypt("World".data(using: .utf8)!, for: TestConstants.testPeerID1) + _ = try aliceManager.decrypt(message2, from: TestConstants.testPeerID2) + + // Simulate Bob restart by creating new manager with same key + let bobManagerRestarted = NoiseSessionManager(localStaticKey: bobKey) + + // Bob initiates new handshake after restart + let newHandshake1 = try bobManagerRestarted.initiateHandshake(with: TestConstants.testPeerID1) + + // Alice should accept the new handshake (clearing old session) + let newHandshake2 = try aliceManager.handleIncomingHandshake(from: TestConstants.testPeerID2, message: newHandshake1) + XCTAssertNotNil(newHandshake2) + + // Complete the new handshake + let newHandshake3 = try bobManagerRestarted.handleIncomingHandshake(from: TestConstants.testPeerID1, message: newHandshake2!) + XCTAssertNotNil(newHandshake3) + _ = try aliceManager.handleIncomingHandshake(from: TestConstants.testPeerID2, message: newHandshake3!) + + // Should be able to exchange messages with new sessions + let testMessage = "After restart".data(using: .utf8)! + let encrypted = try bobManagerRestarted.encrypt(testMessage, for: TestConstants.testPeerID1) + let decrypted = try aliceManager.decrypt(encrypted, from: TestConstants.testPeerID2) + XCTAssertEqual(decrypted, testMessage) + } + + func testNonceDesynchronizationRecovery() throws { + // Create two sessions + aliceSession = NoiseSession(peerID: TestConstants.testPeerID2, role: .initiator, localStaticKey: aliceKey) + bobSession = NoiseSession(peerID: TestConstants.testPeerID1, role: .responder, localStaticKey: bobKey) + + // Establish sessions + try performHandshake(initiator: aliceSession, responder: bobSession) + + // Exchange messages to advance nonces + for i in 0..<5 { + let msg = try aliceSession.encrypt("Message \(i)".data(using: .utf8)!) + _ = try bobSession.decrypt(msg) + } + + // Simulate desynchronization by encrypting but not decrypting + for i in 0..<3 { + _ = try aliceSession.encrypt("Lost message \(i)".data(using: .utf8)!) + } + + // Next message from Alice should fail to decrypt (nonce mismatch) + let desyncMessage = try aliceSession.encrypt("This will fail".data(using: .utf8)!) + XCTAssertThrowsError(try bobSession.decrypt(desyncMessage)) + } + + func testConcurrentEncryption() throws { + // Test thread safety of encryption operations + let aliceManager = NoiseSessionManager(localStaticKey: aliceKey) + let bobManager = NoiseSessionManager(localStaticKey: bobKey) + + try establishManagerSessions(aliceManager: aliceManager, bobManager: bobManager) + + let messageCount = 100 + let expectation = XCTestExpectation(description: "All messages encrypted and decrypted") + expectation.expectedFulfillmentCount = messageCount * 2 + + let group = DispatchGroup() + var encryptedMessages: [Int: Data] = [:] + let encryptionQueue = DispatchQueue(label: "test.encryption", attributes: .concurrent) + let lock = NSLock() + + // Encrypt messages concurrently + for i in 0..