Add private key rotation

7 files changed, 612 insertions, 309 deletions

diff --git a/ios/MullvadVPN/AutomaticKeyRotationManager.swift b/ios/MullvadVPN/AutomaticKeyRotationManager.swift
new file mode 100644
index 0000000000..ee7cc6e4c7
--- /dev/null
+++ b/ios/MullvadVPN/AutomaticKeyRotationManager.swift
@@ -0,0 +1,221 @@
+//
+//  AutomaticKeyRotationManager.swift
+//  MullvadVPN
+//
+//  Created by pronebird on 05/05/2020.
+//  Copyright © 2020 Mullvad VPN AB. All rights reserved.
+//
+
+import Combine
+import Foundation
+import os
+
+/// A private key rotation retry interval on failure (in seconds)
+private let retryIntervalOnFailure = 300
+
+/// A private key rotation interval (in days)
+private let rotationInterval = 1
+
+class AutomaticKeyRotationManager {
+
+    enum Error: Swift.Error {
+        /// An RPC failure
+        case rpc(MullvadRpc.Error)
+
+        /// A failure to read the tunnel configuration
+        case readTunnelConfiguration(TunnelConfigurationManager.Error)
+
+        /// A failure to update tunnel configuration
+        case updateTunnelConfiguration(TunnelConfigurationManager.Error)
+
+        var localizedDescription: String {
+            switch self {
+            case .rpc(let error):
+                return "Rpc error: \(error.localizedDescription)"
+            case .readTunnelConfiguration(let error):
+                return "Read configuration error: \(error.localizedDescription)"
+            case .updateTunnelConfiguration(let error):
+                return "Update configuration error: \(error.localizedDescription)"
+            }
+        }
+    }
+
+    struct KeyRotationEvent {
+        var isNew: Bool
+        var creationDate: Date
+        var publicKey: WireguardPublicKey
+    }
+
+    private let rpc = MullvadRpc()
+    private let persistentKeychainReference: Data
+    private var rotateKeySubscriber: AnyCancellable?
+
+    private let dispatchQueue = DispatchQueue(label: "net.mullvad.vpn.key-manager", qos: .background)
+    private var retryWorkItem: DispatchWorkItem?
+
+    private var isAutomaticRotationEnabled = false
+
+    private let lock = NSLock()
+    private var _keyRotationEventHandler: ((KeyRotationEvent) -> Void)?
+    var keyRotationEventHandler: ((KeyRotationEvent) -> Void)? {
+        get {
+            lock.withCriticalBlock {
+                self._keyRotationEventHandler
+            }
+        }
+        set {
+            lock.withCriticalBlock {
+                self._keyRotationEventHandler = newValue
+            }
+        }
+    }
+
+    init(persistentKeychainReference: Data) {
+        self.persistentKeychainReference = persistentKeychainReference
+    }
+
+    func startAutomaticRotation() {
+        dispatchQueue.async {
+            guard !self.isAutomaticRotationEnabled else { return }
+
+            os_log(.default, log: tunnelProviderLog, "Start automatic key rotation")
+
+            self.isAutomaticRotationEnabled = true
+            self.performKeyRotation()
+        }
+    }
+
+    func stopAutomaticRotation() {
+        dispatchQueue.async {
+            guard self.isAutomaticRotationEnabled else { return }
+
+            os_log(.default, log: tunnelProviderLog, "Stop automatic key rotation")
+
+            self.isAutomaticRotationEnabled = false
+            self.rotateKeySubscriber?.cancel()
+            self.retryWorkItem?.cancel()
+        }
+    }
+
+    private func performKeyRotation() {
+        rotateKeySubscriber = tryRotatingPrivateKey()
+            .receive(on: dispatchQueue)
+            .sink(receiveCompletion: { [weak self] (completion) in
+                guard let self = self else { return }
+
+                switch completion {
+                case .finished:
+                    break
+
+                case .failure(let error):
+                    os_log(.error, log: tunnelProviderLog,
+                           "Failed to rotate the private key: %{public}s. Retry in %d seconds.",
+                           error.localizedDescription,
+                           retryIntervalOnFailure)
+
+                    self.scheduleRetry(deadline: .now() + .seconds(retryIntervalOnFailure))
+                }
+            }) { [weak self] (keyRotationEvent) in
+                guard let self = self else { return }
+
+                if keyRotationEvent.isNew {
+                    os_log(.default, log: tunnelProviderLog, "Finished private key rotation")
+
+                    self.keyRotationEventHandler?(keyRotationEvent)
+                }
+
+                if let rotationDate = Self.nextRotation(creationDate: keyRotationEvent.creationDate) {
+                    let interval = rotationDate.timeIntervalSinceNow
+
+                    os_log(.default, log: tunnelProviderLog,
+                           "Next private key rotation on %{public}s", "\(rotationDate)")
+
+                    self.scheduleRetry(deadline: .now() + .seconds(Int(interval)))
+                } else {
+                    os_log(.error, log: tunnelProviderLog,
+                           "Failed to compute the next private rotation date. Retry in %d seconds.")
+
+                    self.scheduleRetry(deadline: .now() + .seconds(retryIntervalOnFailure))
+                }
+        }
+    }
+
+    private func scheduleRetry(deadline: DispatchTime) {
+        // cancel any pending work
+        retryWorkItem?.cancel()
+
+        // schedule new work
+        let workItem = DispatchWorkItem { [weak self] in
+            self?.performKeyRotation()
+        }
+        retryWorkItem = workItem
+        dispatchQueue.asyncAfter(deadline: deadline, execute: workItem)
+    }
+
+    private func tryRotatingPrivateKey() -> AnyPublisher<KeyRotationEvent, Error> {
+        return TunnelConfigurationManager
+            .load(searchTerm: .persistentReference(persistentKeychainReference))
+            .mapError { .readTunnelConfiguration($0) }
+            .publisher
+            .flatMap { (keychainEntry) -> AnyPublisher<KeyRotationEvent, Error> in
+                let currentPrivateKey = keychainEntry.tunnelConfiguration.interface.privateKey
+
+                if Self.shouldRotateKey(creationDate: currentPrivateKey.creationDate) {
+                    return self.replaceWireguardKey(
+                        accountToken: keychainEntry.accountToken,
+                        oldPublicKey: currentPrivateKey.publicKey
+                    ).map({ (newTunnelConfiguration) -> KeyRotationEvent in
+                        let newPrivateKey = newTunnelConfiguration.interface.privateKey
+
+                        return KeyRotationEvent(
+                            isNew: true,
+                            creationDate: newPrivateKey.creationDate,
+                            publicKey: newPrivateKey.publicKey
+                        )
+                    }).eraseToAnyPublisher()
+                } else {
+                    let result = KeyRotationEvent(
+                        isNew: false,
+                        creationDate: currentPrivateKey.creationDate,
+                        publicKey: currentPrivateKey.publicKey
+                    )
+
+                    return Result.Publisher(result).eraseToAnyPublisher()
+                }
+        }.eraseToAnyPublisher()
+    }
+
+    private func replaceWireguardKey(accountToken: String, oldPublicKey: WireguardPublicKey)
+        -> AnyPublisher<TunnelConfiguration, Error>
+    {
+        let newPrivateKey = WireguardPrivateKey()
+
+        return rpc.replaceWireguardKey(
+            accountToken: accountToken,
+            oldPublicKey: oldPublicKey.rawRepresentation,
+            newPublicKey: newPrivateKey.publicKey.rawRepresentation)
+            .mapError {  .rpc($0) }
+            .flatMap { (addresses) in
+                TunnelConfigurationManager
+                    .update(searchTerm: .persistentReference(self.persistentKeychainReference))
+                    { (tunnelConfiguration) in
+                        tunnelConfiguration.interface.privateKey = newPrivateKey
+                        tunnelConfiguration.interface.addresses = [
+                            addresses.ipv4Address,
+                            addresses.ipv6Address
+                        ]
+                }
+                .mapError { .updateTunnelConfiguration($0) }
+                .publisher
+        }.eraseToAnyPublisher()
+    }
+
+    class func nextRotation(creationDate: Date) -> Date? {
+        return Calendar.current.date(byAdding: .day, value: rotationInterval, to: creationDate)
+    }
+
+    class func shouldRotateKey(creationDate: Date) -> Bool {
+        return nextRotation(creationDate: creationDate)
+            .map { $0 <= Date() } ?? false
+    }
+}
diff --git a/ios/MullvadVPN/KeychainItemRevision.swift b/ios/MullvadVPN/KeychainItemRevision.swift
new file mode 100644
index 0000000000..6e73d1a6a8
--- /dev/null
+++ b/ios/MullvadVPN/KeychainItemRevision.swift
@@ -0,0 +1,70 @@
+//
+//  KeychainItemRevision.swift
+//  MullvadVPN
+//
+//  Created by pronebird on 07/05/2020.
+//  Copyright © 2020 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+
+/// A struct that helps to organize revisions for Keychain items
+/// Uses `Keychain.Attributes.generic` field for storing `Int64` revision number.
+struct KeychainItemRevision {
+    let revision: Int64
+
+    /// Returns the `KeychainItemRevision` initialized with the next in sequence revision
+    var nextRevision: KeychainItemRevision {
+        return .init(revision: revision + 1)
+    }
+
+    /// Initialize a struct with the given revision number
+    init(revision: Int64) {
+        self.revision = revision
+    }
+
+    /// Initialize a struct from the given `Keychain.Attributes`
+    ///
+    /// Returns `nil` when `Keychain.Attributes.generic` is `nil` or when it's set to data that
+    /// cannot be interpreted as `Int64`
+    init?(attributes: Keychain.Attributes) {
+        let revision = attributes.generic.flatMap { Self.parseRevision(from: $0) }
+        if let revision = revision {
+            self.revision = revision
+        } else {
+            return nil
+        }
+    }
+
+    /// Store the revision number in the given `Keychain.Attributes`
+    func store(in attributes: inout Keychain.Attributes) {
+        attributes.generic = asData()
+    }
+
+    /// A convenience method to initialize the first revision in sequence
+    static func firstRevision() -> Self {
+        return .init(revision: 1)
+    }
+
+    /// Serialize the revision number as `Data`
+    private func asData() -> Data {
+        return withUnsafeBytes(of: revision) { Data($0) }
+    }
+
+    /// Private helper to parse `Data` into `Int64`
+    /// Returns `nil` if the given raw data does not match the size of `Int64`
+    private static func parseRevision(from rawData: Data) -> Int64? {
+        var value: Int64 = 0
+
+        // Ensure that the buffer has as many bytes as needed to fill the `Int64`
+        guard rawData.count == MemoryLayout.size(ofValue: value) else {
+            return nil
+        }
+
+        _ = withUnsafeMutableBytes(of: &value) { (valuePointer) in
+            rawData.copyBytes(to: valuePointer)
+        }
+
+        return value
+    }
+}
diff --git a/ios/MullvadVPN/TunnelConfigurationCoder.swift b/ios/MullvadVPN/TunnelConfigurationCoder.swift
deleted file mode 100644
index 97500eccbb..0000000000
--- a/ios/MullvadVPN/TunnelConfigurationCoder.swift
+++ /dev/null
@@ -1,30 +0,0 @@
-//
-//  TunnelConfigurationCoder.swift
-//  MullvadVPN
-//
-//  Created by pronebird on 02/10/2019.
-//  Copyright © 2019 Mullvad VPN AB. All rights reserved.
-//
-
-import Foundation
-
-/// Tunnel configuration encoding and decoding helper
-enum TunnelConfigurationCoder {}
-
-extension TunnelConfigurationCoder {
-
-    enum Error: Swift.Error {
-        case encode(Swift.Error)
-        case decode(Swift.Error)
-    }
-
-    static func decode(data: Data) -> Result<TunnelConfiguration, Error> {
-        return Result { try JSONDecoder().decode(TunnelConfiguration.self, from: data) }
-            .mapError { Error.decode($0) }
-    }
-
-    static func encode(tunnelConfig: TunnelConfiguration) -> Result<Data, Error> {
-        return Result { try JSONEncoder().encode(tunnelConfig) }
-            .mapError { Error.encode($0) }
-    }
-}
diff --git a/ios/MullvadVPN/TunnelConfigurationManager.swift b/ios/MullvadVPN/TunnelConfigurationManager.swift
index 563216a202..29b30befa2 100644
--- a/ios/MullvadVPN/TunnelConfigurationManager.swift
+++ b/ios/MullvadVPN/TunnelConfigurationManager.swift
@@ -12,203 +12,247 @@ import Security
 /// Service name used for keychain items
 private let kServiceName = "Mullvad VPN"
 
