Refactor FXIOS-11415 logins keychain calls (#24845)

* Replaced logins MZKeychainWrapper calls with RustKeychain calls

* Added RustKeychain logins tests, removed unnecessary RustLogins variable

Author: lougeniaC64

firefox-ios/Client/Application/AppDelegate.swift

@@ -22,10 +22,17 @@ class AppDelegate: UIResponder, UIApplicationDelegate, FeatureFlaggable {
         .value()
         .creditCardAutofillStatus
 
+    private let rustKeychainEnabled = FxNimbus.shared
+        .features
+        .rustKeychainRefactor
+        .value()
+        .rustKeychainEnabled
+
     lazy var profile: Profile = BrowserProfile(
         localName: "profile",
         fxaCommandsDelegate: UIApplication.shared.fxaCommandsDelegate,
-        creditCardAutofillEnabled: creditCardAutofillStatus
+        creditCardAutofillEnabled: creditCardAutofillStatus,
+        rustKeychainEnabled: rustKeychainEnabled
     )
 
     lazy var themeManager: ThemeManager = DefaultThemeManager(sharedContainerIdentifier: AppInfo.sharedContainerIdentifier)

firefox-ios/Client/FeatureFlags/NimbusFlaggableFeature.swift

@@ -58,6 +58,7 @@ enum NimbusFeatureFlagID: String, CaseIterable {
     case toolbarNavigationHint
     case tosFeature
     case trackingProtectionRefactor
+    case useRustKeychain
     case zoomFeature
 
     // Add flags here if you want to toggle them in the `FeatureFlagsDebugViewController`. Add in alphabetical order.
@@ -80,7 +81,8 @@ enum NimbusFeatureFlagID: String, CaseIterable {
                 .pdfRefactor,
                 .downloadLiveActivities,
                 .unifiedAds,
-                .unifiedSearch:
+                .unifiedSearch,
+                .useRustKeychain:
             return rawValue + PrefsKeys.FeatureFlags.DebugSuffixKey
         default:
             return nil
@@ -158,6 +160,7 @@ struct NimbusFlaggableFeature: HasNimbusSearchBar {
                 .toolbarNavigationHint,
                 .tosFeature,
                 .trackingProtectionRefactor,
+                .useRustKeychain,
                 .zoomFeature:
             return nil
         }

firefox-ios/Client/Frontend/Settings/Main/Debug/FeatureFlags/FeatureFlagsDebugViewController.swift

@@ -143,6 +143,13 @@ final class FeatureFlagsDebugViewController: SettingsTableViewController, Featur
                 ) { [weak self] _ in
                     self?.reloadView()
                 },
+                FeatureFlagsBoolSetting(
+                    with: .useRustKeychain,
+                    titleText: format(string: "Enable Rust Keychain"),
+                    statusText: format(string: "Toggle to enable Rust Components rust keychain")
+                ) { [weak self] _ in
+                    self?.reloadView()
+                },
                 FeatureFlagsBoolSetting(
                     with: .sentFromFirefox,
                     titleText: format(string: "Enable Sent from Firefox"),

firefox-ios/Client/Nimbus/NimbusFeatureFlagLayer.swift

@@ -145,6 +145,9 @@ final class NimbusFeatureFlagLayer {
         case .trackingProtectionRefactor:
             return checkTrackingProtectionRefactor(from: nimbus)
 
+        case .useRustKeychain:
+            return checkUseRustKeychainFeature(from: nimbus)
+
         case .zoomFeature:
             return checkZoomFeature(from: nimbus)
         }
@@ -443,4 +446,8 @@ final class NimbusFeatureFlagLayer {
     private func checkNICErrorPageFeature(from nimbus: FxNimbus) -> Bool {
         return nimbus.features.nativeErrorPageFeature.value().noInternetConnectionError
     }
+
+    private func checkUseRustKeychainFeature(from nimbus: FxNimbus) -> Bool {
+        return nimbus.features.rustKeychainRefactor.value().rustKeychainEnabled
+    }
 }

firefox-ios/Providers/Profile.swift

@@ -230,8 +230,10 @@ open class BrowserProfile: Profile {
             fatalError("Could not create directory at root path: \(error)")
         }
     }()
+    private var rustKeychainEnabled = false
     fileprivate let name: String
-    fileprivate let keychain: MZKeychainWrapper
+    fileprivate let keychain: RustKeychain
+    fileprivate let legacyKeychain: MZKeychainWrapper
     var isShutdown = false
 
     internal let files: FileAccessor
@@ -257,14 +259,17 @@ open class BrowserProfile: Profile {
     init(localName: String,
          fxaCommandsDelegate: FxACommandsDelegate? = nil,
          creditCardAutofillEnabled: Bool = false,
+         rustKeychainEnabled: Bool = false,
          clear: Bool = false,
          logger: Logger = DefaultLogger.shared) {
         logger.log("Initing profile \(localName) on thread \(Thread.current).",
                    level: .debug,
                    category: .setup)
         self.name = localName
         self.files = ProfileFileAccessor(localName: localName)
-        self.keychain = MZKeychainWrapper.sharedClientAppContainerKeychain
+        self.rustKeychainEnabled = rustKeychainEnabled
+        self.keychain = KeychainManager.shared
+        self.legacyKeychain = KeychainManager.legacyShared
         self.logger = logger
         self.fxaCommandsDelegate = fxaCommandsDelegate
 
@@ -301,7 +306,11 @@ open class BrowserProfile: Profile {
             logger.log("New profile. Removing old Keychain/Prefs data.",
                        level: .info,
                        category: .setup)
-            MZKeychainWrapper.wipeKeychain()
+            if rustKeychainEnabled {
+                RustKeychain.wipeKeychain()
+            } else {
+                MZKeychainWrapper.wipeKeychain()
+            }
             prefs.clearAll()
         }
 
@@ -311,7 +320,8 @@ open class BrowserProfile: Profile {
         // Initiating the sync manager has to happen prior to the databases being opened,
         // because opening them can trigger events to which the SyncManager listens.
         self.syncManager = RustSyncManager(profile: self,
-                                           creditCardAutofillEnabled: creditCardAutofillEnabled)
+                                           creditCardAutofillEnabled: creditCardAutofillEnabled,
+                                           rustKeychainEnabled: rustKeychainEnabled)
 
         let notificationCenter = NotificationCenter.default
 
@@ -676,7 +686,7 @@ open class BrowserProfile: Profile {
             fileURLWithPath: directory,
             isDirectory: true
         ).appendingPathComponent("loginsPerField.db").path
-        return RustLogins(databasePath: databasePath)
+        return RustLogins(databasePath: databasePath, rustKeychainEnabled: self.rustKeychainEnabled)
     }()
 
     lazy var remoteSettingsService: RemoteSettingsService? = {
@@ -755,24 +765,39 @@ open class BrowserProfile: Profile {
         prefs.removeObjectForKey(PrefsKeys.KeyLastRemoteTabSyncTime)
 
         // Save the keys that will be restored
-        let rustAutofillKey = RustAutofillEncryptionKeys()
-        let creditCardKey = keychain.string(forKey: rustAutofillKey.ccKeychainKey)
         let rustLoginsKeys = RustLoginEncryptionKeys()
-        let perFieldKey = keychain.string(forKey: rustLoginsKeys.loginPerFieldKeychainKey)
-        // Remove all items, removal is not key-by-key specific (due to the risk of failing to delete something),
-        // simply restore what is needed.
-        keychain.removeAllKeys()
-
-        if let perFieldKey = perFieldKey {
-            keychain.set(
-                perFieldKey,
-                forKey: rustLoginsKeys.loginPerFieldKeychainKey,
-                withAccessibility: .afterFirstUnlock
-            )
+        let rustAutofillKey = RustAutofillEncryptionKeys()
+        var loginsKey: String?
+        let creditCardKey = legacyKeychain.string(forKey: rustAutofillKey.ccKeychainKey)
+
+        if rustKeychainEnabled {
+            (loginsKey, _) = keychain.getLoginsKeyData()
+
+            // Remove all items, removal is not key-by-key specific (due to the risk of failing to delete something),
+            // simply restore what is needed.
+            keychain.removeAllKeys()
+
+            if let loginsKey = loginsKey {
+                keychain.setLoginsKey(loginsKey)
+            }
+        } else {
+            loginsKey = legacyKeychain.string(forKey: rustLoginsKeys.loginPerFieldKeychainKey)
+
+            // Remove all items, removal is not key-by-key specific (due to the risk of failing to delete something),
+            // simply restore what is needed.
+            legacyKeychain.removeAllKeys()
+
+            if let loginsKey = loginsKey {
+                legacyKeychain.set(loginsKey,
+                                   forKey: rustLoginsKeys.loginPerFieldKeychainKey,
+                                   withAccessibility: .afterFirstUnlock)
+            }
         }
 
         if let creditCardKey = creditCardKey {
-            keychain.set(creditCardKey, forKey: rustAutofillKey.ccKeychainKey, withAccessibility: .afterFirstUnlock)
+            legacyKeychain.set(creditCardKey,
+                               forKey: rustAutofillKey.ccKeychainKey,
+                               withAccessibility: .afterFirstUnlock)
         }
 
         // Tell any observers that our account has changed.

firefox-ios/Providers/RustSyncManager.swift

@@ -36,6 +36,7 @@ public class RustSyncManager: NSObject, SyncManager {
     private let fxaDeclinedEngines = "fxa.cwts.declinedSyncEngines"
     private var notificationCenter: NotificationProtocol
     var creditCardAutofillEnabled = false
+    var rustKeychainEnabled = false
 
     let fifteenMinutesInterval = TimeInterval(60 * 15)
 
@@ -68,6 +69,7 @@ public class RustSyncManager: NSObject, SyncManager {
 
     init(profile: BrowserProfile,
          creditCardAutofillEnabled: Bool = false,
+         rustKeychainEnabled: Bool = false,
          logger: Logger = DefaultLogger.shared,
          notificationCenter: NotificationProtocol = NotificationCenter.default) {
         self.profile = profile
@@ -77,6 +79,7 @@ public class RustSyncManager: NSObject, SyncManager {
 
         super.init()
         self.creditCardAutofillEnabled = creditCardAutofillEnabled
+        self.rustKeychainEnabled = rustKeychainEnabled
     }
 
     @objc
@@ -273,9 +276,15 @@ public class RustSyncManager: NSObject, SyncManager {
                     .prefsForSync
                     .branch("scratchpad")
                     .stringForKey("keyLabel") {
-                        MZKeychainWrapper
-                            .sharedClientAppContainerKeychain
-                            .removeObject(forKey: keyLabel)
+                        if self.rustKeychainEnabled {
+                            RustKeychain
+                                .sharedClientAppContainerKeychain
+                                .removeObject(key: keyLabel)
+                        } else {
+                            MZKeychainWrapper
+                                .sharedClientAppContainerKeychain
+                                .removeObject(forKey: keyLabel)
+                        }
                 }
                 self.prefsForSync.clearAll()
             }

firefox-ios/Storage/MockRustKeychain.swift

@@ -2,6 +2,9 @@
 // License, v. 2.0. If a copy of the MPL was not distributed with this
 // file, You can obtain one at http://mozilla.org/MPL/2.0/
 
+import func MozillaAppServices.createCanary
+import func MozillaAppServices.createKey
+
 class MockRustKeychain: RustKeychain {
     static let shared = MockRustKeychain()
 
@@ -11,16 +14,45 @@ class MockRustKeychain: RustKeychain {
         super.init(serviceName: "Test")
     }
 
-    override func getBaseKeychainQuery(key: String) -> [String: Any] {
-        return [:]
+    override public func removeObject(key: String) {
+        _ = storage.removeValue(forKey: key)
+    }
+
+    override public func removeAllKeys() {
+        storage.removeAll()
+    }
+
+    override public func setLoginsKey(_ value: String) {
+        storage[loginsKeyIdentifier] = value
+    }
+
+    override public func getLoginsKeyData() -> (String?, String?) {
+        return (storage[loginsKeyIdentifier], storage[loginsCanaryKeyIdentifier])
+    }
+
+    override public class func wipeKeychain() {}
+
+    override func createLoginsKeyData() throws -> String {
+        guard let keyValue = try? createKey(),
+              let canaryValue = try? createCanary(text: loginsCanaryPhrase, encryptionKey: keyValue) else {
+            throw LoginEncryptionKeyError.noKeyCreated
+        }
+
+        storage[loginsCanaryKeyIdentifier] = canaryValue
+        storage[loginsKeyIdentifier] = keyValue
+        return keyValue
     }
 
     override func queryKeychainForKey(key: String) -> Result<String?, Error> {
         return .success(storage[key])
     }
 
-    override func addOrUpdateKeychainKey(_ value: String, key: String) -> OSStatus {
-        storage[key] = value
-        return errSecSuccess
+    override func createAndStoreKey(canaryPhrase: String, canaryIdentifier: String, keyIdentifier: String) throws -> String {
+        let keyValue = try createKey()
+        let canaryValue = try createCanary(text: canaryPhrase, encryptionKey: keyValue)
+
+        storage[canaryIdentifier] = canaryValue
+        storage[keyIdentifier] = keyValue
+        return keyValue
     }
 }