Change key exchange to be asynchronous

7 files changed, 67 insertions, 41 deletions

diff --git a/ios/MullvadTypes/Promise.swift b/ios/MullvadTypes/Promise.swift
index 48a12f8182..ab80c31695 100644
--- a/ios/MullvadTypes/Promise.swift
+++ b/ios/MullvadTypes/Promise.swift
@@ -48,18 +48,30 @@ public final class Promise<Success, Failure: Error> {
     }
 }
 
-
+// This object can be used like an async semaphore with exactly 1 writer. It
+// allows the waiter to wait to `receive()` from another operation
+// asynchronously. It is important not to forget to call `send`, otherwise this
+// operation will block indefinitely.
 public struct OneshotChannel {
-    private let semaphore = DispatchSemaphore(value: 0)
-    
+    private var continuation: AsyncStream<Void>.Continuation?
+    private var stream: AsyncStream<Void>
+
     public init() {
+        var ownedContinuation: AsyncStream<Void>.Continuation?
+        stream = AsyncStream { continuation in
+            ownedContinuation = continuation
+        }
+        self.continuation = ownedContinuation
     }
-    
-    public mutating func send() {
-        semaphore.signal()
+
+    public func send() {
+        continuation?.yield()
+        continuation?.finish()
     }
-    
-    public func receive() {
-        semaphore.wait()
+
+    public func receive() async {
+        for await _ in stream {
+            return
+        }
     }
 }
diff --git a/ios/MullvadTypes/Protocols/EphemeralPeerReceiver.swift b/ios/MullvadTypes/Protocols/EphemeralPeerReceiver.swift
index f7fbe8e4a1..7f45669468 100644
--- a/ios/MullvadTypes/Protocols/EphemeralPeerReceiver.swift
+++ b/ios/MullvadTypes/Protocols/EphemeralPeerReceiver.swift
@@ -30,13 +30,21 @@ public class EphemeralPeerReceiver: EphemeralPeerReceiving, TunnelProvider {
     // MARK: - EphemeralPeerReceiving
 
     public func receivePostQuantumKey(_ key: PreSharedKey, ephemeralKey: PrivateKey) {
-        guard let receiver = tunnelProvider as? EphemeralPeerReceiving else { return }
-        receiver.receivePostQuantumKey(key, ephemeralKey: ephemeralKey)
+        let semaphore = DispatchSemaphore(value: 0)
+        Task {
+            await keyReceiver.receivePostQuantumKey(key, ephemeralKey: ephemeralKey)
+            semaphore.signal()
+        }
+        semaphore.wait()
     }
 
     public func receiveEphemeralPeerPrivateKey(_ ephemeralPeerPrivateKey: PrivateKey) {
-        guard let receiver = tunnelProvider as? EphemeralPeerReceiving else { return }
-        receiver.receiveEphemeralPeerPrivateKey(ephemeralPeerPrivateKey)
+        let semaphore = DispatchSemaphore(value: 0)
+        Task {
+            await keyReceiver.receiveEphemeralPeerPrivateKey(ephemeralPeerPrivateKey)
+            semaphore.signal()
+        }
+        semaphore.wait()
     }
 
     public func ephemeralPeerExchangeFailed() {
diff --git a/ios/MullvadTypes/Protocols/EphemeralPeerReceiving.swift b/ios/MullvadTypes/Protocols/EphemeralPeerReceiving.swift
index d55ec09f1f..cc1c8f0f8a 100644
--- a/ios/MullvadTypes/Protocols/EphemeralPeerReceiving.swift
+++ b/ios/MullvadTypes/Protocols/EphemeralPeerReceiving.swift
@@ -15,11 +15,11 @@ public protocol EphemeralPeerReceiving {
     /// - Parameters:
     ///   - key: The preshared key used by the Ephemeral Peer
     ///   - ephemeralKey: The private key used by the Ephemeral Peer
-    func receivePostQuantumKey(_ key: PreSharedKey, ephemeralKey: PrivateKey)
+    func receivePostQuantumKey(_ key: PreSharedKey, ephemeralKey: PrivateKey) async
 
     /// Called when successfully requesting an ephemeral peer with Daita enabled, and Post Quantum PSK disabled
     /// - Parameter _:_ The private key used by the Ephemeral Peer
-    func receiveEphemeralPeerPrivateKey(_: PrivateKey)
+    func receiveEphemeralPeerPrivateKey(_: PrivateKey) async
 
     /// Called when an ephemeral peer could not be successfully negotiated
     func ephemeralPeerExchangeFailed()
diff --git a/ios/PacketTunnel/PacketTunnelProvider/PacketTunnelProvider.swift b/ios/PacketTunnel/PacketTunnelProvider/PacketTunnelProvider.swift
index 3216773eb5..75678417ef 100644
--- a/ios/PacketTunnel/PacketTunnelProvider/PacketTunnelProvider.swift
+++ b/ios/PacketTunnel/PacketTunnelProvider/PacketTunnelProvider.swift
@@ -111,8 +111,11 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
             ),
             onUpdateConfiguration: { [unowned self] configuration in
                 let channel = OneshotChannel()
-                actor.changeEphemeralPeerNegotiationState(configuration: configuration, reconfigurationSemaphore: channel)
-                channel.receive()
+                actor.changeEphemeralPeerNegotiationState(
+                    configuration: configuration,
+                    reconfigurationSemaphore: channel
+                )
+                await channel.receive()
             }, onFinish: { [unowned self] in
                 actor.notifyEphemeralPeerNegotiated()
             }
@@ -313,7 +316,10 @@ extension PacketTunnelProvider {
                     lastConnectionAttempt = connectionAttempt
 
                 case let .negotiatingEphemeralPeer(observedConnectionState, privateKey):
-                    ephemeralPeerExchangingPipeline.startNegotiation(observedConnectionState, privateKey: privateKey)
+                    await ephemeralPeerExchangingPipeline.startNegotiation(
+                        observedConnectionState,
+                        privateKey: privateKey
+                    )
                 case .initial, .connected, .disconnecting, .disconnected, .error:
                     break
                 }
@@ -370,12 +376,12 @@ extension PacketTunnelProvider {
 }
 
 extension PacketTunnelProvider: EphemeralPeerReceiving {
-    func receivePostQuantumKey(_ key: PreSharedKey, ephemeralKey: PrivateKey) {
-        ephemeralPeerExchangingPipeline.receivePostQuantumKey(key, ephemeralKey: ephemeralKey)
+    func receivePostQuantumKey(_ key: PreSharedKey, ephemeralKey: PrivateKey) async {
+        await ephemeralPeerExchangingPipeline.receivePostQuantumKey(key, ephemeralKey: ephemeralKey)
     }
 
-    public func receiveEphemeralPeerPrivateKey(_ ephemeralPeerPrivateKey: PrivateKey) {
-        ephemeralPeerExchangingPipeline.receiveEphemeralPeerPrivateKey(ephemeralPeerPrivateKey)
+    public func receiveEphemeralPeerPrivateKey(_ ephemeralPeerPrivateKey: PrivateKey) async {
+        await ephemeralPeerExchangingPipeline.receiveEphemeralPeerPrivateKey(ephemeralPeerPrivateKey)
     }
 
     func ephemeralPeerExchangeFailed() {
diff --git a/ios/PacketTunnel/PostQuantum/EphemeralPeerExchangingPipeline.swift b/ios/PacketTunnel/PostQuantum/EphemeralPeerExchangingPipeline.swift
index 0d0da931f0..1619889003 100644
--- a/ios/PacketTunnel/PostQuantum/EphemeralPeerExchangingPipeline.swift
+++ b/ios/PacketTunnel/PostQuantum/EphemeralPeerExchangingPipeline.swift
@@ -13,14 +13,14 @@ import WireGuardKitTypes
 
 final public class EphemeralPeerExchangingPipeline {
     let keyExchanger: EphemeralPeerExchangeActorProtocol
-    let onUpdateConfiguration: (EphemeralPeerNegotiationState) -> Void
+    let onUpdateConfiguration: (EphemeralPeerNegotiationState) async -> Void
     let onFinish: () -> Void
 
     private var ephemeralPeerExchanger: EphemeralPeerExchangingProtocol!
 
     public init(
         _ keyExchanger: EphemeralPeerExchangeActorProtocol,
-        onUpdateConfiguration: @escaping (EphemeralPeerNegotiationState) -> Void,
+        onUpdateConfiguration: @escaping (EphemeralPeerNegotiationState) async -> Void,
         onFinish: @escaping () -> Void
     ) {
         self.keyExchanger = keyExchanger
@@ -28,7 +28,7 @@ final public class EphemeralPeerExchangingPipeline {
         self.onFinish = onFinish
     }
 
-    public func startNegotiation(_ connectionState: ObservedConnectionState, privateKey: PrivateKey) {
+    public func startNegotiation(_ connectionState: ObservedConnectionState, privateKey: PrivateKey) async {
         keyExchanger.reset()
         let entryPeer = connectionState.selectedRelays.entry
         let exitPeer = connectionState.selectedRelays.exit
@@ -56,14 +56,14 @@ final public class EphemeralPeerExchangingPipeline {
                 onFinish: onFinish
             )
         }
-        ephemeralPeerExchanger.start()
+        await ephemeralPeerExchanger.start()
     }
 
-    public func receivePostQuantumKey(_ key: PreSharedKey, ephemeralKey: PrivateKey) {
-        ephemeralPeerExchanger.receivePostQuantumKey(key, ephemeralKey: ephemeralKey)
+    public func receivePostQuantumKey(_ key: PreSharedKey, ephemeralKey: PrivateKey) async {
+        await ephemeralPeerExchanger.receivePostQuantumKey(key, ephemeralKey: ephemeralKey)
     }
 
-    public func receiveEphemeralPeerPrivateKey(_ ephemeralPeerPrivateKey: PrivateKey) {
-        ephemeralPeerExchanger.receiveEphemeralPeerPrivateKey(ephemeralPeerPrivateKey)
+    public func receiveEphemeralPeerPrivateKey(_ ephemeralPeerPrivateKey: PrivateKey) async {
+        await ephemeralPeerExchanger.receiveEphemeralPeerPrivateKey(ephemeralPeerPrivateKey)
     }
 }
diff --git a/ios/PacketTunnel/PostQuantum/SingleHopEphemeralPeerExchanger.swift b/ios/PacketTunnel/PostQuantum/SingleHopEphemeralPeerExchanger.swift
index f959a335ef..3c6ab5631a 100644
--- a/ios/PacketTunnel/PostQuantum/SingleHopEphemeralPeerExchanger.swift
+++ b/ios/PacketTunnel/PostQuantum/SingleHopEphemeralPeerExchanger.swift
@@ -18,7 +18,7 @@ struct SingleHopEphemeralPeerExchanger: EphemeralPeerExchangingProtocol {
     let keyExchanger: EphemeralPeerExchangeActorProtocol
     let devicePrivateKey: PrivateKey
     let onFinish: () -> Void
-    let onUpdateConfiguration: (EphemeralPeerNegotiationState) -> Void
+    let onUpdateConfiguration: (EphemeralPeerNegotiationState) async -> Void
     let enablePostQuantum: Bool
     let enableDaita: Bool
 
@@ -28,7 +28,7 @@ struct SingleHopEphemeralPeerExchanger: EphemeralPeerExchangingProtocol {
         keyExchanger: EphemeralPeerExchangeActorProtocol,
         enablePostQuantum: Bool,
         enableDaita: Bool,
-        onUpdateConfiguration: @escaping (EphemeralPeerNegotiationState) -> Void,
+        onUpdateConfiguration: @escaping (EphemeralPeerNegotiationState) async -> Void,
         onFinish: @escaping () -> Void
     ) {
         self.devicePrivateKey = devicePrivateKey
@@ -40,8 +40,8 @@ struct SingleHopEphemeralPeerExchanger: EphemeralPeerExchangingProtocol {
         self.onFinish = onFinish
     }
 
-    func start() {
-        onUpdateConfiguration(.single(EphemeralPeerRelayConfiguration(
+    func start() async {
+        await onUpdateConfiguration(.single(EphemeralPeerRelayConfiguration(
             relay: exit,
             configuration: EphemeralPeerConfiguration(
                 privateKey: devicePrivateKey,
@@ -55,8 +55,8 @@ struct SingleHopEphemeralPeerExchanger: EphemeralPeerExchangingProtocol {
         )
     }
 
-    public func receiveEphemeralPeerPrivateKey(_ ephemeralKey: PrivateKey) {
-        onUpdateConfiguration(.single(EphemeralPeerRelayConfiguration(
+    public func receiveEphemeralPeerPrivateKey(_ ephemeralKey: PrivateKey) async {
+        await onUpdateConfiguration(.single(EphemeralPeerRelayConfiguration(
             relay: exit,
             configuration: EphemeralPeerConfiguration(
                 privateKey: ephemeralKey,
@@ -73,8 +73,8 @@ struct SingleHopEphemeralPeerExchanger: EphemeralPeerExchangingProtocol {
     func receivePostQuantumKey(
         _ preSharedKey: WireGuardKitTypes.PreSharedKey,
         ephemeralKey: WireGuardKitTypes.PrivateKey
-    ) {
-        onUpdateConfiguration(.single(EphemeralPeerRelayConfiguration(
+    ) async {
+        await onUpdateConfiguration(.single(EphemeralPeerRelayConfiguration(
             relay: exit,
             configuration: EphemeralPeerConfiguration(
                 privateKey: ephemeralKey,
diff --git a/ios/PacketTunnelCore/Actor/Protocols/EphemeralPeerExchangingProtocol.swift b/ios/PacketTunnelCore/Actor/Protocols/EphemeralPeerExchangingProtocol.swift
index bffc4f7a21..a0d596fd9e 100644
--- a/ios/PacketTunnelCore/Actor/Protocols/EphemeralPeerExchangingProtocol.swift
+++ b/ios/PacketTunnelCore/Actor/Protocols/EphemeralPeerExchangingProtocol.swift
@@ -9,7 +9,7 @@
 import WireGuardKitTypes
 
 public protocol EphemeralPeerExchangingProtocol {
-    func start()
-    func receivePostQuantumKey(_ preSharedKey: PreSharedKey, ephemeralKey: PrivateKey)
-    func receiveEphemeralPeerPrivateKey(_: PrivateKey)
+    func start() async
+    func receivePostQuantumKey(_ preSharedKey: PreSharedKey, ephemeralKey: PrivateKey) async
+    func receiveEphemeralPeerPrivateKey(_: PrivateKey) async
 }