Pass queue for completion callbacks

6 files changed, 134 insertions, 88 deletions

diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index cab698a459..f27bb0566a 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -61,6 +61,8 @@
 		582BB1B3229574F40055B6EF /* SettingsAccountCell.swift in Sources */ = {isa = PBXBuildFile; fileRef = 582BB1B2229574F40055B6EF /* SettingsAccountCell.swift */; };
 		582BB1B52295780F0055B6EF /* AccountExpiry.swift in Sources */ = {isa = PBXBuildFile; fileRef = 582BB1B42295780F0055B6EF /* AccountExpiry.swift */; };
 		5835B7CC233B76CB0096D79F /* TunnelManager.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5835B7CB233B76CB0096D79F /* TunnelManager.swift */; };
+		583BC70724FE4DC500C9DE04 /* Optional+DispatchQueue.swift in Sources */ = {isa = PBXBuildFile; fileRef = 583BC70624FE4DC400C9DE04 /* Optional+DispatchQueue.swift */; };
+		583BC70824FE4DC500C9DE04 /* Optional+DispatchQueue.swift in Sources */ = {isa = PBXBuildFile; fileRef = 583BC70624FE4DC400C9DE04 /* Optional+DispatchQueue.swift */; };
 		5840250122B1124600E4CFEC /* IpAddress+Codable.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5840250022B1124600E4CFEC /* IpAddress+Codable.swift */; };
 		5840250222B1124600E4CFEC /* IpAddress+Codable.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5840250022B1124600E4CFEC /* IpAddress+Codable.swift */; };
 		5840250422B11AB700E4CFEC /* MullvadEndpoint.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5840250322B11AB700E4CFEC /* MullvadEndpoint.swift */; };
@@ -288,6 +290,7 @@
 		582BB1B2229574F40055B6EF /* SettingsAccountCell.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SettingsAccountCell.swift; sourceTree = "<group>"; };
 		582BB1B42295780F0055B6EF /* AccountExpiry.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AccountExpiry.swift; sourceTree = "<group>"; };
 		5835B7CB233B76CB0096D79F /* TunnelManager.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TunnelManager.swift; sourceTree = "<group>"; };
+		583BC70624FE4DC400C9DE04 /* Optional+DispatchQueue.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "Optional+DispatchQueue.swift"; sourceTree = "<group>"; };
 		5840250022B1124600E4CFEC /* IpAddress+Codable.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "IpAddress+Codable.swift"; sourceTree = "<group>"; };
 		5840250322B11AB700E4CFEC /* MullvadEndpoint.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MullvadEndpoint.swift; sourceTree = "<group>"; };
 		5845F841236CBACD00B2D93C /* PacketTunnelIpc.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = PacketTunnelIpc.swift; sourceTree = "<group>"; };
@@ -612,6 +615,7 @@
 				58C6B35322BB87C4003C19AD /* WireguardPrivateKey.swift */,
 				58F3C098249B978C003E76BE /* x25519.c */,
 				58F3C097249B978C003E76BE /* x25519.h */,
+				583BC70624FE4DC400C9DE04 /* Optional+DispatchQueue.swift */,
 			);
 			path = MullvadVPN;
 			sourceTree = "<group>";
@@ -1002,6 +1006,7 @@
 				580EE22124B3240100F9D8A1 /* TransformOperationObserver.swift in Sources */,
 				582BB1AF229566420055B6EF /* SettingsCell.swift in Sources */,
 				5873884D239E6D7E00E96C4E /* EmbeddedViewContainerView.swift in Sources */,
+				583BC70724FE4DC500C9DE04 /* Optional+DispatchQueue.swift in Sources */,
 				58F3C0A4249CB069003E76BE /* HeaderBarView.swift in Sources */,
 				585FE2F124E1365400439C50 /* LogStreamer.swift in Sources */,
 				58B9EB132488ED2100095626 /* AlertPresenter.swift in Sources */,
@@ -1096,6 +1101,7 @@
 				581503A424D6F1EC00C9C50E /* ChainedError+Logger.swift in Sources */,
 				5815039824D6ECAE00C9C50E /* CustomFormatLogHandler.swift in Sources */,
 				5840250522B11AB700E4CFEC /* MullvadEndpoint.swift in Sources */,
