Change key exchange to be asynchronous

pinkisemils · pinkisemils · commit ee5c10bf7e2c · 2024-12-19T17:05:28.000+01:00
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -13,22 +13,22 @@ 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
         self.onUpdateConfiguration = onUpdateConfiguration
         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
@@ -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
@@ -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
 }

Original file line number	Diff line number	Diff line change
`@@ -13,22 +13,22 @@ import WireGuardKitTypes`
`13`	`13`
`14`	`14`	`final public class EphemeralPeerExchangingPipeline {`
`15`	`15`	`let keyExchanger: EphemeralPeerExchangeActorProtocol`
`16`		`- let onUpdateConfiguration: (EphemeralPeerNegotiationState) -> Void`
	`16`	`+ let onUpdateConfiguration: (EphemeralPeerNegotiationState) async -> Void`
`17`	`17`	`let onFinish: () -> Void`
`18`	`18`
`19`	`19`	`private var ephemeralPeerExchanger: EphemeralPeerExchangingProtocol!`
`20`	`20`
`21`	`21`	`public init(`
`22`	`22`	`_ keyExchanger: EphemeralPeerExchangeActorProtocol,`
`23`		`- onUpdateConfiguration: @escaping (EphemeralPeerNegotiationState) -> Void,`
	`23`	`+ onUpdateConfiguration: @escaping (EphemeralPeerNegotiationState) async -> Void,`
`24`	`24`	`onFinish: @escaping () -> Void`
`25`	`25`	`) {`
`26`	`26`	`self.keyExchanger = keyExchanger`
`27`	`27`	`self.onUpdateConfiguration = onUpdateConfiguration`
`28`	`28`	`self.onFinish = onFinish`
`29`	`29`	`}`
`30`	`30`
`31`		`- public func startNegotiation(_ connectionState: ObservedConnectionState, privateKey: PrivateKey) {`
	`31`	`+ public func startNegotiation(_ connectionState: ObservedConnectionState, privateKey: PrivateKey) async {`
`32`	`32`	`keyExchanger.reset()`
`33`	`33`	`let entryPeer = connectionState.selectedRelays.entry`
`34`	`34`	`let exitPeer = connectionState.selectedRelays.exit`
`@@ -56,14 +56,14 @@ final public class EphemeralPeerExchangingPipeline {`
`56`	`56`	`onFinish: onFinish`
`57`	`57`	`)`
`58`	`58`	`}`
`59`		`- ephemeralPeerExchanger.start()`
	`59`	`+ await ephemeralPeerExchanger.start()`
`60`	`60`	`}`
`61`	`61`
`62`		`- public func receivePostQuantumKey(_ key: PreSharedKey, ephemeralKey: PrivateKey) {`
`63`		`- ephemeralPeerExchanger.receivePostQuantumKey(key, ephemeralKey: ephemeralKey)`
	`62`	`+ public func receivePostQuantumKey(_ key: PreSharedKey, ephemeralKey: PrivateKey) async {`
	`63`	`+ await ephemeralPeerExchanger.receivePostQuantumKey(key, ephemeralKey: ephemeralKey)`
`64`	`64`	`}`
`65`	`65`
`66`		`- public func receiveEphemeralPeerPrivateKey(_ ephemeralPeerPrivateKey: PrivateKey) {`
`67`		`- ephemeralPeerExchanger.receiveEphemeralPeerPrivateKey(ephemeralPeerPrivateKey)`
	`66`	`+ public func receiveEphemeralPeerPrivateKey(_ ephemeralPeerPrivateKey: PrivateKey) async {`
	`67`	`+ await ephemeralPeerExchanger.receiveEphemeralPeerPrivateKey(ephemeralPeerPrivateKey)`
`68`	`68`	`}`
`69`	`69`	`}`