Use IAN TCP connection for ephemeral peer exchange

5 files changed, 124 insertions, 261 deletions

diff --git a/ios/MullvadRustRuntime/EphemeralPeerExchangeActor.swift b/ios/MullvadRustRuntime/EphemeralPeerExchangeActor.swift
index 397b656d61..3f38ccb762 100644
--- a/ios/MullvadRustRuntime/EphemeralPeerExchangeActor.swift
+++ b/ios/MullvadRustRuntime/EphemeralPeerExchangeActor.swift
@@ -21,13 +21,9 @@ public protocol EphemeralPeerExchangeActorProtocol {
 public class EphemeralPeerExchangeActor: EphemeralPeerExchangeActorProtocol {
     struct Negotiation {
         var negotiator: EphemeralPeerNegotiating
-        var inTunnelTCPConnection: NWTCPConnection
-        var tcpConnectionObserver: NSKeyValueObservation
 
         func cancel() {
             negotiator.cancelKeyNegotiation()
-            tcpConnectionObserver.invalidate()
-            inTunnelTCPConnection.cancel()
         }
     }
 
@@ -54,15 +50,6 @@ public class EphemeralPeerExchangeActor: EphemeralPeerExchangeActorProtocol {
         self.keyExchangeRetriesIterator = iteratorProvider()
     }
 
-    private func createTCPConnection(_ gatewayEndpoint: NWHostEndpoint) -> NWTCPConnection {
-        self.packetTunnel.createTCPConnectionThroughTunnel(
-            to: gatewayEndpoint,
-            enableTLS: false,
-            tlsParameters: nil,
-            delegate: nil
-        )
-    }
-
     /// Starts a new key exchange.
     ///
     /// Any ongoing key negotiation is stopped before starting a new one.
@@ -75,49 +62,46 @@ public class EphemeralPeerExchangeActor: EphemeralPeerExchangeActorProtocol {
         endCurrentNegotiation()
         let negotiator = negotiationProvider.init()
 
-        let gatewayAddress = LocalNetworkIPs.gatewayAddress.rawValue
-        let IPv4Gateway = IPv4Address(gatewayAddress)!
-        let endpoint = NWHostEndpoint(hostname: gatewayAddress, port: "\(CONFIG_SERVICE_PORT)")
-        let inTunnelTCPConnection = createTCPConnection(endpoint)
-
         // 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() + tcpConnectionTimeout)
-
-        let tcpConnectionObserver = inTunnelTCPConnection.observe(\.isViable, options: [
-            .initial,
-            .new,
-        ]) { [weak self] observedConnection, _ in
-            guard let self, observedConnection.isViable else { return }
-            self.negotiation?.tcpConnectionObserver.invalidate()
-            self.timer?.cancel()
+        let peerParameters = EphemeralPeerParameters(
+            peer_exchange_timeout: UInt64(tcpConnectionTimeout.timeInterval),
+            enable_post_quantum: enablePostQuantum,
+            enable_daita: enableDaita,
+            funcs: mapWgFunctions(functions: packetTunnel.wgFunctions())
+        )
 
-            if !negotiator.startNegotiation(
-                gatewayIP: IPv4Gateway,
-                devicePublicKey: privateKey.publicKey,
-                presharedKey: ephemeralSharedKey,
-                peerReceiver: packetTunnel,
-                tcpConnection: inTunnelTCPConnection,
-                peerExchangeTimeout: tcpConnectionTimeout,
-                enablePostQuantum: enablePostQuantum,
-                enableDaita: enableDaita
-            ) {
-                // 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()
-            }
+        if !negotiator.startNegotiation(
+            devicePublicKey: privateKey.publicKey,
+            presharedKey: ephemeralSharedKey,
+            peerReceiver: packetTunnel,
+            ephemeralPeerParams: peerParameters
+        ) {
+            // 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()
         }
+
         negotiation = Negotiation(
-            negotiator: negotiator,
-            inTunnelTCPConnection: inTunnelTCPConnection,
-            tcpConnectionObserver: tcpConnectionObserver
+            negotiator: negotiator
         )
     }
 
+    private func mapWgFunctions(functions: WgFunctionPointers) -> WgTcpConnectionFunctions {
+        var mappedFunctions = WgTcpConnectionFunctions()
+
+        mappedFunctions.close_fn = functions.close
+        mappedFunctions.open_fn = functions.open
+        mappedFunctions.send_fn = functions.send
+        mappedFunctions.recv_fn = functions.receive
+
+        return mappedFunctions
+    }
+
     /// Cancels the ongoing key exchange.
     public func endCurrentNegotiation() {
         negotiation?.cancel()
@@ -129,19 +113,4 @@ public class EphemeralPeerExchangeActor: EphemeralPeerExchangeActorProtocol {
         keyExchangeRetriesIterator = iteratorProvider()
         endCurrentNegotiation()
     }
-
-    private func scheduleInTunnelConnectionTimeout(startTime: DispatchWallTime) {
-        let newTimer = DispatchSource.makeTimerSource()
-
-        newTimer.setEventHandler { [weak self] in
-            self?.onFailure()
-            self?.timer?.cancel()
-        }
-
-        newTimer.schedule(wallDeadline: startTime)
-        newTimer.activate()
-
-        timer?.cancel()
-        timer = newTimer
-    }
 }
diff --git a/ios/MullvadRustRuntime/EphemeralPeerNegotiator.swift b/ios/MullvadRustRuntime/EphemeralPeerNegotiator.swift
index ffc0dc15b3..8346b2686d 100644
--- a/ios/MullvadRustRuntime/EphemeralPeerNegotiator.swift
+++ b/ios/MullvadRustRuntime/EphemeralPeerNegotiator.swift
@@ -14,14 +14,10 @@ import WireGuardKitTypes
 // swiftlint:disable function_parameter_count
 public protocol EphemeralPeerNegotiating {
     func startNegotiation(
-        gatewayIP: IPv4Address,
         devicePublicKey: PublicKey,
         presharedKey: PrivateKey,
         peerReceiver: any TunnelProvider,
-        tcpConnection: NWTCPConnection,
-        peerExchangeTimeout: Duration,
-        enablePostQuantum: Bool,
-        enableDaita: Bool
+        ephemeralPeerParams: EphemeralPeerParameters
     ) -> Bool
 
     func cancelKeyNegotiation()
@@ -33,35 +29,30 @@ public protocol EphemeralPeerNegotiating {
 public class EphemeralPeerNegotiator: EphemeralPeerNegotiating {
     required public init() {}
 
-    var cancelToken: EphemeralPeerCancelToken?
+    var cancelToken: OpaquePointer?
 
     public func startNegotiation(
-        gatewayIP: IPv4Address,
         devicePublicKey: PublicKey,
         presharedKey: PrivateKey,
         peerReceiver: any TunnelProvider,
-        tcpConnection: NWTCPConnection,
-        peerExchangeTimeout: Duration,
-        enablePostQuantum: Bool,
-        enableDaita: Bool
+        ephemeralPeerParams: EphemeralPeerParameters
     ) -> Bool {
         // swiftlint:disable:next force_cast
         let ephemeralPeerReceiver = Unmanaged.passUnretained(peerReceiver as! EphemeralPeerReceiver)
             .toOpaque()
-        let opaqueConnection = Unmanaged.passUnretained(tcpConnection).toOpaque()
-        var cancelToken = EphemeralPeerCancelToken()
 
-        let result = request_ephemeral_peer(
+        guard let tunnelHandle = try? peerReceiver.tunnelHandle() else {
+            return false
+        }
+
+        let cancelToken = request_ephemeral_peer(
             devicePublicKey.rawValue.map { $0 },
             presharedKey.rawValue.map { $0 },
             ephemeralPeerReceiver,
-            opaqueConnection,
-            &cancelToken,
-            UInt64(peerExchangeTimeout.timeInterval),
-            enablePostQuantum,
-            enableDaita
+            tunnelHandle,
+            ephemeralPeerParams
         )
-        guard result == 0 else {
+        guard let cancelToken else {
             return false
         }
         self.cancelToken = cancelToken
@@ -69,13 +60,14 @@ public class EphemeralPeerNegotiator: EphemeralPeerNegotiating {
     }
 
     public func cancelKeyNegotiation() {
-        guard var cancelToken else { return }
-        cancel_ephemeral_peer_exchange(&cancelToken)
+        guard let cancelToken else { return }
+        cancel_ephemeral_peer_exchange(cancelToken)
+        self.cancelToken = nil
     }
 
     deinit {
-        guard var cancelToken else { return }
-        drop_ephemeral_peer_exchange_token(&cancelToken)
+        guard let cancelToken else { return }
+        drop_ephemeral_peer_exchange_token(cancelToken)
     }
 }
 
diff --git a/ios/MullvadRustRuntime/EphemeralPeerReceiver.swift b/ios/MullvadRustRuntime/EphemeralPeerReceiver.swift
new file mode 100644
index 0000000000..2b1b4adac5
--- /dev/null
+++ b/ios/MullvadRustRuntime/EphemeralPeerReceiver.swift
@@ -0,0 +1,52 @@
+//
+//  EphemeralPeerReceiver.swift
+//  PacketTunnel
+//
+//  Created by Marco Nikic on 2024-02-15.
+//  Copyright © 2023 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+import MullvadRustRuntimeProxy
+import MullvadTypes
+import NetworkExtension
+import WireGuardKitTypes
+
+/// End sequence of an ephemeral peer exchange.
+///
+/// This FFI function is called by Rust when an ephemeral peer negotiation succeeded or failed.
+/// When both the `rawPresharedKey` and the `rawEphemeralKey` are raw pointers to 32 bytes data arrays,
+/// the quantum-secure key exchange is considered successful.
+/// If the `rawPresharedKey` is nil, but there is a valid `rawEphemeralKey`, it means a Daita peer has been negotiated with.
+/// If `rawEphemeralKey` is nil, the negotiation is considered failed.
+///
+/// - Parameters:
+///   - rawEphemeralPeerReceiver: A raw pointer to the running instance of `NEPacketTunnelProvider`
+///   - rawPresharedKey: A raw pointer to the quantum-secure pre shared key
+///   - rawEphemeralKey: A raw pointer to the ephemeral private key of the device
+@_cdecl("swift_ephemeral_peer_ready")
+func receivePostQuantumKey(
+    rawEphemeralPeerReceiver: UnsafeMutableRawPointer?,
+    rawPresharedKey: UnsafeMutableRawPointer?,
+    rawEphemeralKey: UnsafeMutableRawPointer?
+) {
+    guard let rawEphemeralPeerReceiver else { return }
+    let ephemeralPeerReceiver = Unmanaged<EphemeralPeerReceiver>.fromOpaque(rawEphemeralPeerReceiver)
+        .takeUnretainedValue()
+
+    // If there are no private keys for the ephemeral peer, then the negotiation either failed, or timed out.
+    guard let rawEphemeralKey,
+          let ephemeralKey = PrivateKey(rawValue: Data(bytes: rawEphemeralKey, count: 32)) else {
+        ephemeralPeerReceiver.ephemeralPeerExchangeFailed()
+        return
+    }
+
+    // If there is a pre-shared key, an ephemeral peer was negotiated with Post Quantum options
+    // Otherwise, a Daita enabled ephemeral peer was requested
+    if let rawPresharedKey, let key = PreSharedKey(rawValue: Data(bytes: rawPresharedKey, count: 32)) {
+        ephemeralPeerReceiver.receivePostQuantumKey(key, ephemeralKey: ephemeralKey)
+    } else {
+        ephemeralPeerReceiver.receiveEphemeralPeerPrivateKey(ephemeralKey)
+    }
+    return
+}
diff --git a/ios/MullvadRustRuntime/PacketTunnelProvider+TCPConnection.swift b/ios/MullvadRustRuntime/PacketTunnelProvider+TCPConnection.swift
deleted file mode 100644
index e19750d4dc..0000000000
--- a/ios/MullvadRustRuntime/PacketTunnelProvider+TCPConnection.swift
+++ /dev/null
@@ -1,118 +0,0 @@
-//
-//  PacketTunnelProvider+TCPConnection.swift
-//  PacketTunnel
-//
-//  Created by Marco Nikic on 2024-02-15.
-//  Copyright © 2023 Mullvad VPN AB. All rights reserved.
-//
-
-import Foundation
-import MullvadRustRuntimeProxy
-import MullvadTypes
-import NetworkExtension
-import WireGuardKitTypes
-
-/// Writes data to the in-tunnel TCP connection
-///
-/// This FFI function is called by Rust whenever there is data to be written to the in-tunnel TCP connection when exchanging
-/// quantum-resistant pre shared keys.
-///
-/// Whenever the flow control is given back from the connection, acknowledge that data was written using `rawWriteAcknowledgement`.
-/// - Parameters:
-///   - rawConnection: A raw pointer to the in-tunnel TCP connection
-///   - rawData: A raw pointer to the data to write in the connection
-///   - dataLength: The length of data to write in the connection
-///   - rawWriteAcknowledgement: An opaque pointer needed for write acknowledgement
-@_cdecl("swift_nw_tcp_connection_send")
-func tcpConnectionSend(
-    rawConnection: UnsafeMutableRawPointer?,
-    rawData: UnsafeMutableRawPointer,
-    dataLength: UInt,
-    rawWriteAcknowledgement: UnsafeMutableRawPointer?
-) {
-    guard let rawConnection, let rawWriteAcknowledgement else {
-        handle_sent(0, rawWriteAcknowledgement)
-        return
-    }
-    let tcpConnection = Unmanaged<NWTCPConnection>.fromOpaque(rawConnection).takeUnretainedValue()
-    let data = Data(bytes: rawData, count: Int(dataLength))
-
-    // The guarantee that all writes are sequential is done by virtue of not returning the execution context
-    // to Rust before this closure is done executing.
-    tcpConnection.write(data, completionHandler: { maybeError in
-        if maybeError != nil {
-            handle_sent(0, rawWriteAcknowledgement)
-        } else {
-            handle_sent(dataLength, rawWriteAcknowledgement)
-        }
-    })
-}
-
-/// Reads data to the in-tunnel TCP connection
-///
-/// This FFI function is called by Rust whenever there is data to be read from the in-tunnel TCP connection when exchanging
-/// quantum-resistant pre shared keys.
-///
-/// Whenever the flow control is given back from the connection, acknowledge that data was read using `rawReadAcknowledgement`.
-/// - Parameters:
-///   - rawConnection: A raw pointer to the in-tunnel TCP connection
-///   - rawReadAcknowledgement: An opaque pointer needed for read acknowledgement
-@_cdecl("swift_nw_tcp_connection_read")
-func tcpConnectionReceive(
-    rawConnection: UnsafeMutableRawPointer?,
-    rawReadAcknowledgement: UnsafeMutableRawPointer?
-) {
-    guard let rawConnection, let rawReadAcknowledgement else {
-        handle_recv(nil, 0, rawReadAcknowledgement)
-        return
-    }
-    let tcpConnection = Unmanaged<NWTCPConnection>.fromOpaque(rawConnection).takeUnretainedValue()
-    tcpConnection.readMinimumLength(0, maximumLength: Int(UInt16.max)) { data, maybeError in
-        if let data {
-            if maybeError != nil {
-                handle_recv(nil, 0, rawReadAcknowledgement)
-            } else {
-                handle_recv(data.map { $0 }, UInt(data.count), rawReadAcknowledgement)
-            }
-        }
-    }
-}
-
-/// End sequence of an ephemeral peer exchange.
-///
-/// This FFI function is called by Rust when an ephemeral peer negotiation succeeded or failed.
-/// When both the `rawPresharedKey` and the `rawEphemeralKey` are raw pointers to 32 bytes data arrays,
-/// the quantum-secure key exchange is considered successful.
-/// If the `rawPresharedKey` is nil, but there is a valid `rawEphemeralKey`, it means a Daita peer has been negotiated with.
-/// If `rawEphemeralKey` is nil, the negotiation is considered failed.
-///
-/// - Parameters:
-///   - rawEphemeralPeerReceiver: A raw pointer to the running instance of `NEPacketTunnelProvider`
-///   - rawPresharedKey: A raw pointer to the quantum-secure pre shared key
-///   - rawEphemeralKey: A raw pointer to the ephemeral private key of the device
-@_cdecl("swift_ephemeral_peer_ready")
-func receivePostQuantumKey(
-    rawEphemeralPeerReceiver: UnsafeMutableRawPointer?,
-    rawPresharedKey: UnsafeMutableRawPointer?,
-    rawEphemeralKey: UnsafeMutableRawPointer?
-) {
-    guard let rawEphemeralPeerReceiver else { return }
-    let ephemeralPeerReceiver = Unmanaged<EphemeralPeerReceiver>.fromOpaque(rawEphemeralPeerReceiver)
-        .takeUnretainedValue()
-
-    // If there are no private keys for the ephemeral peer, then the negotiation either failed, or timed out.
-    guard let rawEphemeralKey,
-          let ephemeralKey = PrivateKey(rawValue: Data(bytes: rawEphemeralKey, count: 32)) else {
-        ephemeralPeerReceiver.ephemeralPeerExchangeFailed()
-        return
-    }
-
-    // If there is a pre-shared key, an ephemeral peer was negotiated with Post Quantum options
-    // Otherwise, a Daita enabled ephemeral peer was requested
-    if let rawPresharedKey, let key = PreSharedKey(rawValue: Data(bytes: rawPresharedKey, count: 32)) {
-        ephemeralPeerReceiver.receivePostQuantumKey(key, ephemeralKey: ephemeralKey)
-    } else {
-        ephemeralPeerReceiver.receiveEphemeralPeerPrivateKey(ephemeralKey)
-    }
-    return
-}
diff --git a/ios/MullvadRustRuntime/include/mullvad_rust_runtime.h b/ios/MullvadRustRuntime/include/mullvad_rust_runtime.h
index a45c6ed6c3..93c04587f1 100644
--- a/ios/MullvadRustRuntime/include/mullvad_rust_runtime.h
+++ b/ios/MullvadRustRuntime/include/mullvad_rust_runtime.h
@@ -20,14 +20,26 @@ typedef uint8_t TunnelObfuscatorProtocol;
  */
 typedef struct EncryptedDnsProxyState EncryptedDnsProxyState;
 
+typedef struct ExchangeCancelToken ExchangeCancelToken;
+
 typedef struct ProxyHandle {
   void *context;
   uint16_t port;
 } ProxyHandle;
 
-typedef struct EphemeralPeerCancelToken {
-  void *context;
-} EphemeralPeerCancelToken;
+typedef struct WgTcpConnectionFunctions {
+  int32_t (*open_fn)(int32_t tunnelHandle, const char *address, uint64_t timeout);
+  int32_t (*close_fn)(int32_t tunnelHandle, int32_t socketHandle);
+  int32_t (*recv_fn)(int32_t tunnelHandle, int32_t socketHandle, uint8_t *data, int32_t len);
+  int32_t (*send_fn)(int32_t tunnelHandle, int32_t socketHandle, const uint8_t *data, int32_t len);
+} WgTcpConnectionFunctions;
+
+typedef struct EphemeralPeerParameters {
+  uint64_t peer_exchange_timeout;
+  bool enable_post_quantum;
+  bool enable_daita;
+  struct WgTcpConnectionFunctions funcs;
+} EphemeralPeerParameters;
 
 extern const uint16_t CONFIG_SERVICE_PORT;
 
@@ -84,43 +96,17 @@ int32_t encrypted_dns_proxy_stop(struct ProxyHandle *proxy_config);
  * `sender` must be pointing to a valid instance of a `EphemeralPeerCancelToken` created by the
  * `PacketTunnelProvider`.
  */
-void cancel_ephemeral_peer_exchange(const struct EphemeralPeerCancelToken *sender);
+void cancel_ephemeral_peer_exchange(struct ExchangeCancelToken *sender);
 
 /**
- * Called by the Swift side to signal that the Rust `EphemeralPeerCancelToken` can be safely dropped
- * from memory.
+ * Called by the Swift side to signal that the Rust `EphemeralPeerCancelToken` can be safely
+ * dropped from memory.
  *
  * # Safety
  * `sender` must be pointing to a valid instance of a `EphemeralPeerCancelToken` created by the
  * `PacketTunnelProvider`.
  */
-void drop_ephemeral_peer_exchange_token(const struct EphemeralPeerCancelToken *sender);
-
-/**
- * Called by Swift whenever data has been written to the in-tunnel TCP connection when exchanging
- * quantum-resistant pre shared keys, or ephemeral peers.
- *
- * If `bytes_sent` is 0, this indicates that the connection was closed or that an error occurred.
- *
- * # Safety
- * `sender` must be pointing to a valid instance of a `write_tx` created by the `IosTcpProvider`
- * Callback to call when the TCP connection has written data.
- */
-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, or ephemeral peers.
- *
- * 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 drop_ephemeral_peer_exchange_token(struct ExchangeCancelToken *sender);
 
 /**
  * Entry point for requesting ephemeral peers on iOS.
@@ -128,33 +114,15 @@ void handle_recv(const uint8_t *data, uintptr_t data_len, const void *sender);
  * # 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.
- * `cancel_token` should be owned by the caller of this function.
- */
-int32_t request_ephemeral_peer(const uint8_t *public_key,
-                               const uint8_t *ephemeral_key,
-                               const void *packet_tunnel,
-                               const void *tcp_connection,
-                               struct EphemeralPeerCancelToken *cancel_token,
-                               uint64_t peer_exchange_timeout,
-                               bool enable_post_quantum,
-                               bool enable_daita);
-
-/**
- * Called when there is data to send on the TCP connection.
- * The TCP connection must write data on the wire, then call the `handle_sent` function.
- */
-extern void swift_nw_tcp_connection_send(const void *connection,
-                                         const void *data,
-                                         uintptr_t data_len,
-                                         const void *sender);
-
-/**
- * Called when there is data to read on the TCP connection.
- * The TCP connection must read data from the wire, then call the `handle_read` function.
+ * `packet_tunnel` must be valid pointers to a packet tunnel, the packet tunnel pointer must
+ * outlive the ephemeral peer exchange. `cancel_token` should be owned by the caller of this
+ * function.
  */
-extern void swift_nw_tcp_connection_read(const void *connection, const void *sender);
+struct ExchangeCancelToken *request_ephemeral_peer(const uint8_t *public_key,
+                                                   const uint8_t *ephemeral_key,
+                                                   const void *packet_tunnel,
+                                                   int32_t tunnel_handle,
+                                                   struct EphemeralPeerParameters peer_parameters);
 
 /**
  * Called when the preshared post quantum key is ready,