Add a backing off timeout when negotiating PQ PSK

5 files changed, 56 insertions, 22 deletions

diff --git a/ios/MullvadPostQuantum/PostQuantumKeyNegotiator.swift b/ios/MullvadPostQuantum/PostQuantumKeyNegotiator.swift
index 16f7a24931..ebc494cf0f 100644
--- a/ios/MullvadPostQuantum/PostQuantumKeyNegotiator.swift
+++ b/ios/MullvadPostQuantum/PostQuantumKeyNegotiator.swift
@@ -7,6 +7,7 @@
 //
 
 import Foundation
+import MullvadTypes
 import NetworkExtension
 import TalpidTunnelConfigClientProxy
 import WireGuardKitTypes
@@ -24,7 +25,8 @@ public class PostQuantumKeyNegotiator {
         devicePublicKey: PublicKey,
         presharedKey: PrivateKey,
         packetTunnel: NEPacketTunnelProvider,
-        tcpConnection: NWTCPConnection
+        tcpConnection: NWTCPConnection,
+        postQuantumKeyExchangeTimeout: Duration
     ) -> Bool {
         let packetTunnelPointer = Unmanaged.passUnretained(packetTunnel).toOpaque()
         let opaqueConnection = Unmanaged.passUnretained(tcpConnection).toOpaque()
@@ -35,7 +37,8 @@ public class PostQuantumKeyNegotiator {
             presharedKey.rawValue.map { $0 },
             packetTunnelPointer,
             opaqueConnection,
-            &cancelToken
+            &cancelToken,
+            UInt64(postQuantumKeyExchangeTimeout.timeInterval)
         )
         guard result == 0 else {
             return false
diff --git a/ios/MullvadPostQuantum/talpid-tunnel-config-client/include/talpid_tunnel_config_client.h b/ios/MullvadPostQuantum/talpid-tunnel-config-client/include/talpid_tunnel_config_client.h
index 04180db289..31d31748ae 100644
--- a/ios/MullvadPostQuantum/talpid-tunnel-config-client/include/talpid_tunnel_config_client.h
+++ b/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,17 +47,15 @@ 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.
@@ -62,15 +63,16 @@ void handle_recv(const uint8_t *data,
  * # 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.
diff --git a/ios/MullvadREST/RetryStrategy/RetryStrategy.swift b/ios/MullvadREST/RetryStrategy/RetryStrategy.swift
index 18e3cd69f3..82e31abd2b 100644
--- a/ios/MullvadREST/RetryStrategy/RetryStrategy.swift
+++ b/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 {
diff --git a/ios/PacketTunnel/PacketTunnelProvider/PacketTunnelProvider.swift b/ios/PacketTunnel/PacketTunnelProvider/PacketTunnelProvider.swift
index aae7677d9a..ab464a2274 100644
--- a/ios/PacketTunnel/PacketTunnelProvider/PacketTunnelProvider.swift
+++ b/ios/PacketTunnel/PacketTunnelProvider/PacketTunnelProvider.swift
@@ -35,7 +35,6 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
 
     override init() {
         Self.configureLogging()
-
         providerLogger = Logger(label: "PacketTunnelProvider")
 
         let containerURL = ApplicationConfiguration.containerURL
@@ -46,7 +45,6 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
             relayCache: RelayCache(cacheDirectory: containerURL),
             ipOverrideRepository: IPOverrideRepository()
         )
-
         multihopUpdater = MultihopUpdater(listener: multihopStateListener)
 
         super.init()
@@ -100,7 +98,6 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
         )
 
         let urlRequestProxy = URLRequestProxy(dispatchQueue: internalQueue, transportProvider: transportProvider)
-
         appMessageHandler = AppMessageHandler(packetTunnelActor: actor, urlRequestProxy: urlRequestProxy)
     }
 
@@ -264,7 +261,6 @@ extension PacketTunnelProvider {
                     lastConnectionAttempt = connectionAttempt
 
                 case let .negotiatingPostQuantumKey(_, privateKey):
-                    postQuantumActor.endCurrentNegotiation()
                     postQuantumActor.startNegotiation(with: privateKey)
 
                 case .initial, .connected, .disconnecting, .disconnected, .error:
@@ -311,12 +307,11 @@ extension PacketTunnelProvider {
 
 extension PacketTunnelProvider: PostQuantumKeyReceiving {
     func receivePostQuantumKey(_ key: PreSharedKey, ephemeralKey: PrivateKey) {
+        postQuantumActor.reset()
         actor.replacePreSharedKey(key, ephemeralKey: ephemeralKey)
-        postQuantumActor.endCurrentNegotiation()
     }
 
     func keyExchangeFailed() {
-        postQuantumActor.endCurrentNegotiation()
         // Do not try reconnecting to the `.current` relay, else the actor's `State` equality check will fail
         // and it will not try to reconnect
         actor.reconnect(to: .random, reconnectReason: .connectionLoss)
diff --git a/ios/PacketTunnel/PostQuantumKeyExchangeActor.swift b/ios/PacketTunnel/PostQuantumKeyExchangeActor.swift
index 33291d8ea5..2e52eefce3 100644
--- a/ios/PacketTunnel/PostQuantumKeyExchangeActor.swift
+++ b/ios/PacketTunnel/PostQuantumKeyExchangeActor.swift
@@ -8,6 +8,8 @@
 
 import Foundation
 import MullvadPostQuantum
+import MullvadREST
+import MullvadTypes
 import NetworkExtension
 import WireGuardKitTypes
 
@@ -27,11 +29,15 @@ class PostQuantumKeyExchangeActor {
     unowned let packetTunnel: PacketTunnelProvider
     private var negotiation: Negotiation?
     private var timer: DispatchSourceTimer?
+    private var keyExchangeRetriesIterator = REST.RetryStrategy.postQuantumKeyExchange.makeDelayIterator()
 
     // Callback in the event of the negotiation failing on startup
     var onFailure: () -> Void
 
-    init(packetTunnel: PacketTunnelProvider, onFailure: @escaping (() -> Void)) {
+    init(
+        packetTunnel: PacketTunnelProvider,
+        onFailure: @escaping (() -> Void)
+    ) {
         self.packetTunnel = packetTunnel
         self.onFailure = onFailure
     }
@@ -45,7 +51,16 @@ class PostQuantumKeyExchangeActor {
         )
     }
 
+    /// Starts a new key exchange.
+    ///
+    /// Any ongoing key negotiation is stopped before starting a new one.
+    /// An exponential backoff timer is used to stop the exchange if it takes too long,
+    /// or if the TCP connection takes too long to become ready.
+    /// It is reset after every successful key exchange.
+    ///
+    /// - Parameter privateKey: The device's current private key
     func startNegotiation(with privateKey: PrivateKey) {
+        endCurrentNegotiation()
         let negotiator = PostQuantumKeyNegotiator()
 
         let gatewayAddress = "10.64.0.1"
@@ -56,8 +71,9 @@ class PostQuantumKeyExchangeActor {
         // This will become the new private key of the device
         let ephemeralSharedKey = PrivateKey()
 
+        let tcpConnectionTimeout = keyExchangeRetriesIterator.next() ?? .seconds(10)
         // If the connection never becomes viable, force a reconnection after 10 seconds
-        scheduleInTunnelConnectionTimeout(startTime: .now() + 10)
+        scheduleInTunnelConnectionTimeout(startTime: .now() + tcpConnectionTimeout)
 
         let tcpConnectionObserver = inTunnelTCPConnection.observe(\.isViable, options: [
             .initial,
@@ -72,7 +88,8 @@ class PostQuantumKeyExchangeActor {
                 devicePublicKey: privateKey.publicKey,
                 presharedKey: ephemeralSharedKey,
                 packetTunnel: packetTunnel,
-                tcpConnection: inTunnelTCPConnection
+                tcpConnection: inTunnelTCPConnection,
+                postQuantumKeyExchangeTimeout: tcpConnectionTimeout
             ) {
                 self.negotiation = nil
                 self.onFailure()
@@ -85,11 +102,18 @@ class PostQuantumKeyExchangeActor {
         )
     }
 
+    /// Cancels the ongoing key exchange.
     func endCurrentNegotiation() {
         negotiation?.cancel()
         negotiation = nil
     }
 
+    /// Resets the exponential timeout for successful key exchanges, and ends the current key exchange.
+    func reset() {
+        keyExchangeRetriesIterator = REST.RetryStrategy.postQuantumKeyExchange.makeDelayIterator()
+        endCurrentNegotiation()
+    }
+
     private func scheduleInTunnelConnectionTimeout(startTime: DispatchWallTime) {
         let newTimer = DispatchSource.makeTimerSource()