1 files changed, 10 insertions, 2 deletions
diff --git a/ios/MullvadRustRuntime/PostQuantumKeyExchangeActor.swift b/ios/MullvadRustRuntime/PostQuantumKeyExchangeActor.swift
index 8028c1ba3e..a854d83cd8 100644
--- a/ios/MullvadRustRuntime/PostQuantumKeyExchangeActor.swift
+++ b/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()
             }