Merge branch 'daita-is-not-actually-used-when-multihop-is-used-ios-1111'

2 files changed, 47 insertions, 5 deletions

diff --git a/ios/PacketTunnel/PostQuantum/MultiHopEphemeralPeerExchanger.swift b/ios/PacketTunnel/PostQuantum/MultiHopEphemeralPeerExchanger.swift
index e3be56e267..3bf3d34fb0 100644
--- a/ios/PacketTunnel/PostQuantum/MultiHopEphemeralPeerExchanger.swift
+++ b/ios/PacketTunnel/PostQuantum/MultiHopEphemeralPeerExchanger.swift
@@ -71,8 +71,8 @@ final class MultiHopEphemeralPeerExchanger: EphemeralPeerExchangingProtocol {
         _ ephemeralPeerPrivateKey: PrivateKey,
         daitaParameters: DaitaV2Parameters?
     ) async {
-        self.daitaParameters = daitaParameters
         if state == .negotiatingWithEntry {
+            self.daitaParameters = daitaParameters
             entryPeerKey = EphemeralPeerKey(ephemeralKey: ephemeralPeerPrivateKey)
             await negotiateBetweenEntryAndExit()
         } else if state == .negotiatingBetweenEntryAndExit {
@@ -86,8 +86,8 @@ final class MultiHopEphemeralPeerExchanger: EphemeralPeerExchangingProtocol {
         ephemeralKey: PrivateKey,
         daitaParameters: DaitaV2Parameters?
     ) async {
-        self.daitaParameters = daitaParameters
         if state == .negotiatingWithEntry {
+            self.daitaParameters = daitaParameters
             entryPeerKey = EphemeralPeerKey(preSharedKey: preSharedKey, ephemeralKey: ephemeralKey)
             await negotiateBetweenEntryAndExit()
         } else if state == .negotiatingBetweenEntryAndExit {
@@ -130,7 +130,7 @@ final class MultiHopEphemeralPeerExchanger: EphemeralPeerExchangingProtocol {
                 configuration: EphemeralPeerConfiguration(
                     privateKey: devicePrivateKey,
                     allowedIPs: defaultGatewayAddressRange,
-                    daitaParameters: self.daitaParameters
+                    daitaParameters: nil
                 )
             )
         ))
@@ -160,7 +160,7 @@ final class MultiHopEphemeralPeerExchanger: EphemeralPeerExchangingProtocol {
                     privateKey: exitPeerKey.ephemeralKey,
                     preSharedKey: exitPeerKey.preSharedKey,
                     allowedIPs: allTrafficRange,
-                    daitaParameters: self.daitaParameters
+                    daitaParameters: nil
                 )
             )
         ))
diff --git a/ios/PacketTunnelCoreTests/MultiHopEphemeralPeerExchangerTests.swift b/ios/PacketTunnelCoreTests/MultiHopEphemeralPeerExchangerTests.swift
index 90d05566c2..4e77a6b617 100644
--- a/ios/PacketTunnelCoreTests/MultiHopEphemeralPeerExchangerTests.swift
+++ b/ios/PacketTunnelCoreTests/MultiHopEphemeralPeerExchangerTests.swift
@@ -160,7 +160,49 @@ final class MultiHopEphemeralPeerExchangerTests: XCTestCase {
             keyExchanger: peerExchangeActor,
             enablePostQuantum: false,
             enableDaita: true
-        ) { _ in
+        ) { params in
+            reconfigurationExpectation.fulfill()
+        } onFinish: {
+            negotiationSuccessful.fulfill()
+        }
+
+        peerExchangeActor.delegate = KeyExchangingResultStub(onReceiveEphemeralPeerPrivateKey: { ephemeralKey, daita in
+            await multiHopPeerExchanger.receiveEphemeralPeerPrivateKey(ephemeralKey, daitaParameters: daita)
+        })
+        await multiHopPeerExchanger.start()
+
+        wait(
+            for: [unexpectedNegotiationFailure, reconfigurationExpectation, negotiationSuccessful],
+            timeout: .UnitTest.invertedTimeout
+        )
+    }
+
+    func testEphemeralPeerExchangeSuccessPassesDaitaParameters() async throws {
+        let unexpectedNegotiationFailure = expectation(description: "Negotiation failed.")
+        unexpectedNegotiationFailure.isInverted = true
+
+        let reconfigurationExpectation = expectation(description: "Tunnel reconfiguration took place")
+        reconfigurationExpectation.expectedFulfillmentCount = 3
+
+        let negotiationSuccessful = expectation(description: "Negotiation succeeded.")
+        negotiationSuccessful.expectedFulfillmentCount = 1
+
+        let peerExchangeActor = EphemeralPeerExchangeActorStub()
+        let preSharedKey = try XCTUnwrap(PreSharedKey(hexKey: PrivateKey().hexKey))
+        peerExchangeActor.result = .success((preSharedKey, PrivateKey()))
+
+        let multiHopPeerExchanger = MultiHopEphemeralPeerExchanger(
+            entry: entryRelay,
+            exit: exitRelay,
+            devicePrivateKey: PrivateKey(),
+            keyExchanger: peerExchangeActor,
+            enablePostQuantum: false,
+            enableDaita: true
+        ) { params in
+            if case let .multi(entry, exit) = params {
+                XCTAssertNotNil(entry.configuration.daitaParameters)
+                XCTAssertNil(exit.configuration.daitaParameters)
+            }
             reconfigurationExpectation.fulfill()
         } onFinish: {
             negotiationSuccessful.fulfill()