Add timeout to the tcp connection for the pq key negotiation ios 701

ios/MullvadPostQuantum/PostQuantumKeyNegotiator.swift

@@ -7,26 +7,46 @@
 //
 
 import Foundation
+import MullvadTypes
 import NetworkExtension
+import PacketTunnelCore
 import TalpidTunnelConfigClientProxy
 import WireGuardKitTypes
 
+// swiftlint:disable function_parameter_count
+public protocol PostQuantumKeyNegotiating {
+    func startNegotiation(
+        gatewayIP: IPv4Address,
+        devicePublicKey: PublicKey,
+        presharedKey: PrivateKey,
+        packetTunnel: any TunnelProvider,
+        tcpConnection: NWTCPConnection,
+        postQuantumKeyExchangeTimeout: Duration
+    ) -> Bool
+
+    func cancelKeyNegotiation()
+
+    init()
+}
+
 /**
  Attempt to start the asynchronous process of key negotiation. Returns true if successfully started, false if failed.
  */
-public class PostQuantumKeyNegotiator {
-    public init() {}
+public class PostQuantumKeyNegotiator: PostQuantumKeyNegotiating {
+    required public init() {}
 
     var cancelToken: PostQuantumCancelToken?
 
     public func startNegotiation(
         gatewayIP: IPv4Address,
         devicePublicKey: PublicKey,
         presharedKey: PrivateKey,
-        packetTunnel: NEPacketTunnelProvider,
-        tcpConnection: NWTCPConnection
+        packetTunnel: any TunnelProvider,
+        tcpConnection: NWTCPConnection,
+        postQuantumKeyExchangeTimeout: Duration
     ) -> Bool {
-        let packetTunnelPointer = Unmanaged.passUnretained(packetTunnel).toOpaque()
+        // swiftlint:disable:next force_cast
+        let packetTunnelPointer = Unmanaged.passUnretained(packetTunnel as! NEPacketTunnelProvider).toOpaque()
         let opaqueConnection = Unmanaged.passUnretained(tcpConnection).toOpaque()
         var cancelToken = PostQuantumCancelToken()
 
@@ -35,7 +55,8 @@ public class PostQuantumKeyNegotiator {
             presharedKey.rawValue.map { $0 },
             packetTunnelPointer,
             opaqueConnection,
-            &cancelToken
+            &cancelToken,
+            UInt64(postQuantumKeyExchangeTimeout.timeInterval)
         )
         guard result == 0 else {
             return false
@@ -54,3 +75,5 @@ public class PostQuantumKeyNegotiator {
         drop_post_quantum_key_exchange_token(&cancelToken)
     }
 }
+
+// swiftlint:enable function_parameter_count

ios/MullvadPostQuantum/talpid-tunnel-config-client/include/talpid_tunnel_config_client.h

@@ -16,15 +16,18 @@ typedef struct PostQuantumCancelToken {
  * Called by the Swift side to signal that the quantum-secure key exchange should be cancelled.
  *
  * # Safety
- * `sender` must be pointing to a valid instance of a `PostQuantumCancelToken` created by the `PacketTunnelProvider`.
+ * `sender` must be pointing to a valid instance of a `PostQuantumCancelToken` created by the
+ * `PacketTunnelProvider`.
  */
 void cancel_post_quantum_key_exchange(const struct PostQuantumCancelToken *sender);
 
 /**
- * Called by the Swift side to signal that the Rust `PostQuantumCancelToken` can be safely dropped from memory.
+ * Called by the Swift side to signal that the Rust `PostQuantumCancelToken` can be safely dropped
+ * from memory.
  *
  * # Safety
- * `sender` must be pointing to a valid instance of a `PostQuantumCancelToken` created by the `PacketTunnelProvider`.
+ * `sender` must be pointing to a valid instance of a `PostQuantumCancelToken` created by the
+ * `PacketTunnelProvider`.
  */
 void drop_post_quantum_key_exchange_token(const struct PostQuantumCancelToken *sender);
 
@@ -44,33 +47,32 @@ void handle_sent(uintptr_t bytes_sent, const void *sender);
  * Called by Swift whenever data has been read from the in-tunnel TCP connection when exchanging
  * quantum-resistant pre shared keys.
  *
- * If `data` is null or empty, this indicates that the connection was closed or that an error occurred.
- * An empty buffer is sent to the underlying reader to signal EOF.
+ * If `data` is null or empty, this indicates that the connection was closed or that an error
+ * occurred. An empty buffer is sent to the underlying reader to signal EOF.
  *
  * # Safety
  * `sender` must be pointing to a valid instance of a `read_tx` created by the `IosTcpProvider`
  *
  * Callback to call when the TCP connection has received data.
  */
-void handle_recv(const uint8_t *data,
-                 uintptr_t data_len,
-                 const void *sender);
+void handle_recv(const uint8_t *data, uintptr_t data_len, const void *sender);
 
 /**
  * Entry point for exchanging post quantum keys on iOS.
  * The TCP connection must be created to go through the tunnel.
  * # Safety
  * `public_key` and `ephemeral_key` must be valid respective `PublicKey` and `PrivateKey` types.
  * They will not be valid after this function is called, and thus must be copied here.
- * `packet_tunnel` and `tcp_connection` must be valid pointers to a packet tunnel and a TCP connection
- * instances.
+ * `packet_tunnel` and `tcp_connection` must be valid pointers to a packet tunnel and a TCP
+ * connection instances.
  * `cancel_token` should be owned by the caller of this function.
  */
 int32_t negotiate_post_quantum_key(const uint8_t *public_key,
                                    const uint8_t *ephemeral_key,
                                    const void *packet_tunnel,
                                    const void *tcp_connection,
-                                   struct PostQuantumCancelToken *cancel_token);
+                                   struct PostQuantumCancelToken *cancel_token,
+                                   uint64_t post_quantum_key_exchange_timeout);
 
 /**
  * Called when there is data to send on the TCP connection.

ios/MullvadPostQuantumTests/MullvadPostQuantum+Stubs.swift

@@ -0,0 +1,95 @@
+//
+//  MullvadPostQuantum+Stubs.swift
+//  MullvadPostQuantumTests
+//
+//  Created by Marco Nikic on 2024-06-12.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+@testable import MullvadPostQuantum
+@testable import MullvadTypes
+import NetworkExtension
+@testable import PacketTunnelCore
+@testable import WireGuardKitTypes
+
+// swiftlint:disable function_parameter_count
+class NWTCPConnectionStub: NWTCPConnection {
+    var _isViable = false
+    override var isViable: Bool {
+        _isViable
+    }
+
+    func becomeViable() {
+        willChangeValue(for: \.isViable)
+        _isViable = true
+        didChangeValue(for: \.isViable)
+    }
+}
+
+class TunnelProviderStub: TunnelProvider {
+    let tcpConnection: NWTCPConnectionStub
+
+    init(tcpConnection: NWTCPConnectionStub) {
+        self.tcpConnection = tcpConnection
+    }
+
+    func createTCPConnectionThroughTunnel(
+        to remoteEndpoint: NWEndpoint,
+        enableTLS: Bool,
+        tlsParameters TLSParameters: NWTLSParameters?,
+        delegate: Any?
+    ) -> NWTCPConnection {
+        tcpConnection
+    }
+}
+
+class FailedNegotiatorStub: PostQuantumKeyNegotiating {
+    var onCancelKeyNegotiation: (() -> Void)?
+
+    required init() {
+        onCancelKeyNegotiation = nil
+    }
+
+    init(onCancelKeyNegotiation: (() -> Void)? = nil) {
+        self.onCancelKeyNegotiation = onCancelKeyNegotiation
+    }
+
+    func startNegotiation(
+        gatewayIP: IPv4Address,
+        devicePublicKey: WireGuardKitTypes.PublicKey,
+        presharedKey: WireGuardKitTypes.PrivateKey,
+        packetTunnel: PacketTunnelCore.TunnelProvider,
+        tcpConnection: NWTCPConnection,
+        postQuantumKeyExchangeTimeout: MullvadTypes.Duration
+    ) -> Bool { false }
+
+    func cancelKeyNegotiation() {
+        onCancelKeyNegotiation?()
+    }
+}
+
+class SuccessfulNegotiatorStub: PostQuantumKeyNegotiating {
+    var onCancelKeyNegotiation: (() -> Void)?
+    required init() {
+        onCancelKeyNegotiation = nil
+    }
+
+    init(onCancelKeyNegotiation: (() -> Void)? = nil) {
+        self.onCancelKeyNegotiation = onCancelKeyNegotiation
+    }
+
+    func startNegotiation(
+        gatewayIP: IPv4Address,
+        devicePublicKey: WireGuardKitTypes.PublicKey,
+        presharedKey: WireGuardKitTypes.PrivateKey,
+        packetTunnel: PacketTunnelCore.TunnelProvider,
+        tcpConnection: NWTCPConnection,
+        postQuantumKeyExchangeTimeout: MullvadTypes.Duration
+    ) -> Bool { true }
+
+    func cancelKeyNegotiation() {
+        onCancelKeyNegotiation?()
+    }
+}
+
+// swiftlint:enable function_parameter_count

ios/MullvadPostQuantumTests/PostQuantumKeyExchangeActorTests.swift

@@ -0,0 +1,91 @@
+//
+//  PostQuantumKeyExchangeActorTests.swift
+//  PostQuantumKeyExchangeActorTests
+//
+//  Created by Marco Nikic on 2024-06-12.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+@testable import MullvadPostQuantum
+@testable import MullvadTypes
+import NetworkExtension
+@testable import PacketTunnelCore
+@testable import WireGuardKitTypes
+import XCTest
+
+class MullvadPostQuantumTests: XCTestCase {
+    var tcpConnection: NWTCPConnectionStub!
+    var tunnelProvider: TunnelProviderStub!
+
+    override func setUpWithError() throws {
+        tcpConnection = NWTCPConnectionStub()
+        tunnelProvider = TunnelProviderStub(tcpConnection: tcpConnection)
+    }
+
+    func testKeyExchangeFailsWhenNegotiationCannotStart() {
+        let negotiationFailure = expectation(description: "Negotiation failed")
+
+        let keyExchangeActor = PostQuantumKeyExchangeActor(
+            packetTunnel: tunnelProvider,
+            onFailure: {
+                negotiationFailure.fulfill()
+            },
+            negotiationProvider: FailedNegotiatorStub.self
+        )
+
+        let privateKey = PrivateKey()
+        keyExchangeActor.startNegotiation(with: privateKey)
+        tcpConnection.becomeViable()
+
+        wait(for: [negotiationFailure])
+    }
+
+    func testKeyExchangeFailsWhenTCPConnectionIsNotReadyInTime() {
+        let negotiationFailure = expectation(description: "Negotiation failed")
+
+        // Setup the actor to wait only 10 milliseconds before declaring the TCP connection is not ready in time.
+        let keyExchangeActor = PostQuantumKeyExchangeActor(
+            packetTunnel: tunnelProvider,
+            onFailure: {
+                negotiationFailure.fulfill()
+            },
+            negotiationProvider: FailedNegotiatorStub.self,
+            keyExchangeRetriesIterator: AnyIterator { .milliseconds(10) }
+        )
+
+        let privateKey = PrivateKey()
+        keyExchangeActor.startNegotiation(with: privateKey)
+
+        wait(for: [negotiationFailure])
+    }
+
+    func testResetEndsTheCurrentNegotiation() throws {
+        let unexpectedNegotiationFailure = expectation(description: "Unexpected negotiation failure")
+        unexpectedNegotiationFailure.isInverted = true
+
+        let keyExchangeActor = PostQuantumKeyExchangeActor(
+            packetTunnel: tunnelProvider,
+            onFailure: {
+                unexpectedNegotiationFailure.fulfill()
+            },
+            negotiationProvider: SuccessfulNegotiatorStub.self
+        )
+
+        let privateKey = PrivateKey()
+        keyExchangeActor.startNegotiation(with: privateKey)
+
+        let negotiationProvider = try XCTUnwrap(
+            keyExchangeActor.negotiation?
+                .negotiator as? SuccessfulNegotiatorStub
+        )
+
+        let negotationCancelledExpectation = expectation(description: "Negotiation cancelled")
+        negotiationProvider.onCancelKeyNegotiation = {
+            negotationCancelledExpectation.fulfill()
+        }
+
+        keyExchangeActor.reset()
+
+        wait(for: [negotationCancelledExpectation, unexpectedNegotiationFailure], timeout: 0.5)
+    }
+}

ios/MullvadREST/RetryStrategy/RetryStrategy.swift

@@ -68,6 +68,16 @@ extension REST {
             multiplier: 2,
             maxDelay: .seconds(8)
         )
+
+        public static var postQuantumKeyExchange = RetryStrategy(
+            maxRetryCount: 10,
+            delay: .exponentialBackoff(
+                initial: .seconds(10),
+                multiplier: UInt64(2),
+                maxDelay: .seconds(30)
+            ),
+            applyJitter: true
+        )
     }
 
     public enum RetryDelay: Equatable {