-enum TunnelConfigurationManagerError: Error {
-    case encode(TunnelConfigurationCoder.Error)
-    case decode(TunnelConfigurationCoder.Error)
-    case addToKeychain(KeychainError)
-    case updateKeychain(KeychainError)
-    case removeKeychainItem(KeychainError)
-    case getFromKeychain(KeychainError)
-    case getPersistentKeychainRef(KeychainError)
-}
-
 enum TunnelConfigurationManager {}
 
 extension TunnelConfigurationManager {
 
-    static func save(configuration: TunnelConfiguration, account: String) -> Result<(), TunnelConfigurationManagerError> {
-        TunnelConfigurationCoder.encode(tunnelConfig: configuration)
-            .mapError { .encode($0) }
-            .flatMap { (data) -> Result<(), TunnelConfigurationManagerError> in
-                Keychain.updateItem(account: account, data: data)
-                    .flatMapError { (keychainError) -> Result<(), TunnelConfigurationManagerError> in
-                        if case .itemNotFound = keychainError {
-                            return Keychain.addItem(account: account, data: data)
-                                .mapError { .addToKeychain($0) }
-                        } else {
-                            return .failure(.updateKeychain(keychainError))
-                        }
-                }
-        }
-    }
+    enum Error: Swift.Error {
+        case encode(Swift.Error)
+        case decode(Swift.Error)
+        case addToKeychain(Keychain.Error)
+        case updateKeychain(Keychain.Error)
+        case removeKeychainItem(Keychain.Error)
+        case getFromKeychain(Keychain.Error)
+        case getPersistentKeychainRef(Keychain.Error)
 
-    static func load(account: String) -> Result<TunnelConfiguration, TunnelConfigurationManagerError> {
-        Keychain.getItemData(account: account)
-            .mapError { .getFromKeychain($0) }
-            .flatMap { (data) in
-                TunnelConfigurationCoder.decode(data: data)
-                    .mapError { .decode($0) }
+        var localizedDescription: String {
+            switch self {
+            case .encode(let error):
+                return error.localizedDescription
+            case .decode(let error):
+                return error.localizedDescription
+            case .addToKeychain(let error):
+                return error.localizedDescription
+            case .updateKeychain(let error):
+                return error.localizedDescription
+            case .removeKeychainItem(let error):
+                return error.localizedDescription
+            case .getFromKeychain(let error):
+                return error.localizedDescription
+            case .getPersistentKeychainRef(let error):
+                return error.localizedDescription
+            }
         }
     }
 
-    static func load(persistentKeychainRef: Data) -> Result<TunnelConfiguration, TunnelConfigurationManagerError> {
-        Keychain.getItemData(persistentKeychainRef: persistentKeychainRef)
-            .mapError { .getFromKeychain($0) }
-            .flatMap { (data) in
-                TunnelConfigurationCoder.decode(data: data)
-                    .mapError { .decode($0) }
+    typealias Result<T> = Swift.Result<T, Error>
+
+    /// Keychain access level that should be used for all items containing tunnel configuration
+    private static let keychainAccessibleLevel = Keychain.Accessible.afterFirstUnlock
+
+    enum KeychainSearchTerm {
+        case accountToken(String)
+        case persistentReference(Data)
+
+        /// Returns `Keychain.Attributes` appropriate for adding or querying the item
+        fileprivate func makeKeychainAttributes() -> Keychain.Attributes {
+            var attributes = Keychain.Attributes()
+            attributes.class = .genericPassword
+
+            switch self {
+            case .accountToken(let accountToken):
+                attributes.account = accountToken
+                attributes.service = kServiceName
+
+            case .persistentReference(let persistentReference):
+                attributes.valuePersistentReference = persistentReference
+            }
+
+            return attributes
         }
     }
 
-    static func remove(account: String) -> Result<(), TunnelConfigurationManagerError> {
-        Keychain.removeItem(account: account)
-            .mapError { .removeKeychainItem($0) }
+    struct KeychainEntry {
+        let accountToken: String
+        let tunnelConfiguration: TunnelConfiguration
     }
 
-    static func getPersistentKeychainRef(account: String) -> Result<Data, TunnelConfigurationManagerError> {
-        Keychain.getPersistentRef(account: account)
-            .mapError { .getPersistentKeychainRef($0) }
-    }
+    static func load(searchTerm: KeychainSearchTerm) -> Result<KeychainEntry> {
+        var query = searchTerm.makeKeychainAttributes()
+        query.return = [.data, .attributes]
 
-}
+        return Keychain.findFirst(query: query)
+            .mapError { .getFromKeychain($0) }
+            .flatMap { (attributes) in
+                let attributes = attributes!
+                let account = attributes.account!
+                let data = attributes.valueData!
 
-private enum Keychain {}
+                return Self.decode(data: data)
+                    .map { KeychainEntry(accountToken: account, tunnelConfiguration: $0) }
+        }
+    }
 
-private extension Keychain {
+    static func add(configuration: TunnelConfiguration, account: String) -> Result<()> {
+        Self.encode(tunnelConfig: configuration)
+            .flatMap { (data) -> Result<()> in
+                var attributes = KeychainSearchTerm.accountToken(account)
+                    .makeKeychainAttributes()
 
-    /// A Keychain Result type
-    typealias Result<T> = Swift.Result<T, KeychainError>
+                // Share the item with the application group
+                attributes.accessGroup = ApplicationConfiguration.securityGroupIdentifier
 
-    /// List all of the account tokens in Keychain
-    static func listAccounts() -> Result<[String]> {
-        let query: [CFString: Any] = [
-            kSecClass: kSecClassGenericPassword,
-            kSecAttrService: kServiceName,
-            kSecReturnAttributes: true,
-            kSecMatchLimit: kSecMatchLimitAll,
-        ]
+                // Make sure the keychain item is available after the first unlock to enable
+                // automatic key rotation in background (from the packet tunnel process)
+                attributes.accessible = Self.keychainAccessibleLevel
 
-        return executeSecCopyMatching(query: query)
-            .map { (result) in
-                let attrs = result as! [[CFString: Any]]
-                let accountTokens = attrs.compactMap { (dict) in
-                    dict[kSecAttrAccount] as? String
-                }
+                // Store value
+                attributes.valueData = data
 
-                return accountTokens
+                // Add revision
+                KeychainItemRevision.firstRevision().store(in: &attributes)
+
+                return Keychain.add(attributes)
+                    .mapError { .addToKeychain($0) }
+                    .map { _ in () }
         }
     }
 
-    /// Get a persistent reference to the Keychain item for the given account token
-    static func getPersistentRef(account: String) -> Result<Data> {
-        let query: [CFString: Any] = [
-            kSecClass: kSecClassGenericPassword,
-            kSecAttrAccount: account,
-            kSecAttrService: kServiceName,
-            kSecReturnPersistentRef: true
-        ]
+    /// This is a migration path for the existing Keychain entries created by 2020.2 or before.
+    ///
+    /// - Set the appropriate `accessible` so that the Packet Tunnel can access the tunnel
+    ///    configuration when the device is locked.
+    /// - Add revision field
+    ///
+    /// - Returns: A boolean that indicates whether the entry was up to date prior to the
+    ///            migration request.
 
-        return executeSecCopyMatching(query: query)
-            .map { $0 as! Data }
-    }
+    static func migrateKeychainEntry(searchTerm: KeychainSearchTerm) -> Result<Bool> {
+        var queryAttributes = searchTerm.makeKeychainAttributes()
+        queryAttributes.return = [.attributes]
 
-    /// Get data associated with the given persistent Keychain reference
-    static func getItemData(persistentKeychainRef: Data) -> Result<Data> {
-        let query: [CFString: Any] = [
-            kSecClass: kSecClassGenericPassword,
-            kSecValuePersistentRef: persistentKeychainRef,
-            kSecReturnData: true
-        ]
+        return Keychain.findFirst(query: queryAttributes)
+            .mapError { .getFromKeychain($0) }
+            .flatMap { (itemAttributes) -> Result<Bool> in
+                let itemAttributes = itemAttributes!
 
-        return executeSecCopyMatching(query: query)
-            .map { $0 as! Data }
-    }
+                let searchAttributes = searchTerm.makeKeychainAttributes()
+                var updateAttributes = Keychain.Attributes()
+
+                // Add revision if it's missing
+                if KeychainItemRevision(attributes: itemAttributes) == nil {
+                    KeychainItemRevision.firstRevision().store(in: &updateAttributes)
+                }
 
-    /// Get data associated with the given account token
-    static func getItemData(account: String) -> Result<Data> {
-        let query: [CFString: Any] = [
-            kSecClass: kSecClassGenericPassword,
-            kSecAttrAccount: account,
-            kSecAttrService: kServiceName,
-            kSecReturnData: true
-        ]
+                // Fix the accessibility permission for the Keychain entry
+                if itemAttributes.accessible != Self.keychainAccessibleLevel {
+                    updateAttributes.accessible = Self.keychainAccessibleLevel
+                }
 
-        return executeSecCopyMatching(query: query)
-            .map { $0 as! Data }
+                // Return immediately if nothing to update (i.e the keychain query is empty)
+                if updateAttributes.keychainRepresentation().isEmpty {
+                    return .success(false)
+                } else {
+                    return Keychain.update(query: searchAttributes, update: updateAttributes)
+                        .mapError { .updateKeychain($0) }
+                        .map { true }
+                }
+        }
     }
 
-    /// Store data in the Keychain and associate it with the given account token
-    static func addItem(account: String, data: Data) -> Result<()> {
-        let attributes: [CFString: Any] = [
-            kSecClass: kSecClassGenericPassword,
-            kSecAttrAccount: account,
-            kSecAttrService: kServiceName,
-            kSecValueData: data,
-            kSecReturnData: false,
+    /// Reads the tunnel configuration from Keychain, then passes it to the given closure for
+    /// modifications, saves the result back to Keychain.
+    ///
+    /// The given block may run multiple times if Keychain entry was changed between read and write
+    /// operations.
+    static func update(searchTerm: KeychainSearchTerm,
+                       using changeConfiguration: (inout TunnelConfiguration) -> Void)
+        -> Result<TunnelConfiguration>
+    {
+        while true {
+            var searchQuery = searchTerm.makeKeychainAttributes()
+            searchQuery.return = [.attributes, .data]
 
-            // Share the item with the application group
-            kSecAttrAccessGroup: ApplicationConfiguration.securityGroupIdentifier,
-        ]
+            let result = Keychain.findFirst(query: searchQuery)
+                .mapError { TunnelConfigurationManager.Error.getFromKeychain($0) }
+                .flatMap { (itemAttributes) -> Result<TunnelConfiguration> in
+                    let itemAttributes = itemAttributes!
+                    let serializedData = itemAttributes.valueData!
+                    let account = itemAttributes.account!
 
-        let status = SecItemAdd(attributes as CFDictionary, nil)
+                    // Parse the current revision from Keychain attributes
+                    let currentRevision = KeychainItemRevision(attributes: itemAttributes)
 
-        return mapSecResult(status: status) {
-            ()
-        }
-    }
+                    // Pick the next revision in sequence
+                    let nextRevision = currentRevision?.nextRevision
+                        ?? KeychainItemRevision.firstRevision()
 
-    /// Replace the data associated with the given account token.
-    static func updateItem(account: String, data: Data) -> Result<()> {
-        let query: [CFString: Any] = [
-            kSecClass: kSecClassGenericPassword,
-            kSecAttrAccount: account,
-            kSecAttrService: kServiceName,
-        ]
+                    return Self.decode(data: serializedData)
+                        .flatMap { (tunnelConfig) -> Result<TunnelConfiguration> in
+                            var tunnelConfig = tunnelConfig
+                            changeConfiguration(&tunnelConfig)
 
-        let update: [CFString: Any] = [
-            kSecValueData: data
-        ]
+                            return Self.encode(tunnelConfig: tunnelConfig)
+                                .flatMap { (newData) -> Result<TunnelConfiguration> in
+                                    // `SecItemUpdate` does not accept query parameters when using
+                                    // persistent reference, so constraint the query to account
+                                    // token instead now when we know it
+                                    var updateQuery = KeychainSearchTerm
+                                        .accountToken(account)
+                                        .makeKeychainAttributes()
 
-        let status = SecItemUpdate(query as CFDictionary, update as CFDictionary)
+                                    // Provide the last known revision via generic field to prevent
+                                    // overwriting the item if it was modified in the meanwhile.
+                                    // This field can be missing for the existing apps on AppStore
+                                    currentRevision?.store(in: &updateQuery)
 
-        return mapSecResult(status: status) {
-            ()
-        }
-    }
+                                    var updateAttributes = Keychain.Attributes()
+                                    updateAttributes.valueData = newData
 
-    /// Remove the data associated with the given account token
-    static func removeItem(account: String) -> Result<()> {
-        let query: [CFString: Any] = [
-            kSecClass: kSecClassGenericPassword,
-            kSecAttrAccount: account,
-            kSecAttrService: kServiceName,
-        ]
+                                    // Add the next revision number
+                                    nextRevision.store(in: &updateAttributes)
 
-        let status = SecItemDelete(query as CFDictionary)
+                                    return Keychain.update(query: updateQuery, update: updateAttributes)
+                                        .mapError { TunnelConfigurationManager.Error.updateKeychain($0) }
+                                        .map { tunnelConfig }
+                            }
+                        }
+            }
 
-        return mapSecResult(status: status) {
-            ()
+            // Retry if Keychain reported that the item was not found when updating
+            if case .failure(.updateKeychain(.itemNotFound)) = result  {
+                continue
+            } else {
+                return result
+            }
         }
     }
 
-    /// A private helper that verifies the given `status` and executes `body` on success
-    static private func mapSecResult<T>(status: OSStatus, body: () -> T) -> Result<T> {
-        if status == errSecSuccess {
-            return .success(body())
-        } else {
-            return .failure(KeychainError(code: status))
-        }
+    static func remove(searchTerm: KeychainSearchTerm) -> Result<()> {
+        return Keychain.delete(query: searchTerm.makeKeychainAttributes())
+            .mapError { .removeKeychainItem($0) }
     }
 
-    /// A private helper to execute the given query using `SecCopyMatching` and map the result to
-    /// the `Result<CFTypeRef?>` type.
-    static private func executeSecCopyMatching(query: [CFString: Any]) -> Result<CFTypeRef?> {
-        var result: CFTypeRef?
-        let status = SecItemCopyMatching(query as CFDictionary, &result)
+    /// Get a persistent reference to the Keychain item for the given account token
+    static func getPersistentKeychainReference(account: String) -> Result<Data> {
+        var query = KeychainSearchTerm.accountToken(account)
+            .makeKeychainAttributes()
+        query.return = [.persistentReference]
 
-        return mapSecResult(status: status) {
-            result
+        return Keychain.findFirst(query: query)
+            .mapError { .getPersistentKeychainRef($0) }
+            .map { (attributes) -> Data in
+                return attributes!.valuePersistentReference!
         }
     }
 
+    private static func encode(tunnelConfig: TunnelConfiguration) -> Result<Data> {
+        return Swift.Result { try JSONEncoder().encode(tunnelConfig) }
+            .mapError { .encode($0) }
+    }
+
+    private static func decode(data: Data) -> Result<TunnelConfiguration> {
+        return Swift.Result { try JSONDecoder().decode(TunnelConfiguration.self, from: data) }
+            .mapError { .decode($0) }
+    }
 }
diff --git a/ios/MullvadVPN/TunnelManager.swift b/ios/MullvadVPN/TunnelManager.swift
index f90140491e..e418061cd8 100644
--- a/ios/MullvadVPN/TunnelManager.swift
+++ b/ios/MullvadVPN/TunnelManager.swift
@@ -32,13 +32,13 @@ enum TunnelManagerError: Error {
     case unsetAccount(UnsetAccountError)
 
     /// A failure to set the relay constraints
-    case setRelayConstraints(UpdateTunnelConfigurationError)
+    case setRelayConstraints(TunnelConfigurationManager.Error)
 
     /// A failure to get the relay constraints
-    case getRelayConstraints(TunnelConfigurationManagerError)
+    case getRelayConstraints(TunnelConfigurationManager.Error)
 
     /// A failure to get a public key used for Wireguard
-    case getWireguardPublicKey(TunnelConfigurationManagerError)
+    case getWireguardPublicKey(TunnelConfigurationManager.Error)
 
     /// A failure to re-generate a private key used for Wireguard
     case regenerateWireguardPrivateKey(RegenerateWireguardPrivateKeyError)
@@ -74,7 +74,7 @@ extension TunnelManagerError: LocalizedError {
         case .setAccount(.pushWireguardKey(let pushError)),
              .regenerateWireguardPrivateKey(.replaceWireguardKey(let pushError)):
             switch pushError {
-            case .transport(.network(let urlError)):
+            case .network(let urlError):
                 return urlError.localizedDescription
 
             case .server(let serverError):
@@ -127,13 +127,13 @@ enum TunnelIpcRequestError: Error {
 
 enum SetAccountError: Error {
     /// A failure to make the tunnel configuration
-    case makeTunnelConfiguration(TunnelConfigurationManagerError)
+    case makeTunnelConfiguration(TunnelConfigurationManager.Error)
 
     /// A failure to update the tunnel configuration
-    case updateTunnelConfiguration(UpdateTunnelConfigurationError)
+    case updateTunnelConfiguration(TunnelConfigurationManager.Error)
 
     /// A failure to push the wireguard key
-    case pushWireguardKey(PushWireguardKeyError)
+    case pushWireguardKey(MullvadRpc.Error)
 
     /// A failure to set up a tunnel
     case setup(SetupTunnelError)
@@ -144,34 +144,18 @@ enum UnsetAccountError: Error {
     case removeTunnel(Error)
 
     /// A failure to remove a tunnel configuration from Keychain
-    case removeTunnelConfiguration(TunnelConfigurationManagerError)
+    case removeTunnelConfiguration(TunnelConfigurationManager.Error)
 }
 
 enum RegenerateWireguardPrivateKeyError: Error {
     /// A failure to read the public Wireguard key from Keychain
-    case readPublicWireguardKey(TunnelConfigurationManagerError)
+    case readPublicWireguardKey(TunnelConfigurationManager.Error)
 
     /// A failure to replace the public Wireguard key
-    case replaceWireguardKey(PushWireguardKeyError)
+    case replaceWireguardKey(MullvadRpc.Error)
 
     /// A failure to update tunnel configuration
-    case updateTunnelConfiguration(UpdateTunnelConfigurationError)
-
-    /// A failure to set up a tunnel
-    case setupTunnel(SetupTunnelError)
-}
-
-enum PushWireguardKeyError: Error {
-    case transport(MullvadAPI.Error)
-    case server(MullvadAPI.ResponseError)
-}
-
-enum UpdateTunnelConfigurationError: Error {
-    /// Unable to load the existing configuration
-    case loadTunnelConfiguration(TunnelConfigurationManagerError)
-
-    /// Unable to save the configuration
-    case saveTunnelConfiguration(TunnelConfigurationManagerError)
+    case updateTunnelConfiguration(TunnelConfigurationManager.Error)
 }
 
 enum StartTunnelError: Error {
@@ -193,7 +177,7 @@ enum SetupTunnelError: Error {
     case reloadTunnel(Error)
 
     /// Unable to obtain the keychain reference for the configuration
-    case obtainKeychainRef(TunnelConfigurationManagerError)
+    case obtainKeychainRef(TunnelConfigurationManager.Error)
 }
 
 enum LoadTunnelError: Error {
@@ -306,7 +290,7 @@ class TunnelManager {
     /// A queue used for access synchronization to the TunnelManager members
     private let executionQueue = DispatchQueue(label: "net.mullvad.vpn.tunnel-manager.execution-queue")
 
-    private let apiClient = MullvadAPI()
+    private let rpc = MullvadRpc()
     private var tunnelProvider: TunnelProviderManagerType?
     private var tunnelIpc: PacketTunnelIpc?
 
@@ -336,6 +320,11 @@ class TunnelManager {
                 .receive(on: self.executionQueue)
                 .flatMap { (tunnels) -> AnyPublisher<(), LoadTunnelError> in
 
+                    // Migrate tunnel configuration if needed
+                    if let accountToken = accountToken {
+                        self.migrateTunnelConfiguration(accountToken: accountToken)
+                    }
+
                     // No tunnels found. Save the account token.
                     guard let tunnelProvider = tunnels?.first else {
                         self.accountToken = accountToken
@@ -427,21 +416,16 @@ class TunnelManager {
                     // Send wireguard key to the server
                     let publicKey = tunnelConfig.interface.privateKey.publicKey.rawRepresentation
 
-                    return self.apiClient.pushWireguardKey(accountToken: accountToken, publicKey: publicKey)
-                        .mapError { (networkError) -> SetAccountError in
-                            return .pushWireguardKey(.transport(networkError))
-                    }.flatMap { (response: MullvadAPI.Response<WireguardAssociatedAddresses>) in
-                        return response.result.publisher
-                            .mapError { (serverError) -> SetAccountError in
-                                return .pushWireguardKey(.server(serverError))
-                        }
-                    }.flatMap { (addresses) in
-                        return self.updateAssociatedAddresses(
-                            accountToken: accountToken,
-                            addresses: addresses
-                        ).mapError { SetAccountError.updateTunnelConfiguration($0) }
-                            .publisher
-                    }.flatMap { setupTunnelPublisher }.eraseToAnyPublisher()
+                    return self.rpc.pushWireguardKey(accountToken: accountToken, publicKey: publicKey)
+                        .mapError { SetAccountError.pushWireguardKey($0) }
+                        .flatMap { (addresses) in
+                            self.updateAssociatedAddresses(
+                                accountToken: accountToken,
+                                addresses: addresses
+                            ).mapError { SetAccountError.updateTunnelConfiguration($0) }
+                                .publisher
+                                .flatMap { _ in setupTunnelPublisher }
+                    }.eraseToAnyPublisher()
             }.mapError { TunnelManagerError.setAccount($0) }
         }.eraseToAnyPublisher()
     }
@@ -458,37 +442,30 @@ class TunnelManager {
                     let removeKeychainConfigPublisher = Deferred {
                         () -> AnyPublisher<(), UnsetAccountError> in
                         // Load existing configuration
-                        switch TunnelConfigurationManager.load(account: accountToken) {
-                        case .success(let tunnelConfig):
-                            let publicKey = tunnelConfig.interface
+                        switch TunnelConfigurationManager.load(searchTerm: .accountToken(accountToken)) {
+                        case .success(let keychainEntry):
+                            let publicKey = keychainEntry.tunnelConfiguration
+                                .interface
                                 .privateKey
                                 .publicKey
                                 .rawRepresentation
 
                             // Remove configuration from Keychain
-                            return TunnelConfigurationManager.remove(account: accountToken)
+                            return TunnelConfigurationManager.remove(searchTerm: .accountToken(accountToken))
                                 .mapError { UnsetAccountError.removeTunnelConfiguration($0) }
                                 .publisher
                                 .flatMap {
                                     // Remove WireGuard key from master
-                                    self.apiClient.removeWireguardKey(
+                                    self.rpc.removeWireguardKey(
                                         accountToken: accountToken,
                                         publicKey: publicKey
                                     )
                                         .retry(1)
-                                        .map({ (response) -> () in
-                                            switch response.result {
-                                            case .success(let isRemoved):
-                                                os_log(.debug, "Removed the WireGuard key from server: %{public}s", "\(isRemoved)")
-
-                                            case .failure(let error):
-                                                os_log(.error, "Failed to remove the WireGuard key from server. Server error: %{public}s", error.localizedDescription)
-                                            }
-
-                                            // Suppress server errors
+                                        .map({ (isRemoved) -> () in
+                                            os_log(.debug, "Removed the WireGuard key from server: %{public}s", "\(isRemoved)")
                                             return ()
                                         }).catch({ (error) -> Result<(), UnsetAccountError>.Publisher in
-                                            os_log(.error, "Failed to remove the Wireguard key from server. Network error: %{public}s", error.localizedDescription)
+                                            os_log(.error, "Failed to remove the Wireguard key from server: %{public}s", error.localizedDescription)
 
                                             // Suppress network errors
                                             return Result.Publisher(())
@@ -558,51 +535,40 @@ class TunnelManager {
                 .flatMap { (accountToken) -> AnyPublisher<(), TunnelManagerError> in
                     let newPrivateKey = WireguardPrivateKey()
 
-                    return TunnelConfigurationManager.load(account: accountToken)
-                        .map { $0.interface.privateKey.publicKey }
+                    return TunnelConfigurationManager.load(searchTerm: .accountToken(accountToken))
+                        .map { $0.tunnelConfiguration.interface.privateKey.publicKey }
                         .mapError { RegenerateWireguardPrivateKeyError.readPublicWireguardKey($0) }
                         .publisher
                         .flatMap { (oldPublicKey) in
-                            self.apiClient.replaceWireguardKey(
+                            self.rpc.replaceWireguardKey(
                                 accountToken: accountToken,
                                 oldPublicKey: oldPublicKey.rawRepresentation,
                                 newPublicKey: newPrivateKey.publicKey.rawRepresentation)
-                                .mapError { (networkError) -> RegenerateWireguardPrivateKeyError in
-                                    return .replaceWireguardKey(.transport(networkError))
-                            }.receive(on: self.executionQueue)
-                                .flatMap { (response: MullvadAPI.Response<WireguardAssociatedAddresses>) in
-                                    return response.result.publisher
-                                        .mapError { (serverError) -> RegenerateWireguardPrivateKeyError in
-                                            return .replaceWireguardKey(.server(serverError))
+                                .mapError { RegenerateWireguardPrivateKeyError.replaceWireguardKey($0) }
+                                .receive(on: self.executionQueue)
+                                .flatMap { (addresses) in
+                                    TunnelConfigurationManager.update(searchTerm: .accountToken(accountToken)) {
+                                        (tunnelConfiguration) in
+                                        tunnelConfiguration.interface.privateKey = newPrivateKey
+                                        tunnelConfiguration.interface.addresses = [
+                                            addresses.ipv4Address,
+                                            addresses.ipv6Address
+                                        ]
                                     }
-                            }
-                            .flatMap { (addresses) in
-                                self.updateTunnelConfiguration(accountToken: accountToken) {
-                                    (tunnelConfiguration) in
-                                    tunnelConfiguration.interface.privateKey = newPrivateKey
-                                    tunnelConfiguration.interface.addresses = [
-                                        addresses.ipv4Address,
-                                        addresses.ipv6Address
-                                    ]
-                                }
-                                .mapError { .updateTunnelConfiguration($0) }
-                                .publisher
-                            }
-                            .flatMap {
-                                self.setupTunnel(accountToken: accountToken)
+                                    .mapError { .updateTunnelConfiguration($0) }
                                     .map { _ in () }
-                                    .mapError { RegenerateWireguardPrivateKeyError.setupTunnel($0) }
+                                    .publisher
                             }.receive(on: self.executionQueue)
                                 .flatMap { _ in
                                     // Ignore Packet Tunnel IPC errors but log them
                                     self.reloadPacketTunnelConfiguration()
-                                        .replaceError(with: ())
-                                        .setFailureType(to: RegenerateWireguardPrivateKeyError.self)
                                         .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)
                                             }
                                         })
+                                        .replaceError(with: ())
+                                        .setFailureType(to: RegenerateWireguardPrivateKeyError.self)
                             }
                     }
                     .mapError { TunnelManagerError.regenerateWireguardPrivateKey($0) }
@@ -617,11 +583,12 @@ class TunnelManager {
                 .setFailureType(to: TunnelManagerError.self)
                 .replaceNil(with: .missingAccount)
                 .flatMap { (accountToken) in
-                    self.updateTunnelConfiguration(accountToken: accountToken) { (tunnelConfig) in
-                        tunnelConfig.relayConstraints = constraints
+                    TunnelConfigurationManager
+                        .update(searchTerm: .accountToken(accountToken)) { (tunnelConfig) in
+                            tunnelConfig.relayConstraints = constraints
                     }.mapError { TunnelManagerError.setRelayConstraints($0) }
                         .publisher
-                        .flatMap {
+                        .flatMap { _ in
                             // Ignore Packet Tunnel IPC errors but log them
                             self.reloadPacketTunnelConfiguration()
                                 .replaceError(with: ())
@@ -642,16 +609,17 @@ class TunnelManager {
                 .setFailureType(to: TunnelManagerError.self)
                 .replaceNil(with: .missingAccount)
                 .flatMap { (accountToken) in
-                    TunnelConfigurationManager.load(account: accountToken)
-                        .map { $0.relayConstraints }
-                        .flatMapError { (error) -> Result<RelayConstraints, TunnelConfigurationManagerError> in
+                    TunnelConfigurationManager.load(searchTerm: .accountToken(accountToken))
+                        .map { $0.tunnelConfiguration.relayConstraints }
+                        .flatMapError { (error) -> Result<RelayConstraints, TunnelConfigurationManager.Error> in
                             // Return default constraints if the config is not found in Keychain
                             if case .getFromKeychain(.itemNotFound) = error {
                                 return .success(TunnelConfiguration().relayConstraints)
                             } else {
                                 return .failure(error)
                             }
-                    }.mapError { TunnelManagerError.getRelayConstraints($0) }.publisher
+                    }.mapError { .getRelayConstraints($0) }
+                        .publisher
             }
         }.eraseToAnyPublisher()
     }
@@ -662,9 +630,10 @@ class TunnelManager {
                 .setFailureType(to: TunnelManagerError.self)
                 .replaceNil(with: .missingAccount)
                 .flatMap { (accountToken) in
-                    TunnelConfigurationManager.load(account: accountToken)
-                        .map { $0.interface.privateKey.publicKey }
-                        .mapError { TunnelManagerError.getWireguardPublicKey($0) }.publisher
+                    TunnelConfigurationManager.load(searchTerm: .accountToken(accountToken))
+                        .map { $0.tunnelConfiguration.interface.privateKey.publicKey }
+                        .mapError { .getWireguardPublicKey($0) }
+                        .publisher
             }
         }.eraseToAnyPublisher()
     }
@@ -673,21 +642,25 @@ class TunnelManager {
 
     /// Tell Packet Tunnel process to reload the tunnel configuration
     private func reloadPacketTunnelConfiguration() -> AnyPublisher<(), TunnelIpcRequestError> {
-        Just(tunnelIpc)
-            .setFailureType(to: TunnelIpcRequestError.self)
-            .replaceNil(with: .missingIpc)
-            .flatMap { (tunnelIpc) in
-                tunnelIpc.reloadConfiguration()
-                    .mapError { .send($0) }
-        }.eraseToAnyPublisher()
+        return executeIpcRequest { $0.reloadConfiguration() }
     }
 
+    /// Ask Packet Tunnel process to return the current tunnel connection info
     private func getTunnelConnectionInfo() -> AnyPublisher<TunnelConnectionInfo, TunnelIpcRequestError> {
+        return executeIpcRequest { $0.getTunnelInformation() }
+    }
+
+    /// IPC interactor helper that automatically maps the `PacketTunnelIpcError` to
+    /// `TunnelIpcRequestError`
+    private func executeIpcRequest<T>(
+        _ body: @escaping (PacketTunnelIpc) -> AnyPublisher<T, PacketTunnelIpcError>
+    ) -> AnyPublisher<T, TunnelIpcRequestError>
+    {
         Just(tunnelIpc)
             .setFailureType(to: TunnelIpcRequestError.self)
             .replaceNil(with: .missingIpc)
             .flatMap { (tunnelIpc) in
-                tunnelIpc.getTunnelInformation()
+                body(tunnelIpc)
                     .mapError { .send($0) }
         }.eraseToAnyPublisher()
     }
@@ -782,13 +755,16 @@ class TunnelManager {
     }
 
     /// Retrieve the existing TunnelConfiguration or create a new one
-    private func makeTunnelConfiguration(accountToken: String) -> Result<TunnelConfiguration, TunnelConfigurationManagerError> {
-        TunnelConfigurationManager.load(account: accountToken)
-            .flatMapError { (error) -> Result<TunnelConfiguration, TunnelConfigurationManagerError> in
+    private func makeTunnelConfiguration(accountToken: String) -> Result<TunnelConfiguration, TunnelConfigurationManager.Error> {
+        TunnelConfigurationManager.load(searchTerm: .accountToken(accountToken))
+            .map { $0.tunnelConfiguration }
+            .flatMapError { (error) -> Result<TunnelConfiguration, TunnelConfigurationManager.Error> in
                 // Return default tunnel configuration if the config is not found in Keychain
                 if case .getFromKeychain(.itemNotFound) = error {
                     let defaultConfiguration = TunnelConfiguration()
-                    return TunnelConfigurationManager.save(configuration: defaultConfiguration, account: accountToken)
+
+                    return TunnelConfigurationManager
+                        .add(configuration: defaultConfiguration, account: accountToken)
                         .map { defaultConfiguration }
                 } else {
                     return .failure(error)
@@ -805,7 +781,7 @@ class TunnelManager {
                 return tunnels?.first ?? TunnelProviderManagerType()
         }
         .flatMap { (tunnelProvider) in
-            TunnelConfigurationManager.getPersistentKeychainRef(account: accountToken)
+            TunnelConfigurationManager.getPersistentKeychainReference(account: accountToken)
                 .mapError { SetupTunnelError.obtainKeychainRef($0) }
                 .map { (tunnelProvider, $0) }
                 .publisher
@@ -855,21 +831,12 @@ class TunnelManager {
         return protocolConfig
     }
 
-    private func updateTunnelConfiguration(accountToken: String, using block: (inout TunnelConfiguration) -> Void) -> Result<(), UpdateTunnelConfigurationError> {
-        TunnelConfigurationManager.load(account: accountToken)
-            .mapError { UpdateTunnelConfigurationError.loadTunnelConfiguration($0) }
-            .flatMap { (tunnelConfig) -> Result<(), UpdateTunnelConfigurationError> in
-                var tunnelConfig = tunnelConfig
-
-                block(&tunnelConfig)
-
-                return TunnelConfigurationManager.save(configuration: tunnelConfig, account: accountToken)
-                    .mapError { .saveTunnelConfiguration($0) }
-        }
-    }
-
-    private func updateAssociatedAddresses(accountToken: String, addresses: WireguardAssociatedAddresses) -> Result<(), UpdateTunnelConfigurationError> {
-        updateTunnelConfiguration(accountToken: accountToken) { (tunnelConfig) in
+    private func updateAssociatedAddresses(
+        accountToken: String,
+        addresses: WireguardAssociatedAddresses
+    ) -> Result<TunnelConfiguration, TunnelConfigurationManager.Error>
+    {
+        TunnelConfigurationManager.update(searchTerm: .accountToken(accountToken)) { (tunnelConfig) in
             tunnelConfig.interface.addresses = [
                 addresses.ipv4Address,
                 addresses.ipv6Address
@@ -877,6 +844,24 @@ class TunnelManager {
         }
     }
 
+    private func migrateTunnelConfiguration(accountToken: String) {
+        let result = TunnelConfigurationManager
+            .migrateKeychainEntry(searchTerm: .accountToken(accountToken))
+
+        switch result {
+        case .success(let migrated):
+            if migrated {
+                os_log("Migrated Keychain tunnel configuration")
+            } else {
+                os_log("Tunnel configuration is up to date. No migration needed.")
+            }
+
+        case .failure(let error):
+            os_log("Failed to migrate tunnel configuration: %{public}s",
+                   error.localizedDescription)
+        }
+    }
+
 }
 
 /// Convenience methods to provide `Future` based alternatives for working with
diff --git a/ios/MullvadVPN/WireguardKeysViewController.swift b/ios/MullvadVPN/WireguardKeysViewController.swift
index 1a50dc9c7e..513e6097db 100644
--- a/ios/MullvadVPN/WireguardKeysViewController.swift
+++ b/ios/MullvadVPN/WireguardKeysViewController.swift
@@ -59,7 +59,8 @@ class WireguardKeysViewController: UIViewController {
     @IBOutlet var verifyKeyButton: UIButton!
     @IBOutlet var wireguardKeyStatusView: WireguardKeyStatusView!
 
-    private var fetchKeySubscriber: AnyCancellable?
+    private var tunnelStateSubscriber: AnyCancellable?
+    private var loadKeySubscriber: AnyCancellable?
     private var verifyKeySubscriber: AnyCancellable?
     private var regenerateKeySubscriber: AnyCancellable?
     private var creationDateTimerSubscriber: AnyCancellable?
@@ -100,6 +101,18 @@ class WireguardKeysViewController: UIViewController {
                 }
         }
 
+        tunnelStateSubscriber = TunnelManager.shared.$tunnelState
+            .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)
+                }
+            })
+
         loadPublicKey(animated: false)
     }
 
@@ -160,7 +173,7 @@ class WireguardKeysViewController: UIViewController {
     }
 
     private func loadPublicKey(animated: Bool) {
-        fetchKeySubscriber = TunnelManager.shared.getWireguardPublicKey()
+        loadKeySubscriber = TunnelManager.shared.getWireguardPublicKey()
             .receive(on: DispatchQueue.main)
             .sink(receiveCompletion: { (completion) in
                 switch completion {
diff --git a/ios/MullvadVPN/WireguardPrivateKey.swift b/ios/MullvadVPN/WireguardPrivateKey.swift
index 293750f893..07defa7cf8 100644
--- a/ios/MullvadVPN/WireguardPrivateKey.swift
+++ b/ios/MullvadVPN/WireguardPrivateKey.swift
@@ -52,7 +52,7 @@ extension WireguardPrivateKey: Equatable {
 }
 
 /// A struct holding a public key used for Wireguard with associated metadata
-struct WireguardPublicKey {
+struct WireguardPublicKey: Codable, Equatable {
     /// Refers to private key creation date
     let creationDate: Date