mirror of
https://github.com/permissionlesstech/bitchat.git
synced 2026-07-25 01:05:19 +00:00
Fix NoiseSessionManager to always accept handshake initiations
Previously, NoiseSessionManager would reject handshake initiations if it had an existing established session. This caused deadlocks when one peer cleared their session (e.g., after decryption failure) but the other peer rejected the new handshake. Changes: - NoiseSessionManager now always accepts handshake initiations, clearing any existing session - Added comprehensive tests for handshake recovery scenarios - Tests verify proper re-establishment after decryption failures and nonce desynchronization
This commit is contained in:
jack
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user