Add a backing off timeout when negotiating PQ PSK

8 files changed, 186 insertions, 104 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()
 
diff --git a/talpid-tunnel-config-client/src/ios_ffi/ios_runtime.rs b/talpid-tunnel-config-client/src/ios_ffi/ios_runtime.rs
index 2989b4fb48..ba93ae2ece 100644
--- a/talpid-tunnel-config-client/src/ios_ffi/ios_runtime.rs
+++ b/talpid-tunnel-config-client/src/ios_ffi/ios_runtime.rs
@@ -1,9 +1,17 @@
 use super::{ios_tcp_connection::*, PostQuantumCancelToken};
 use crate::{request_ephemeral_peer_with, Error, RelayConfigService};
 use libc::c_void;
-use std::{future::Future, io, pin::Pin, ptr, sync::Arc};
+use std::{
+    future::Future,
+    io,
+    pin::Pin,
+    ptr,
+    sync::{Arc, Mutex},
+};
 use talpid_types::net::wireguard::{PrivateKey, PublicKey};
-use tokio::{runtime::Builder, sync::mpsc};
+use tokio::{
+    runtime::Builder,
+};
 use tonic::transport::channel::Endpoint;
 use tower::util::service_fn;
 
@@ -15,10 +23,19 @@ pub unsafe fn run_post_quantum_psk_exchange(
     ephemeral_key: [u8; 32],
     packet_tunnel: *const c_void,
     tcp_connection: *const c_void,
+    post_quantum_key_exchange_timeout: u64,
 ) -> Result<PostQuantumCancelToken, Error> {
-    match unsafe { IOSRuntime::new(pub_key, ephemeral_key, packet_tunnel, tcp_connection) } {
+    match unsafe {
+        IOSRuntime::new(
+            pub_key,
+            ephemeral_key,
+            packet_tunnel,
+            tcp_connection,
+            post_quantum_key_exchange_timeout,
+        )
+    } {
         Ok(runtime) => {
-            let token = runtime.cancel_token_tx.clone();
+            let token = runtime.packet_tunnel.tcp_connection.clone();
 
             runtime.run();
             Ok(PostQuantumCancelToken {
@@ -35,7 +52,7 @@ pub unsafe fn run_post_quantum_psk_exchange(
 #[derive(Clone)]
 pub struct SwiftContext {
     pub packet_tunnel: *const c_void,
-    pub tcp_connection: *const c_void,
+    pub tcp_connection: Arc<Mutex<ConnectionContext>>,
 }
 
 unsafe impl Send for SwiftContext {}
@@ -46,8 +63,7 @@ struct IOSRuntime {
     pub_key: [u8; 32],
     ephemeral_key: [u8; 32],
     packet_tunnel: SwiftContext,
-    cancel_token_tx: Arc<mpsc::UnboundedSender<()>>,
-    cancel_token_rx: mpsc::UnboundedReceiver<()>,
+    post_quantum_key_exchange_timeout: u64,
 }
 
 impl IOSRuntime {
@@ -56,6 +72,7 @@ impl IOSRuntime {
         ephemeral_key: [u8; 32],
         packet_tunnel: *const libc::c_void,
         tcp_connection: *const c_void,
+        post_quantum_key_exchange_timeout: u64,
     ) -> io::Result<Self> {
         let runtime = Builder::new_multi_thread()
             .enable_all()
@@ -64,18 +81,15 @@ impl IOSRuntime {
 
         let context = SwiftContext {
             packet_tunnel,
-            tcp_connection,
+            tcp_connection: Arc::new(Mutex::new(ConnectionContext::new(tcp_connection))),
         };
 
-        let (tx, rx) = mpsc::unbounded_channel();
-
         Ok(Self {
             runtime,
             pub_key,
             ephemeral_key,
             packet_tunnel: context,
-            cancel_token_tx: Arc::new(tx),
-            cancel_token_rx: rx,
+            post_quantum_key_exchange_timeout,
         })
     }
 
@@ -84,9 +98,8 @@ impl IOSRuntime {
             self.run_service_inner();
         });
     }
-
-    /// Creates a `RelayConfigService` using the in-tunnel TCP Connection provided by the Packet Tunnel Provider
-    /// # Safety
+    /// Creates a `RelayConfigService` using the in-tunnel TCP Connection provided by the Packet
+    /// Tunnel Provider # Safety
     /// It is unsafe to call this with an already used `SwiftContext`
     async unsafe fn ios_tcp_client(
         ctx: SwiftContext,
@@ -112,11 +125,7 @@ impl IOSRuntime {
     }
 
     fn run_service_inner(self) {
-        let Self {
-            runtime,
-            mut cancel_token_rx,
-            ..
-        } = self;
+        let Self { runtime, .. } = self;
 
         let packet_tunnel_ptr = self.packet_tunnel.packet_tunnel;
         runtime.block_on(async move {
@@ -158,14 +167,9 @@ impl IOSRuntime {
                     }
                 }
 
-                _ = tokio::time::sleep(std::time::Duration::from_secs(5)) => {
+                _ = tokio::time::sleep(std::time::Duration::from_secs(self.post_quantum_key_exchange_timeout)) => {
+                            shutdown_handle.shutdown();
                             unsafe { swift_post_quantum_key_ready(packet_tunnel_ptr, ptr::null(), ptr::null()); }
-                            shutdown_handle.shutdown()
-                }
-
-                _ = cancel_token_rx.recv() => {
-                    shutdown_handle.shutdown()
-                    // The swift runtime pre emptively cancelled the key exchange, nothing to do here.
                 }
             }
         });
diff --git a/talpid-tunnel-config-client/src/ios_ffi/ios_tcp_connection.rs b/talpid-tunnel-config-client/src/ios_ffi/ios_tcp_connection.rs
index 86a3114edf..656bb142c8 100644
--- a/talpid-tunnel-config-client/src/ios_ffi/ios_tcp_connection.rs
+++ b/talpid-tunnel-config-client/src/ios_ffi/ios_tcp_connection.rs
@@ -1,10 +1,7 @@
 use libc::c_void;
 use std::{
     io::{self, Result},
-    sync::{
-        atomic::{self, AtomicBool},
-        Arc, Mutex, Weak,
-    },
+    sync::{Arc, Mutex, MutexGuard, Weak},
     task::{Poll, Waker},
 };
 use tokio::{
@@ -49,28 +46,29 @@ pub struct IosTcpProvider {
     write_rx: mpsc::UnboundedReceiver<usize>,
     read_tx: Arc<mpsc::UnboundedSender<Box<[u8]>>>,
     read_rx: mpsc::UnboundedReceiver<Box<[u8]>>,
-    tcp_connection: *const c_void,
+    tcp_connection: Arc<Mutex<ConnectionContext>>,
     read_in_progress: bool,
     write_in_progress: bool,
-    shutdown: Arc<AtomicBool>,
-    waker: Arc<Mutex<Option<Waker>>>,
 }
 
 pub struct IosTcpShutdownHandle {
-    shutdown: Arc<AtomicBool>,
-    waker: Arc<Mutex<Option<Waker>>>,
+    context: Arc<Mutex<ConnectionContext>>,
 }
 
+pub struct ConnectionContext {
+    waker: Option<Waker>,
+    tcp_connection: Option<*const c_void>,
+}
+
+unsafe impl Send for ConnectionContext {}
+
 impl IosTcpProvider {
-    /**
-     * # Safety
-     * `tcp_connection` must be pointing to a valid instance of a `NWTCPConnection`, created by the `PacketTunnelProvider`
-     */
-    pub unsafe fn new(tcp_connection: *const c_void) -> (Self, IosTcpShutdownHandle) {
+    /// # Safety
+    /// `tcp_connection` must be pointing to a valid instance of a `NWTCPConnection`, created by the
+    /// `PacketTunnelProvider`
+    pub unsafe fn new(connection: Arc<Mutex<ConnectionContext>>) -> (Self, IosTcpShutdownHandle) {
         let (tx, rx) = mpsc::unbounded_channel();
         let (recv_tx, recv_rx) = mpsc::unbounded_channel();
-        let shutdown = Arc::new(AtomicBool::new(false));
-        let waker = Arc::new(Mutex::new(None));
 
         (
             Self {
@@ -78,32 +76,33 @@ impl IosTcpProvider {
                 write_rx: rx,
                 read_tx: Arc::new(recv_tx),
                 read_rx: recv_rx,
-                tcp_connection,
+                tcp_connection: connection.clone(),
                 read_in_progress: false,
                 write_in_progress: false,
-                shutdown: shutdown.clone(),
-                waker: waker.clone(),
             },
-            IosTcpShutdownHandle { shutdown, waker },
+            IosTcpShutdownHandle {
+                context: connection,
+            },
         )
     }
 
-    fn is_shutdown(&self) -> bool {
-        self.shutdown.load(atomic::Ordering::SeqCst)
-    }
-
-    fn maybe_set_waker(&self, new_waker: Waker) {
-        if let Ok(mut waker) = self.waker.lock() {
-            *waker = Some(new_waker);
-        }
+    fn maybe_set_waker(new_waker: Waker, connection: &mut MutexGuard<'_, ConnectionContext>) {
+        connection.waker = Some(new_waker);
     }
 }
 
 impl IosTcpShutdownHandle {
-    pub fn shutdown(&self) {
-        self.shutdown.store(true, atomic::Ordering::SeqCst);
-        if let Some(waker) = self.waker.lock().ok().and_then(|mut waker| waker.take()) {
-            waker.wake();
+    pub fn shutdown(self) {
+        {
+            let Ok(mut context) = self.context.lock() else {
+                return;
+            };
+
+            context.tcp_connection = None;
+            if let Some(waker) = context.waker.take() {
+                waker.wake();
+            }
+            std::mem::drop(context);
         }
     }
 }
@@ -114,7 +113,15 @@ impl AsyncWrite for IosTcpProvider {
         cx: &mut std::task::Context<'_>,
         buf: &[u8],
     ) -> std::task::Poll<Result<usize>> {
-        self.maybe_set_waker(cx.waker().clone());
+        let connection_lock = self.tcp_connection.clone();
+        let Ok(mut connection) = connection_lock.lock() else {
+            return Poll::Ready(Err(connection_closed_err()));
+        };
+        let Some(tcp_ptr) = connection.tcp_connection else {
+            return Poll::Ready(Err(connection_closed_err()));
+        };
+        Self::maybe_set_waker(cx.waker().clone(), &mut connection);
+
         match self.write_rx.poll_recv(cx) {
             std::task::Poll::Ready(Some(bytes_sent)) => {
                 self.write_in_progress = false;
@@ -125,14 +132,12 @@ impl AsyncWrite for IosTcpProvider {
                 Poll::Ready(Err(connection_closed_err()))
             }
             std::task::Poll::Pending => {
-                if self.is_shutdown() {
-                    return Poll::Ready(Err(connection_closed_err()));
-                }
                 if !self.write_in_progress {
                     let raw_sender = Weak::into_raw(Arc::downgrade(&self.write_tx));
                     unsafe {
                         swift_nw_tcp_connection_send(
-                            self.tcp_connection,
+                            // self.tcp_connection,
+                            tcp_ptr,
                             buf.as_ptr() as _,
                             buf.len(),
                             raw_sender as _,
@@ -165,9 +170,14 @@ impl AsyncRead for IosTcpProvider {
         cx: &mut std::task::Context<'_>,
         buf: &mut tokio::io::ReadBuf<'_>,
     ) -> std::task::Poll<std::io::Result<()>> {
-        if self.is_shutdown() {
+        let connection_lock = self.tcp_connection.clone();
+        let Ok(mut connection) = connection_lock.lock() else {
             return Poll::Ready(Err(connection_closed_err()));
-        }
+        };
+        let Some(tcp_ptr) = connection.tcp_connection else {
+            return Poll::Ready(Err(connection_closed_err()));
+        };
+        Self::maybe_set_waker(cx.waker().clone(), &mut connection);
 
         match self.read_rx.poll_recv(cx) {
             std::task::Poll::Ready(Some(data)) => {
@@ -183,7 +193,8 @@ impl AsyncRead for IosTcpProvider {
                 if !self.read_in_progress {
                     let raw_sender = Weak::into_raw(Arc::downgrade(&self.read_tx));
                     unsafe {
-                        swift_nw_tcp_connection_read(self.tcp_connection, raw_sender as _);
+                        // TODO
+                        swift_nw_tcp_connection_read(tcp_ptr, raw_sender as _);
                     }
                     self.read_in_progress = true;
                 }
@@ -192,3 +203,19 @@ impl AsyncRead for IosTcpProvider {
         }
     }
 }
+
+impl ConnectionContext {
+    pub fn new(tcp_connection: *const c_void) -> Self {
+        Self {
+            tcp_connection: Some(tcp_connection),
+            waker: None,
+        }
+    }
+
+    pub fn shutdown(&mut self) {
+        self.tcp_connection = None;
+        if let Some(waker) = self.waker.take() {
+            waker.wake();
+        }
+    }
+}
diff --git a/talpid-tunnel-config-client/src/ios_ffi/mod.rs b/talpid-tunnel-config-client/src/ios_ffi/mod.rs
index 94dcf95b6f..704332a7bf 100644
--- a/talpid-tunnel-config-client/src/ios_ffi/mod.rs
+++ b/talpid-tunnel-config-client/src/ios_ffi/mod.rs
@@ -2,8 +2,9 @@ pub mod ios_runtime;
 pub mod ios_tcp_connection;
 
 use crate::ios_ffi::ios_runtime::run_post_quantum_psk_exchange;
+use ios_tcp_connection::ConnectionContext;
 use libc::c_void;
-use std::sync::{Arc, Weak};
+use std::sync::{Arc, Mutex, Weak};
 use tokio::sync::mpsc;
 
 use std::sync::Once;
@@ -20,34 +21,43 @@ impl PostQuantumCancelToken {
     /// This function can only be called when the context pointer is valid.
     unsafe fn cancel(&self) {
         // # Safety
-        // Try to take the value, if there is a value, we can safely send the message, otherwise, assume it has been dropped and nothing happens
-        let send_tx: Arc<mpsc::UnboundedSender<()>> = unsafe { Arc::from_raw(self.context as _) };
-        let _ = send_tx.send(());
+        // Try to take the value, if there is a value, we can safely send the message, otherwise,
+        // assume it has been dropped and nothing happens
+        let connection_context: Arc<Mutex<ConnectionContext>> =
+            unsafe { Arc::from_raw(self.context as _) };
+        if let Ok(mut connection) = connection_context.lock() {
+            connection.shutdown();
+        }
+
         // Call std::mem::forget here to avoid dropping the channel.
-        std::mem::forget(send_tx);
+        std::mem::forget(connection_context);
     }
 }
 
 impl Drop for PostQuantumCancelToken {
     fn drop(&mut self) {
-        let _: Arc<mpsc::UnboundedSender<()>> = unsafe { Arc::from_raw(self.context as _) };
+        let _: Arc<Mutex<ConnectionContext>> = unsafe { Arc::from_raw(self.context as _) };
     }
 }
+
 unsafe impl Send for 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`.
 #[no_mangle]
 pub unsafe extern "C" fn cancel_post_quantum_key_exchange(sender: *const PostQuantumCancelToken) {
     let sender = unsafe { &*sender };
     sender.cancel();
 }
-/// 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`.
 #[no_mangle]
 pub unsafe extern "C" fn drop_post_quantum_key_exchange_token(
     sender: *const PostQuantumCancelToken,
@@ -74,8 +84,8 @@ pub unsafe extern "C" fn handle_sent(bytes_sent: usize, sender: *const c_void) {
 /// 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`
@@ -102,8 +112,8 @@ pub unsafe extern "C" fn handle_recv(data: *const u8, mut data_len: usize, sende
 /// # 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.
 #[no_mangle]
 pub unsafe extern "C" fn negotiate_post_quantum_key(
@@ -112,10 +122,11 @@ pub unsafe extern "C" fn negotiate_post_quantum_key(
     packet_tunnel: *const c_void,
     tcp_connection: *const c_void,
     cancel_token: *mut PostQuantumCancelToken,
+    post_quantum_key_exchange_timeout: u64,
 ) -> i32 {
     INIT_LOGGING.call_once(|| {
         let _ = oslog::OsLogger::new("net.mullvad.MullvadVPN.TTCC")
-            .level_filter(log::LevelFilter::Trace)
+            .level_filter(log::LevelFilter::Debug)
             .init();
     });
 
@@ -123,7 +134,13 @@ pub unsafe extern "C" fn negotiate_post_quantum_key(
     let eph_key_copy: [u8; 32] = unsafe { std::ptr::read(ephemeral_key as *const [u8; 32]) };
 
     match unsafe {
-        run_post_quantum_psk_exchange(pub_key_copy, eph_key_copy, packet_tunnel, tcp_connection)
+        run_post_quantum_psk_exchange(
+            pub_key_copy,
+            eph_key_copy,
+            packet_tunnel,
+            tcp_connection,
+            post_quantum_key_exchange_timeout,
+        )
     } {
         Ok(token) => {
             unsafe { std::ptr::write(cancel_token, token) };