+				583BC70824FE4DC500C9DE04 /* Optional+DispatchQueue.swift in Sources */,
 				58BFA5C722A7C97F00A6173D /* RelayCache.swift in Sources */,
 				580EE21024B322E700F9D8A1 /* TransformOperation.swift in Sources */,
 				58906DE02445C7A5002F0673 /* NEProviderStopReason+Debug.swift in Sources */,
diff --git a/ios/MullvadVPN/AutomaticKeyRotationManager.swift b/ios/MullvadVPN/AutomaticKeyRotationManager.swift
index 79e435a2a6..27d480c57c 100644
--- a/ios/MullvadVPN/AutomaticKeyRotationManager.swift
+++ b/ios/MullvadVPN/AutomaticKeyRotationManager.swift
@@ -69,6 +69,9 @@ class AutomaticKeyRotationManager {
     /// A variable backing the `eventHandler` public property
     private var _eventHandler: ((KeyRotationResult) -> Void)?
 
+    /// A dispatch queue used for broadcasting events
+    private let eventQueue: DispatchQueue?
+
     /// An event handler that's invoked when key rotation occurred
     var eventHandler: ((KeyRotationResult) -> Void)? {
         get {
@@ -83,11 +86,12 @@ class AutomaticKeyRotationManager {
         }
     }
 
-    init(persistentKeychainReference: Data) {
+    init(persistentKeychainReference: Data, eventQueue: DispatchQueue?) {
         self.persistentKeychainReference = persistentKeychainReference
+        self.eventQueue = eventQueue
     }
 
-    func startAutomaticRotation(completionHandler: @escaping () -> Void) {
+    func startAutomaticRotation(queue: DispatchQueue?, completionHandler: @escaping () -> Void) {
         dispatchQueue.async {
             guard !self.isAutomaticRotationEnabled else { return }
 
@@ -96,11 +100,11 @@ class AutomaticKeyRotationManager {
             self.isAutomaticRotationEnabled = true
             self.performKeyRotation()
 
-            completionHandler()
+            queue.performOnWrappedOrCurrentQueue(block: completionHandler)
         }
     }
 
-    func stopAutomaticRotation(completionHandler: @escaping () -> Void) {
+    func stopAutomaticRotation(queue: DispatchQueue?, completionHandler: @escaping () -> Void) {
         dispatchQueue.async {
             guard self.isAutomaticRotationEnabled else { return }
 
@@ -113,7 +117,7 @@ class AutomaticKeyRotationManager {
 
             self.timerSource?.cancel()
 
-            completionHandler()
+            queue.performOnWrappedOrCurrentQueue(block: completionHandler)
         }
     }
 
@@ -215,7 +219,9 @@ class AutomaticKeyRotationManager {
             if event.isNew {
                 logger.info("Finished private key rotation")
 
-                eventHandler?(event)
+                eventQueue.performOnWrappedOrCurrentQueue {
+                    self.eventHandler?(event)
+                }
             }
 
             if let rotationDate = Self.nextRotation(creationDate: event.creationDate) {
diff --git a/ios/MullvadVPN/Optional+DispatchQueue.swift b/ios/MullvadVPN/Optional+DispatchQueue.swift
new file mode 100644
index 0000000000..34402977ec
--- /dev/null
+++ b/ios/MullvadVPN/Optional+DispatchQueue.swift
@@ -0,0 +1,22 @@
+//
+//  Optional+DispatchQueue.swift
+//  MullvadVPN
+//
+//  Created by pronebird on 01/09/2020.
+//  Copyright © 2020 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+
+extension Optional where Wrapped == DispatchQueue {
+    /// Unwrap the `DispatchQueue` and perform the block on it, otherwise call the block
+    /// synchronously on the current queue when `Optional` is `.none`.
+    func performOnWrappedOrCurrentQueue(block: @escaping () -> Void) {
+        switch self {
+        case .some(let queue):
+            queue.async(execute: block)
+        case .none:
+            block()
+        }
+    }
+}
diff --git a/ios/MullvadVPN/RelayCache.swift b/ios/MullvadVPN/RelayCache.swift
index 802e278c14..c561303e49 100644
--- a/ios/MullvadVPN/RelayCache.swift
+++ b/ios/MullvadVPN/RelayCache.swift
@@ -110,35 +110,34 @@ class RelayCache {
         self.cacheFileURL = cacheFileURL
     }
 
-    func startPeriodicUpdates(completionHandler: (() -> Void)?) {
+    func startPeriodicUpdates(queue: DispatchQueue?, completionHandler: (() -> Void)?) {
         dispatchQueue.async {
-            guard !self.isPeriodicUpdatesEnabled else {
-                completionHandler?()
-                return
-            }
-
-            self.isPeriodicUpdatesEnabled = true
+            if !self.isPeriodicUpdatesEnabled {
+                self.isPeriodicUpdatesEnabled = true
 
-            switch Self.read(cacheFileURL: self.cacheFileURL) {
-            case .success(let cachedRelayList):
-                if let nextUpdate = Self.nextUpdateDate(lastUpdatedAt: cachedRelayList.updatedAt) {
-                    let startTime = Self.makeWalltime(fromDate: nextUpdate)
-                    self.scheduleRepeatingTimer(startTime: startTime)
-                }
+                switch Self.read(cacheFileURL: self.cacheFileURL) {
+                case .success(let cachedRelayList):
+                    if let nextUpdate = Self.nextUpdateDate(lastUpdatedAt: cachedRelayList.updatedAt) {
+                        let startTime = Self.makeWalltime(fromDate: nextUpdate)
+                        self.scheduleRepeatingTimer(startTime: startTime)
+                    }
 
-            case .failure(let readError):
-                self.logger.error(chainedError: readError, message: "Failed to read the relay cache")
+                case .failure(let readError):
+                    self.logger.error(chainedError: readError, message: "Failed to read the relay cache")
 
-                if Self.shouldDownloadRelaysOnReadFailure(readError) {
-                    self.scheduleRepeatingTimer(startTime: .now())
+                    if Self.shouldDownloadRelaysOnReadFailure(readError) {
+                        self.scheduleRepeatingTimer(startTime: .now())
+                    }
                 }
             }
 
-            completionHandler?()
+            queue.performOnWrappedOrCurrentQueue {
+                completionHandler?()
+            }
         }
     }
 
-    func stopPeriodicUpdates(completionHandler: (() -> Void)?) {
+    func stopPeriodicUpdates(queue: DispatchQueue?, completionHandler: (() -> Void)?) {
         dispatchQueue.async {
             self.isPeriodicUpdatesEnabled = false
 
@@ -146,7 +145,9 @@ class RelayCache {
             self.timerSource = nil
             self.downloadTask?.cancel()
 
-            completionHandler?()
+            queue.performOnWrappedOrCurrentQueue {
+                completionHandler?()
+            }
         }
     }
 
diff --git a/ios/PacketTunnel/PacketTunnelProvider.swift b/ios/PacketTunnel/PacketTunnelProvider.swift
index 94ce8ce42e..3657009909 100644
--- a/ios/PacketTunnel/PacketTunnelProvider.swift
+++ b/ios/PacketTunnel/PacketTunnelProvider.swift
@@ -42,9 +42,11 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
 
     override init() {
         initLoggingSystem(bundleIdentifier: Bundle.main.bundleIdentifier!)
-        WireguardDevice.setTunnelLogger(Logger(label: "WireGuard"))
 
         logger = Logger(label: "PacketTunnelProvider")
+
+        let wireguardLogger = Logger(label: "WireGuard")
+        WireguardDevice.setTunnelLogger(wireguardLogger)
     }
 
     // MARK: - Subclass
@@ -147,43 +149,39 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
                 }
 
                 self.startWireguardDevice(packetFlow: self.packetFlow, configuration: packetTunnelConfig.wireguardConfig) { (result) in
-                    self.dispatchQueue.async {
-                        guard case .success(let device) = result else {
-                            self.tunnelState = .disconnected
+                    guard case .success(let device) = result else {
+                        self.tunnelState = .disconnected
 
-                            completionHandler(result.map { _ in () })
-                            return
-                        }
+                        completionHandler(result.map { _ in () })
+                        return
+                    }
+
+                    let persistentKeychainReference = packetTunnelConfig.persistentKeychainReference
+                    let keyRotationManager = AutomaticKeyRotationManager(persistentKeychainReference: persistentKeychainReference, eventQueue: self.dispatchQueue)
+                    keyRotationManager.eventHandler = { [weak self] (keyRotationEvent) in
+                        guard let self = self else { return }
 
-                        let persistentKeychainReference = packetTunnelConfig.persistentKeychainReference
-                        let keyRotationManager = AutomaticKeyRotationManager(persistentKeychainReference: persistentKeychainReference)
-                        keyRotationManager.eventHandler = { (keyRotationEvent) in
-                            self.dispatchQueue.async {
-                                self.reloadTunnelSettings { (result) in
-                                    switch result {
-                                    case .success:
-                                        break
+                        self.reloadTunnelSettings { (result) in
+                            switch result {
+                            case .success:
+                                break
 
-                                    case .failure(let error):
-                                        self.logger.error(chainedError: error, message: "Failed to reload tunnel settings")
-                                    }
-                                }
+                            case .failure(let error):
+                                self.logger.error(chainedError: error, message: "Failed to reload tunnel settings")
                             }
                         }
+                    }
 
-                        RelayCache.shared.startPeriodicUpdates {
-                            keyRotationManager.startAutomaticRotation {
-                                self.dispatchQueue.async {
-                                    let context = PacketTunnelContext(
-                                        wireguardDevice: device,
-                                        keyRotationManager: keyRotationManager
-                                    )
+                    RelayCache.shared.startPeriodicUpdates(queue: self.dispatchQueue) {
+                        keyRotationManager.startAutomaticRotation(queue: self.dispatchQueue) {
+                            let context = PacketTunnelContext(
+                                wireguardDevice: device,
+                                keyRotationManager: keyRotationManager
+                            )
 
-                                    self.tunnelState = .connected(packetTunnelConfig.selectorResult.tunnelConnectionInfo, context)
+                            self.tunnelState = .connected(packetTunnelConfig.selectorResult.tunnelConnectionInfo, context)
 
-                                    completionHandler(.success(()))
-                                }
-                            }
+                            completionHandler(.success(()))
                         }
                     }
                 }
@@ -200,18 +198,16 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
 
         self.tunnelState = .disconnecting
 
-        RelayCache.shared.stopPeriodicUpdates {
-            context.keyRotationManager.stopAutomaticRotation {
-                context.wireguardDevice.stop { (result) in
-                    self.dispatchQueue.async {
-                        let result = result.mapError({ (error) -> PacketTunnelProviderError in
-                            return .stopWireguardDevice(error)
-                        })
+        RelayCache.shared.stopPeriodicUpdates(queue: self.dispatchQueue) {
+            context.keyRotationManager.stopAutomaticRotation(queue: self.dispatchQueue) {
+                context.wireguardDevice.stop(queue: self.dispatchQueue) { (result) in
+                    let result = result.mapError({ (error) -> PacketTunnelProviderError in
+                        return .stopWireguardDevice(error)
+                    })
 
-                        self.tunnelState = .disconnected
+                    self.tunnelState = .disconnected
 
-                        completionHandler(result)
-                    }
+                    completionHandler(result)
                 }
             }
         }
@@ -262,10 +258,8 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
                     return
                 }
 
-                context.wireguardDevice.setConfiguration(packetTunnelConfig.wireguardConfig) { (result) in
-                    self.dispatchQueue.async {
-                        finishReconnecting(result.mapError { PacketTunnelProviderError.updateWireguardConfiguration($0) })
-                    }
+                context.wireguardDevice.setConfiguration(packetTunnelConfig.wireguardConfig, queue: self.dispatchQueue) { (result) in
+                    finishReconnecting(result.mapError { PacketTunnelProviderError.updateWireguardConfiguration($0) })
                 }
             }
         }
@@ -297,10 +291,8 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
             return
         }
 
-        Self.makePacketTunnelConfig(keychainReference: keychainReference) { (result) in
-            self.dispatchQueue.async {
-                completionHandler(result)
-            }
+        Self.makePacketTunnelConfig(keychainReference: keychainReference, queue: self.dispatchQueue) { (result) in
+            completionHandler(result)
         }
     }
 
@@ -339,7 +331,7 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
     }
 
     /// Returns a `PacketTunnelConfig` that contains the tunnel settings and selected relay
-    private class func makePacketTunnelConfig(keychainReference: Data, completionHandler: @escaping (Result<PacketTunnelConfiguration, PacketTunnelProviderError>) -> Void) {
+    private class func makePacketTunnelConfig(keychainReference: Data, queue: DispatchQueue?, completionHandler: @escaping (Result<PacketTunnelConfiguration, PacketTunnelProviderError>) -> Void) {
         switch Self.readTunnelSettings(keychainReference: keychainReference) {
         case .success(let tunnelSettings):
             Self.selectRelayEndpoint(relayConstraints: tunnelSettings.relayConstraints) { (result) in
@@ -350,11 +342,16 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
                         selectorResult: selectorResult
                     )
                 }
-                completionHandler(result)
+
+                queue.performOnWrappedOrCurrentQueue {
+                    completionHandler(result)
+                }
             }
 
         case .failure(let error):
-            completionHandler(.failure(error))
+            queue.performOnWrappedOrCurrentQueue {
+                completionHandler(.failure(error))
+            }
         }
     }
 
@@ -397,7 +394,7 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
 
         logger.info("Tunnel interface is \(tunnelDeviceName)")
 
-        device.start(configuration: configuration) { (result) in
+        device.start(queue: dispatchQueue, configuration: configuration) { (result) in
             let result = result.map { device }
                 .mapError { PacketTunnelProviderError.startWireguardDevice($0) }
 
diff --git a/ios/PacketTunnel/WireguardDevice.swift b/ios/PacketTunnel/WireguardDevice.swift
index 77ed1c5553..2f46f76f84 100644
--- a/ios/PacketTunnel/WireguardDevice.swift
+++ b/ios/PacketTunnel/WireguardDevice.swift
@@ -141,10 +141,12 @@ class WireguardDevice {
 
     // MARK: - Public methods
 
-    func start(configuration: WireguardConfiguration, completionHandler: @escaping (Result<(), Error>) -> Void) {
+    func start(queue: DispatchQueue?, configuration: WireguardConfiguration, completionHandler: @escaping (Result<(), Error>) -> Void) {
         workQueue.async {
             guard !self.isStarted else {
-                completionHandler(.failure(.alreadyStarted))
+                queue.performOnWrappedOrCurrentQueue {
+                    completionHandler(.failure(.alreadyStarted))
+                }
                 return
             }
 
@@ -160,15 +162,19 @@ class WireguardDevice {
 
                 self.startNetworkMonitor()
 
-                completionHandler(.success(()))
+                queue.performOnWrappedOrCurrentQueue {
+                    completionHandler(.success(()))
+                }
 
             case .failure(let error):
-                completionHandler(.failure(error))
+                queue.performOnWrappedOrCurrentQueue {
+                    completionHandler(.failure(error))
+                }
             }
         }
     }
 
-    func stop(completionHandler: @escaping (Result<(), Error>) -> Void) {
+    func stop(queue: DispatchQueue?, completionHandler: @escaping (Result<(), Error>) -> Void) {
         workQueue.async {
             if self.isStarted {
                 self.networkMonitor?.cancel()
@@ -177,14 +183,18 @@ class WireguardDevice {
                 self.stopWireguardBackend()
                 self.isStarted = false
 
-                completionHandler(.success(()))
+                queue.performOnWrappedOrCurrentQueue {
+                    completionHandler(.success(()))
+                }
             } else {
-                completionHandler(.failure(.notStarted))
+                queue.performOnWrappedOrCurrentQueue {
+                    completionHandler(.failure(.notStarted))
+                }
             }
         }
     }
 
-    func setConfiguration(_ newConfiguration: WireguardConfiguration, completionHandler: @escaping (Result<(), Error>) -> Void) {
+    func setConfiguration(_ newConfiguration: WireguardConfiguration, queue: DispatchQueue?, completionHandler: @escaping (Result<(), Error>) -> Void) {
         workQueue.async {
             if self.isStarted {
                 if let handle = self.wireguardHandle {
@@ -196,9 +206,13 @@ class WireguardDevice {
 
                 self.configuration = newConfiguration
 
-                completionHandler(.success(()))
+                queue.performOnWrappedOrCurrentQueue {
+                    completionHandler(.success(()))
+                }
             } else {
-                completionHandler(.failure(.notStarted))
+                queue.performOnWrappedOrCurrentQueue {
+                    completionHandler(.failure(.notStarted))
+                }
             }
         }
     }