Negotiate with both peers when using multihop and pq ios 746

ios/MullvadRustRuntime/PostQuantumKeyExchangeActor.swift

@@ -12,7 +12,13 @@ import MullvadTypes
 import NetworkExtension
 import WireGuardKitTypes
 
-public class PostQuantumKeyExchangeActor {
+public protocol PostQuantumKeyExchangeActorProtocol {
+    func startNegotiation(with privateKey: PrivateKey)
+    func endCurrentNegotiation()
+    func reset()
+}
+
+public class PostQuantumKeyExchangeActor: PostQuantumKeyExchangeActorProtocol {
     struct Negotiation {
         var negotiator: PostQuantumKeyNegotiating
         var inTunnelTCPConnection: NWTCPConnection
@@ -69,7 +75,7 @@ public class PostQuantumKeyExchangeActor {
         endCurrentNegotiation()
         let negotiator = negotiationProvider.init()
 
-        let gatewayAddress = "10.64.0.1"
+        let gatewayAddress = LocalNetworkIPs.gatewayAddress.rawValue
         let IPv4Gateway = IPv4Address(gatewayAddress)!
         let endpoint = NWHostEndpoint(hostname: gatewayAddress, port: "\(CONFIG_SERVICE_PORT)")
         let inTunnelTCPConnection = createTCPConnection(endpoint)
@@ -97,6 +103,8 @@ public class PostQuantumKeyExchangeActor {
                 tcpConnection: inTunnelTCPConnection,
                 postQuantumKeyExchangeTimeout: tcpConnectionTimeout
             ) {
+                // Cancel the negotiation to shut down any remaining use of the TCP connection on the Rust side
+                self.negotiation?.cancel()
                 self.negotiation = nil
                 self.onFailure()
             }

ios/MullvadRustRuntimeTests/MullvadPostQuantum+Stubs.swift

@@ -58,7 +58,7 @@ class FailedNegotiatorStub: PostQuantumKeyNegotiating {
         gatewayIP: IPv4Address,
         devicePublicKey: WireGuardKitTypes.PublicKey,
         presharedKey: WireGuardKitTypes.PrivateKey,
-        packetTunnel: PacketTunnelCore.TunnelProvider,
+        postQuantumKeyReceiver packetTunnel: any MullvadTypes.TunnelProvider,
         tcpConnection: NWTCPConnection,
         postQuantumKeyExchangeTimeout: MullvadTypes.Duration
     ) -> Bool { false }
@@ -82,7 +82,7 @@ class SuccessfulNegotiatorStub: PostQuantumKeyNegotiating {
         gatewayIP: IPv4Address,
         devicePublicKey: WireGuardKitTypes.PublicKey,
         presharedKey: WireGuardKitTypes.PrivateKey,
-        packetTunnel: PacketTunnelCore.TunnelProvider,
+        postQuantumKeyReceiver packetTunnel: any MullvadTypes.TunnelProvider,
         tcpConnection: NWTCPConnection,
         postQuantumKeyExchangeTimeout: MullvadTypes.Duration
     ) -> Bool { true }

ios/MullvadRustRuntimeTests/PostQuantumKeyExchangeActorTests.swift → ios/MullvadRustRuntimeTests/MullvadPostQuantumTests.swift

@@ -1,6 +1,6 @@
 //
-//  PostQuantumKeyExchangeActorTests.swift
-//  PostQuantumKeyExchangeActorTests
+//  MullvadPostQuantumTests.swift
+//  MullvadPostQuantumTests
 //
 //  Created by Marco Nikic on 2024-06-12.
 //  Copyright © 2024 Mullvad VPN AB. All rights reserved.
@@ -71,7 +71,7 @@ class MullvadPostQuantumTests: XCTestCase {
                 unexpectedNegotiationFailure.fulfill()
             },
             negotiationProvider: SuccessfulNegotiatorStub.self,
-            iteratorProvider: { AnyIterator { .milliseconds(10) } }
+            iteratorProvider: { AnyIterator { .seconds(1) } }
         )
 
         let privateKey = PrivateKey()

ios/MullvadTypes/LocalNetworkIPs.swift

@@ -0,0 +1,15 @@
+//
+//  LocalNetworkIPs.swift
+//  MullvadTypes
+//
+//  Created by Mojgan on 2024-07-26.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+
+public enum LocalNetworkIPs: String {
+    case gatewayAddress = "10.64.0.1"
+    case defaultRouteIpV4 = "0.0.0.0"
+    case defaultRouteIpV6 = "::"
+}

ios/MullvadTypes/Protocols/PostQuantumKeyReceiving.swift

@@ -9,7 +9,7 @@
 import Foundation
 import WireGuardKitTypes
 
-public protocol PostQuantumKeyReceiving: AnyObject {
+public protocol PostQuantumKeyReceiving {
     func receivePostQuantumKey(_ key: PreSharedKey, ephemeralKey: PrivateKey)
     func keyExchangeFailed()
 }

ios/PacketTunnel/PacketTunnelProvider/PacketTunnelProvider.swift

@@ -21,12 +21,12 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
     private let providerLogger: Logger
 
     private var actor: PacketTunnelActor!
-    private var postQuantumActor: PostQuantumKeyExchangeActor!
     private var appMessageHandler: AppMessageHandler!
     private var stateObserverTask: AnyTask?
     private var deviceChecker: DeviceChecker!
     private var adapter: WgAdapter!
     private var relaySelector: RelaySelectorWrapper!
+    private var postQuantumKeyExchangingPipeline: PostQuantumKeyExchangingPipeline!
 
     private let multihopStateListener = MultihopStateListener()
     private let multihopUpdater: MultihopUpdater
@@ -94,14 +94,21 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
             protocolObfuscator: ProtocolObfuscator<UDPOverTCPObfuscator>()
         )
 
-        postQuantumActor = PostQuantumKeyExchangeActor(
-            packetTunnel: postQuantumReceiver,
-            onFailure: self.keyExchangeFailed,
-            iteratorProvider: { REST.RetryStrategy.postQuantumKeyExchange.makeDelayIterator() }
-        )
-
         let urlRequestProxy = URLRequestProxy(dispatchQueue: internalQueue, transportProvider: transportProvider)
         appMessageHandler = AppMessageHandler(packetTunnelActor: actor, urlRequestProxy: urlRequestProxy)
+
+        postQuantumKeyExchangingPipeline = PostQuantumKeyExchangingPipeline(
+            PostQuantumKeyExchangeActor(
+                packetTunnel: postQuantumReceiver,
+                onFailure: self.keyExchangeFailed,
+                iteratorProvider: { REST.RetryStrategy.postQuantumKeyExchange.makeDelayIterator() }
+            ),
+            onUpdateConfiguration: { [unowned self] configuration in
+                actor.replaceKeyWithPQ(configuration: configuration)
+            }, onFinish: { [unowned self] in
+                actor.notifyPostQuantumKeyExchanged()
+            }
+        )
     }
 
     override func startTunnel(options: [String: NSObject]? = nil) async throws {
@@ -251,8 +258,8 @@ extension PacketTunnelProvider {
                 }
 
                 switch newState {
-                case let .reconnecting(connState), let .connecting(connState):
-                    let connectionAttempt = connState.connectionAttemptCount
+                case let .reconnecting(observedConnectionState), let .connecting(observedConnectionState):
+                    let connectionAttempt = observedConnectionState.connectionAttemptCount
 
                     // Start device check every second failure attempt to connect.
                     if lastConnectionAttempt != connectionAttempt, connectionAttempt > 0,
@@ -263,9 +270,8 @@ extension PacketTunnelProvider {
                     // Cache last connection attempt to filter out repeating calls.
                     lastConnectionAttempt = connectionAttempt
 
-                case let .negotiatingPostQuantumKey(_, privateKey):
-                    postQuantumActor.startNegotiation(with: privateKey)
-
+                case let .negotiatingPostQuantumKey(observedConnectionState, privateKey):
+                    postQuantumKeyExchangingPipeline.startNegotiation(observedConnectionState, privateKey: privateKey)
                 case .initial, .connected, .disconnecting, .disconnected, .error:
                     break
                 }
@@ -310,8 +316,7 @@ extension PacketTunnelProvider {
 
 extension PacketTunnelProvider: PostQuantumKeyReceiving {
     func receivePostQuantumKey(_ key: PreSharedKey, ephemeralKey: PrivateKey) {
-        postQuantumActor.reset()
-        actor.replacePreSharedKey(key, ephemeralKey: ephemeralKey)
+        postQuantumKeyExchangingPipeline.receivePostQuantumKey(key, ephemeralKey: ephemeralKey)
     }
 
     func keyExchangeFailed() {

ios/PacketTunnel/PostQuantum/MultiHopPostQuantumKeyExchanging.swift

@@ -0,0 +1,132 @@
+//
+//  MultiHopPostQuantumKeyExchanging.swift
+//  PacketTunnel
+//
+//  Created by Mojgan on 2024-07-15.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import MullvadREST
+import MullvadRustRuntime
+import MullvadSettings
+import MullvadTypes
+import PacketTunnelCore
+import WireGuardKitTypes
+
+final class MultiHopPostQuantumKeyExchanging: PostQuantumKeyExchangingProtocol {
+    let entry: SelectedRelay
+    let exit: SelectedRelay
+    let keyExchanger: PostQuantumKeyExchangeActorProtocol
+    let devicePrivateKey: PrivateKey
+    let onFinish: () -> Void
+    let onUpdateConfiguration: (PostQuantumNegotiationState) -> Void
+
+    private var entryPostQuantumKey: PostQuantumKey!
+    private var exitPostQuantumKey: PostQuantumKey!
+
+    private let defaultGatewayAddressRange = [IPAddressRange(from: "\(LocalNetworkIPs.gatewayAddress.rawValue)/32")!]
+    private let allTrafficRange = [
+        IPAddressRange(from: "\(LocalNetworkIPs.defaultRouteIpV4.rawValue)/0")!,
+        IPAddressRange(from: "\(LocalNetworkIPs.defaultRouteIpV6.rawValue)/0")!,
+    ]
+
+    private var state: StateMachine = .initial
+
+    enum StateMachine {
+        case initial
+        case negotiatingWithEntry
+        case negotiatingBetweenEntryAndExit
+        case makeConnection
+    }
+
+    init(
+        entry: SelectedRelay,
+        exit: SelectedRelay,
+        devicePrivateKey: PrivateKey,
+        keyExchanger: PostQuantumKeyExchangeActorProtocol,
+        onUpdateConfiguration: @escaping (PostQuantumNegotiationState) -> Void,
+        onFinish: @escaping () -> Void
+    ) {
+        self.entry = entry
+        self.exit = exit
+        self.devicePrivateKey = devicePrivateKey
+        self.keyExchanger = keyExchanger
+        self.onUpdateConfiguration = onUpdateConfiguration
+        self.onFinish = onFinish
+    }
+
+    func start() {
+        guard state == .initial else { return }
+        negotiateWithEntry()
+    }
+
+    func receivePostQuantumKey(
+        _ preSharedKey: PreSharedKey,
+        ephemeralKey: PrivateKey
+    ) {
+        if state == .negotiatingWithEntry {
+            entryPostQuantumKey = PostQuantumKey(preSharedKey: preSharedKey, ephemeralKey: ephemeralKey)
+            negotiateBetweenEntryAndExit()
+        } else if state == .negotiatingBetweenEntryAndExit {
+            exitPostQuantumKey = PostQuantumKey(preSharedKey: preSharedKey, ephemeralKey: ephemeralKey)
+            makeConnection()
+        }
+    }
+
+    private func negotiateWithEntry() {
+        state = .negotiatingWithEntry
+        onUpdateConfiguration(.single(PostQuantumConfigurationRelay(
+            relay: entry,
+            configuration: PostQuantumConfiguration(
+                privateKey: devicePrivateKey,
+                allowedIPs: defaultGatewayAddressRange
+            )
+        )))
+        keyExchanger.startNegotiation(with: devicePrivateKey)
+    }
+
+    private func negotiateBetweenEntryAndExit() {
+        state = .negotiatingBetweenEntryAndExit
+        onUpdateConfiguration(.multi(
+            entry: PostQuantumConfigurationRelay(
+                relay: entry,
+                configuration: PostQuantumConfiguration(
+                    privateKey: entryPostQuantumKey.ephemeralKey,
+                    preSharedKey: entryPostQuantumKey.preSharedKey,
+                    allowedIPs: [IPAddressRange(from: "\(exit.endpoint.ipv4Relay.ip)/32")!]
+                )
+            ),
+            exit: PostQuantumConfigurationRelay(
+                relay: exit,
+                configuration: PostQuantumConfiguration(
+                    privateKey: devicePrivateKey,
+                    allowedIPs: defaultGatewayAddressRange
+                )
+            )
+        ))
+        keyExchanger.startNegotiation(with: devicePrivateKey)
+    }
+
+    private func makeConnection() {
+        state = .makeConnection
+        onUpdateConfiguration(.multi(
+            entry: PostQuantumConfigurationRelay(
+                relay: entry,
+                configuration: PostQuantumConfiguration(
+                    privateKey: entryPostQuantumKey.ephemeralKey,
+                    preSharedKey: entryPostQuantumKey.preSharedKey,
+                    allowedIPs: [IPAddressRange(from: "\(exit.endpoint.ipv4Relay.ip)/32")!]
+                )
+            ),
+            exit: PostQuantumConfigurationRelay(
+                relay: exit,
+                configuration: PostQuantumConfiguration(
+                    privateKey: exitPostQuantumKey.ephemeralKey,
+                    preSharedKey: exitPostQuantumKey.preSharedKey,
+                    allowedIPs: allTrafficRange
+                )
+            )
+        ))
+        self.onFinish()
+    }
+}