Propagate Wireguard public key via TunnelManager.publicKey and keep it up to date

2 files changed, 70 insertions, 70 deletions

diff --git a/ios/MullvadVPN/TunnelManager.swift b/ios/MullvadVPN/TunnelManager.swift
index 08158e0c2d..8103edf373 100644
--- a/ios/MullvadVPN/TunnelManager.swift
+++ b/ios/MullvadVPN/TunnelManager.swift
@@ -309,6 +309,9 @@ class TunnelManager {
 
     @Published private(set) var tunnelState = TunnelState.disconnected
 
+    /// A last known public key
+    @Published private(set) var publicKey: WireguardPublicKey?
+
     /// Initialize the TunnelManager with the tunnel from the system
     ///
     /// The given account token is used to ensure that the system tunnel was configured for the same
@@ -320,9 +323,12 @@ class TunnelManager {
                 .receive(on: self.executionQueue)
                 .flatMap { (tunnels) -> AnyPublisher<(), LoadTunnelError> in
 
-                    // Migrate tunnel configuration if needed
                     if let accountToken = accountToken {
+                        // Migrate tunnel configuration if needed
                         self.migrateTunnelConfiguration(accountToken: accountToken)
+
+                        // Load last known public key
+                        self.loadPublicKey(accountToken: accountToken)
                     }
 
                     // No tunnels found. Save the account token.
@@ -364,6 +370,11 @@ class TunnelManager {
             if let status = self.tunnelProvider?.connection.status {
                 self.updateTunnelState(connectionStatus: status)
             }
+
+            // Reload the last known public key
+            if let accountToken = self.accountToken {
+                self.loadPublicKey(accountToken: accountToken)
+            }
             return Result.Publisher(()).eraseToAnyPublisher()
         }.eraseToAnyPublisher()
     }
@@ -421,15 +432,21 @@ class TunnelManager {
                             .mapError { SetAccountError.setup($0) }
                     }
 
+                    let publicKey = tunnelConfig.interface.privateKey.publicKey
+
+                    // Save the last known public key
+                    self.publicKey = publicKey
+
                     // Make sure to avoid pushing the wireguard keys when addresses are assigned
                     guard tunnelConfig.interface.addresses.isEmpty else {
                         return setupTunnelPublisher.eraseToAnyPublisher()
                     }
 
                     // Send wireguard key to the server
-                    let publicKey = tunnelConfig.interface.privateKey.publicKey.rawRepresentation
-
-                    return self.rpc.pushWireguardKey(accountToken: accountToken, publicKey: publicKey)
+                    return self.rpc.pushWireguardKey(
+                        accountToken: accountToken,
+                        publicKey: publicKey.rawRepresentation
+                    )
                         .mapError { SetAccountError.pushWireguardKey($0) }
                         .flatMap { (addresses) in
                             self.updateAssociatedAddresses(
@@ -525,6 +542,7 @@ class TunnelManager {
             .handleEvents(receiveCompletion: { (completion) in
                 if case .finished = completion {
                     self.accountToken = nil
+                    self.publicKey = nil
                     self.tunnelProvider = nil
                     self.tunnelIpc = nil
 
@@ -572,9 +590,12 @@ class TunnelManager {
                                     .map { _ in () }
                                     .publisher
                             }.receive(on: self.executionQueue)
-                                .flatMap { _ in
+                                .flatMap { _ -> AnyPublisher<(), RegenerateWireguardPrivateKeyError> in
+                                    // Save new public key
+                                    self.publicKey = newPrivateKey.publicKey
+
                                     // Ignore Packet Tunnel IPC errors but log them
-                                    self.reloadPacketTunnelConfiguration()
+                                    return self.reloadPacketTunnelConfiguration()
                                         .handleEvents(receiveCompletion: { (completion) in
                                             if case .failure(let error) = completion {
                                                 os_log(.error, "Failed to tell the tunnel to reload configuration: %{public}s", error.localizedDescription)
@@ -582,6 +603,7 @@ class TunnelManager {
                                         })
                                         .replaceError(with: ())
                                         .setFailureType(to: RegenerateWireguardPrivateKeyError.self)
+                                        .eraseToAnyPublisher()
                             }
                     }
                     .mapError { TunnelManagerError.regenerateWireguardPrivateKey($0) }
@@ -637,20 +659,6 @@ class TunnelManager {
         }.eraseToAnyPublisher()
     }
 
-    func getWireguardPublicKey() -> AnyPublisher<WireguardPublicKey, TunnelManagerError> {
-        MutuallyExclusive(exclusivityQueue: exclusivityQueue, executionQueue: executionQueue) {
-            Just(self.accountToken)
-                .setFailureType(to: TunnelManagerError.self)
-                .replaceNil(with: .missingAccount)
-                .flatMap { (accountToken) in
-                    TunnelConfigurationManager.load(searchTerm: .accountToken(accountToken))
-                        .map { $0.tunnelConfiguration.interface.privateKey.publicKey }
-                        .mapError { .getWireguardPublicKey($0) }
-                        .publisher
-            }
-        }.eraseToAnyPublisher()
-    }
-
     // MARK: - Private
 
     /// Tell Packet Tunnel process to reload the tunnel configuration
@@ -708,6 +716,18 @@ class TunnelManager {
         updateTunnelState(connectionStatus: connection.status)
     }
 
+    private func loadPublicKey(accountToken: String) {
+        switch TunnelConfigurationManager.load(searchTerm: .accountToken(accountToken)) {
+        case .success(let entry):
+            self.publicKey = entry.tunnelConfiguration.interface.privateKey.publicKey
+
+        case .failure(let error):
+            os_log(.error, "Failed to load the public key: %{public}s", error.localizedDescription)
+
+            self.publicKey = nil
+        }
+    }
+
     /// Initiates the `tunnelState` update
     private func updateTunnelState(connectionStatus: NEVPNStatus) {
         os_log(.default, "VPN Status: %{public}s", "\(connectionStatus)")
@@ -751,6 +771,12 @@ class TunnelManager {
                         .eraseToAnyPublisher()
 
                 case .reasserting:
+                    // Refresh the last known public key on reconnect to cover the possibility of
+                    // the key being changed due to key rotation.
+                    if let accountToken = self.accountToken {
+                        self.loadPublicKey(accountToken: accountToken)
+                    }
+
                     return self.getTunnelConnectionInfo()
                         .mapError { .ipcRequest($0) }
                         .map { .reconnecting($0) }
diff --git a/ios/MullvadVPN/WireguardKeysViewController.swift b/ios/MullvadVPN/WireguardKeysViewController.swift
index 5b01dd541c..c677c75eee 100644
--- a/ios/MullvadVPN/WireguardKeysViewController.swift
+++ b/ios/MullvadVPN/WireguardKeysViewController.swift
@@ -59,7 +59,7 @@ class WireguardKeysViewController: UIViewController {
     @IBOutlet var verifyKeyButton: UIButton!
     @IBOutlet var wireguardKeyStatusView: WireguardKeyStatusView!
 
-    private var tunnelStateSubscriber: AnyCancellable?
+    private var publicKeySubscriber: AnyCancellable?
     private var loadKeySubscriber: AnyCancellable?
     private var verifyKeySubscriber: AnyCancellable?
     private var regenerateKeySubscriber: AnyCancellable?
@@ -67,7 +67,6 @@ class WireguardKeysViewController: UIViewController {
     private var copyToPasteboardSubscriber: AnyCancellable?
 
     private let rpc = MullvadRpc()
-    private var publicKey: WireguardPublicKey?
 
     private var state: WireguardKeysViewState = .default {
         didSet {
@@ -78,39 +77,29 @@ class WireguardKeysViewController: UIViewController {
     override func viewDidLoad() {
         super.viewDidLoad()
 
-        // Reset Storyboard placeholders
-        setPublicKeyTitle(string: "-", animated: false)
-        creationDateLabel.text = "-"
-
         creationDateTimerSubscriber = Timer.publish(every: kCreationDateRefreshInterval, on: .main, in: .common)
             .autoconnect()
             .sink { [weak self] _ in
-                guard let self = self else { return }
+                let publicKey = TunnelManager.shared.publicKey
 
-                if let creationDate = self.publicKey?.creationDate {
-                    self.updateCreationDateLabel(with: creationDate)
-                }
+                self?.updatePublicKey(publicKey: publicKey, animated: true)
         }
 
-        tunnelStateSubscriber = TunnelManager.shared.$tunnelState
+        publicKeySubscriber = TunnelManager.shared.$publicKey
+            .dropFirst()
             .receive(on: DispatchQueue.main)
-            .sink(receiveValue: {  [weak self] (tunnelState) in
-                guard let self = self else { return }
-
-                // Reload the public key when the tunnel is reconnecting
-                // Normally this may happen in response to private key change
-                if case .reconnecting = tunnelState {
-                    self.loadPublicKey(animated: true)
-                }
+            .sink(receiveValue: {  [weak self] (publicKey) in
+                self?.updatePublicKey(publicKey: publicKey, animated: true)
             })
 
-        loadPublicKey(animated: false)
+        // Set public key title without animation
+        updatePublicKey(publicKey: TunnelManager.shared.publicKey, animated: false)
     }
 
     // MARK: - IBActions
 
     @IBAction func copyPublicKey(_ sender: Any) {
-        guard let publicKey = self.publicKey else { return }
+        guard let publicKey = TunnelManager.shared.publicKey else { return }
 
         UIPasteboard.general.string = publicKey.stringRepresentation()
 
@@ -121,12 +110,11 @@ class WireguardKeysViewController: UIViewController {
         copyToPasteboardSubscriber =
             Just(()).cancellableDelay(for: .seconds(3), scheduler: DispatchQueue.main)
                 .sink(receiveValue: { [weak self] () in
-                    guard let self = self, let publicKey = self.publicKey else { return }
+                    guard let self = self else { return }
 
-                    let displayKey = publicKey
-                        .stringRepresentation(maxLength: kDisplayPublicKeyMaxLength)
+                    let publicKey = TunnelManager.shared.publicKey
 
-                    self.setPublicKeyTitle(string: displayKey, animated: true)
+                    self.updatePublicKey(publicKey: publicKey, animated: true)
                 })
     }
 
@@ -136,7 +124,7 @@ class WireguardKeysViewController: UIViewController {
 
     @IBAction func handleVerifyKey(_ sender: Any) {
         guard let accountToken = Account.shared.token,
-            let publicKey = publicKey else { return }
+            let publicKey = TunnelManager.shared.publicKey else { return }
 
         verifyKey(accountToken: accountToken, publicKey: publicKey)
     }
@@ -157,30 +145,16 @@ class WireguardKeysViewController: UIViewController {
         creationDateLabel.text = formatKeyGenerationElapsedTime(with: creationDate) ?? "-"
     }
 
-    private func loadPublicKey(animated: Bool) {
-        loadKeySubscriber = TunnelManager.shared.getWireguardPublicKey()
-            .receive(on: DispatchQueue.main)
-            .sink(receiveCompletion: { (completion) in
-                switch completion {
-                case .finished:
-                    break
-
-                case .failure(let error):
-                    os_log(.error, "Failed to receive the public key for Wireguard: %{public}s",
-                           error.localizedDescription)
-
-                    self.presentError(error, preferredStyle: .alert)
-                }
-            }) { [weak self] (publicKey) in
-                guard let self = self else { return }
+    private func updatePublicKey(publicKey: WireguardPublicKey?, animated: Bool) {
+        if let publicKey = publicKey {
+            let displayKey = publicKey
+                .stringRepresentation(maxLength: kDisplayPublicKeyMaxLength)
 
-                let displayKey = publicKey
-                    .stringRepresentation(maxLength: kDisplayPublicKeyMaxLength)
-
-                self.setPublicKeyTitle(string: displayKey, animated: animated)
-                self.updateCreationDateLabel(with: publicKey.creationDate)
-
-                self.publicKey = publicKey
+            setPublicKeyTitle(string: displayKey, animated: animated)
+            updateCreationDateLabel(with: publicKey.creationDate)
+        } else {
+            setPublicKeyTitle(string: "-", animated: animated)
+            creationDateLabel.text = "-"
         }
     }
 
@@ -245,7 +219,7 @@ class WireguardKeysViewController: UIViewController {
             .sink { (completion) in
                 switch completion {
                 case .finished:
-                    self.loadPublicKey(animated: true)
+                    break
 
                 case .failure(let error):
                     os_log(.error, "Failed to re-generate the private key: %{public}s",