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
2025-07-23 19:03:41 +02:00
parent a84d6f22ef
commit 2759202616
3 changed files with 252 additions and 5 deletions
+7 -5
View File
@@ -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)
@@ -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 {