Refactor FXIOS-11628 autofill keychain calls (#25152)

lougeniaC64 · web-flow · commit 55dde867aada · 2025-03-14T21:09:58.000+01:00
Replaced credit cards MZKeychainWrapper calls with RustKeychain calls
diff --git a/firefox-ios/Client/Frontend/Settings/Main/Debug/FeatureFlags/FeatureFlagsDebugViewController.swift b/firefox-ios/Client/Frontend/Settings/Main/Debug/FeatureFlags/FeatureFlagsDebugViewController.swift
@@ -146,7 +146,7 @@ final class FeatureFlagsDebugViewController: SettingsTableViewController, Featur
                 FeatureFlagsBoolSetting(
                     with: .useRustKeychain,
                     titleText: format(string: "Enable Rust Keychain"),
-                    statusText: format(string: "Toggle to enable Rust Components rust keychain")
+                    statusText: format(string: "Toggle to enable rust keychain")
                 ) { [weak self] _ in
                     self?.reloadView()
                 },
diff --git a/firefox-ios/Providers/Profile.swift b/firefox-ios/Providers/Profile.swift
@@ -482,7 +482,7 @@ open class BrowserProfile: Profile {
         isDirectory: true
     ).appendingPathComponent("autofill.db").path
 
-    lazy var autofill = RustAutofill(databasePath: autofillDbPath)
+    lazy var autofill = RustAutofill(databasePath: autofillDbPath, rustKeychainEnabled: rustKeychainEnabled)
 
     #if !MOZ_TARGET_NOTIFICATIONSERVICE && !MOZ_TARGET_SHARETO && !MOZ_TARGET_CREDENTIAL_PROVIDER
     lazy var searchEnginesManager: SearchEnginesManager = {
@@ -768,38 +768,47 @@ open class BrowserProfile: Profile {
         let rustLoginsKeys = RustLoginEncryptionKeys()
         let rustAutofillKey = RustAutofillEncryptionKeys()
         var loginsKey: String?
-        let creditCardKey = legacyKeychain.string(forKey: rustAutofillKey.ccKeychainKey)
+        var loginsCanary: String?
+
+        var creditCardKey: String?
+        var creditCardCanary: String?
 
         if rustKeychainEnabled {
-            (loginsKey, _) = keychain.getLoginsKeyData()
+            (creditCardKey, creditCardCanary) = keychain.getCreditCardKeyData()
+            (loginsKey, loginsCanary) = 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)
+            if let creditCardKey = creditCardKey, let creditCardCanary = creditCardCanary {
+                keychain.setCreditCardsKeyData(keyValue: creditCardKey, canaryValue: creditCardCanary)
+            }
+
+            if let loginsKey = loginsKey, let loginsCanary = loginsCanary {
+                keychain.setLoginsKeyData(keyValue: loginsKey, canaryValue: loginsCanary)
             }
         } else {
+            creditCardKey = legacyKeychain.string(forKey: rustAutofillKey.ccKeychainKey)
             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 creditCardKey = creditCardKey {
+                legacyKeychain.set(creditCardKey,
+                                   forKey: rustAutofillKey.ccKeychainKey,
+                                   withAccessibility: .afterFirstUnlock)
+            }
+
             if let loginsKey = loginsKey {
                 legacyKeychain.set(loginsKey,
                                    forKey: rustLoginsKeys.loginPerFieldKeychainKey,
                                    withAccessibility: .afterFirstUnlock)
             }
         }
 
-        if let creditCardKey = creditCardKey {
-            legacyKeychain.set(creditCardKey,
-                               forKey: rustAutofillKey.ccKeychainKey,
-                               withAccessibility: .afterFirstUnlock)
-        }
-
         // Tell any observers that our account has changed.
         NotificationCenter.default.post(name: .FirefoxAccountChanged, object: nil)
 
diff --git a/firefox-ios/Storage/MockRustKeychain.swift b/firefox-ios/Storage/MockRustKeychain.swift
@@ -13,7 +13,6 @@ class MockRustKeychain: RustKeychain {
     private init() {
         super.init(serviceName: "Test")
     }
-
     override public func removeObject(key: String) {
         _ = storage.removeValue(forKey: key)
     }
@@ -22,14 +21,24 @@ class MockRustKeychain: RustKeychain {
         storage.removeAll()
     }
 
-    override public func setLoginsKey(_ value: String) {
-        storage[loginsKeyIdentifier] = value
+    override public func setLoginsKeyData(keyValue: String, canaryValue: String) {
+        storage[loginsKeyIdentifier] = keyValue
+        storage[loginsCanaryKeyIdentifier] = canaryValue
+    }
+
+    override public func setCreditCardsKeyData(keyValue: String, canaryValue: String) {
+        storage[creditCardKeyIdentifier] = keyValue
+        storage[creditCardCanaryKeyIdentifier] = canaryValue
     }
 
     override public func getLoginsKeyData() -> (String?, String?) {
         return (storage[loginsKeyIdentifier], storage[loginsCanaryKeyIdentifier])
     }
 
+    override public func getCreditCardKeyData() -> (String?, String?) {
+        return (storage[creditCardKeyIdentifier], storage[creditCardCanaryKeyIdentifier])
+    }
+
     override public class func wipeKeychain() {}
 
     override func createLoginsKeyData() throws -> String {
diff --git a/firefox-ios/Storage/Rust/RustAutofill.swift b/firefox-ios/Storage/Rust/RustAutofill.swift
@@ -34,6 +34,7 @@ public class RustAutofill {
 
     private(set) var isOpen = false
     private var didAttemptToMoveToBackup = false
+    private var rustKeychainEnabled = false
     private let logger: Logger
 
     // MARK: - Initialization
@@ -43,10 +44,13 @@ public class RustAutofill {
     /// - Parameters:
     ///   - databasePath: The path to the Autofill database file.
     ///   - logger: An optional logger for recording informational and error messages. Default is shared DefaultLogger.
-    public init(databasePath: String, logger: Logger = DefaultLogger.shared) {
+    public init(databasePath: String,
+                logger: Logger = DefaultLogger.shared,
+                rustKeychainEnabled: Bool = false) {
         self.databasePath = databasePath
         queue = DispatchQueue(label: "RustAutofill queue: \(databasePath)")
         self.logger = logger
+        self.rustKeychainEnabled = rustKeychainEnabled
     }
 
     // MARK: - Database Operations
@@ -140,7 +144,7 @@ public class RustAutofill {
             return nil
         }
         let keys = RustAutofillEncryptionKeys()
-        return keys.decryptCreditCardNum(encryptedCCNum: encryptedCCNum)
+        return keys.decryptCreditCardNum(encryptedCCNum: encryptedCCNum, rustKeychainEnabled: rustKeychainEnabled)
     }
 
     /// Retrieves a credit card from the database by its identifier.
@@ -519,7 +523,7 @@ public class RustAutofill {
                     // There are no records in the database so we don't need to scrub any
                     // existing credit card records. We just need to create a new key.
                     do {
-                        let key = try rustKeys.createAndStoreKey()
+                        let key = try self.createKey(rustKeys: rustKeys)
                         completion(.success(key))
                     } catch let error as NSError {
                         completion(.failure(error))
@@ -570,9 +574,14 @@ public class RustAutofill {
     private func getKeychainData(rustKeys: RustAutofillEncryptionKeys,
                                  completion: @escaping (String?, String?) -> Void) {
         DispatchQueue.global(qos: .background).sync {
-            let key = rustKeys.keychain.string(forKey: rustKeys.ccKeychainKey)
-            let encryptedCanaryPhrase = rustKeys.keychain.string(forKey: rustKeys.ccCanaryPhraseKey)
-            completion(key, encryptedCanaryPhrase)
+            if rustKeychainEnabled {
+                let (key, encryptedCanaryPhrase) = rustKeys.keychain.getCreditCardKeyData()
+                completion(key, encryptedCanaryPhrase)
+            } else {
+                let key = rustKeys.legacyKeychain.string(forKey: rustKeys.ccKeychainKey)
+                let encryptedCanaryPhrase = rustKeys.legacyKeychain.string(forKey: rustKeys.ccCanaryPhraseKey)
+                completion(key, encryptedCanaryPhrase)
+            }
         }
     }
 
@@ -582,7 +591,7 @@ public class RustAutofill {
             switch result {
             case .success(()):
                 do {
-                    let key = try rustKeys.createAndStoreKey()
+                    let key = try self.createKey(rustKeys: rustKeys)
                     completion(.success(key))
                 } catch let error as NSError {
                     self.logger.log("Error creating credit card encryption key",
@@ -597,6 +606,12 @@ public class RustAutofill {
         }
     }
 
+    private func createKey(rustKeys: RustAutofillEncryptionKeys) throws -> String {
+        return self.rustKeychainEnabled ?
+            try rustKeys.keychain.createCreditCardsKeyData() :
+            try rustKeys.createAndStoreKey()
+    }
+
     private func hasCreditCards(completion: @escaping (Result<Bool, Error>) -> Void) {
         return listCreditCards { (creditCards, err) in
             guard err == nil else {
diff --git a/firefox-ios/Storage/Rust/RustAutofillEncryptionKeys.swift b/firefox-ios/Storage/Rust/RustAutofillEncryptionKeys.swift
@@ -27,7 +27,8 @@ public extension AutofillApiError {
 public class RustAutofillEncryptionKeys {
     public let ccKeychainKey = "appservices.key.creditcard.perfield"
 
-    let keychain = MZKeychainWrapper.sharedClientAppContainerKeychain
+    let legacyKeychain = MZKeychainWrapper.sharedClientAppContainerKeychain
+    let keychain = RustKeychain.sharedClientAppContainerKeychain
     let ccCanaryPhraseKey = "creditCardCanaryPhrase"
     let canaryPhrase = "a string for checking validity of the key"
 
@@ -43,20 +44,20 @@ public class RustAutofillEncryptionKeys {
             let canary = try self.createCanary(text: canaryPhrase, key: secret)
 
             DispatchQueue.global(qos: .background).sync {
-                keychain.set(secret,
-                             forKey: ccKeychainKey,
-                             withAccessibility: MZKeychainItemAccessibility.afterFirstUnlock)
-                keychain.set(canary,
-                             forKey: ccCanaryPhraseKey,
-                             withAccessibility: MZKeychainItemAccessibility.afterFirstUnlock)
+                legacyKeychain.set(secret,
+                                   forKey: ccKeychainKey,
+                                   withAccessibility: MZKeychainItemAccessibility.afterFirstUnlock)
+                legacyKeychain.set(canary,
+                                   forKey: ccCanaryPhraseKey,
+                                   withAccessibility: MZKeychainItemAccessibility.afterFirstUnlock)
             }
 
             return secret
         } catch let err as NSError {
             if let autofillStoreError = err as? AutofillApiError {
-                logAutofillStoreError(err: autofillStoreError,
-                                      errorDomain: err.domain,
-                                      errorMessage: "Error while creating and storing credit card key")
+                keychain.logAutofillStoreError(err: autofillStoreError,
+                                               errorDomain: err.domain,
+                                               errorMessage: "Error while creating and storing credit card key")
 
                 throw AutofillEncryptionKeyError.noKeyCreated
             } else {
@@ -70,16 +71,24 @@ public class RustAutofillEncryptionKeys {
         }
     }
 
-    func decryptCreditCardNum(encryptedCCNum: String) -> String? {
-        guard let key = self.keychain.string(forKey: self.ccKeychainKey) else { return nil }
+    func decryptCreditCardNum(encryptedCCNum: String, rustKeychainEnabled: Bool) -> String? {
+        var keyValue: String?
+
+        if rustKeychainEnabled {
+            (keyValue, _) = keychain.getCreditCardKeyData()
+        } else {
+            keyValue = legacyKeychain.string(forKey: self.ccKeychainKey)
+        }
+
+        guard let key = keyValue else { return nil }
 
         do {
             return try decryptString(key: key, ciphertext: encryptedCCNum)
         } catch let err as NSError {
             if let autofillStoreError = err as? AutofillApiError {
-                logAutofillStoreError(err: autofillStoreError,
-                                      errorDomain: err.domain,
-                                      errorMessage: "Error while decrypting credit card")
+                keychain.logAutofillStoreError(err: autofillStoreError,
+                                               errorDomain: err.domain,
+                                               errorMessage: "Error while decrypting credit card")
             } else {
                 logger.log("Unknown error while decrypting credit card",
                            level: .warning,
@@ -100,25 +109,4 @@ public class RustAutofillEncryptionKeys {
                               key: String) throws -> String {
         return try encryptString(key: key, cleartext: text)
     }
-
-    private func logAutofillStoreError(err: AutofillApiError,
-                                       errorDomain: String,
-                                       errorMessage: String) {
-        var message: String {
-            switch err {
-            case .SqlError(let message),
-                    .CryptoError(let message),
-                    .NoSuchRecord(let message),
-                    .UnexpectedAutofillApiError(let message):
-                return message
-            case .InterruptedError:
-                return "Interrupted Error"
-            }
-        }
-
-        logger.log(errorMessage,
-                   level: .warning,
-                   category: .storage,
-                   description: "\(errorDomain) - \(err.descriptionValue): \(message)")
-    }
 }
diff --git a/firefox-ios/Storage/RustKeychain.swift b/firefox-ios/Storage/RustKeychain.swift
@@ -45,6 +45,10 @@ open class RustKeychain {
     public let loginsCanaryKeyIdentifier = "canaryPhrase"
     let loginsCanaryPhrase = "a string for checking validity of the key"
 
+    public let creditCardKeyIdentifier = "appservices.key.creditcard.perfield"
+    public let creditCardCanaryKeyIdentifier = "creditCardCanaryPhrase"
+    let creditCardCanaryPhrase = "a string for checking validity of the key"
+
     public static var sharedClientAppContainerKeychain: RustKeychain {
         let baseBundleIdentifier = AppInfo.baseBundleIdentifier
 
@@ -85,14 +89,25 @@ open class RustKeychain {
         logErrorFromStatus(status, errMsg: "Failed to remove all keys")
     }
 
-    public func setLoginsKey(_ value: String) {
-        addOrUpdateKeychainKey(value, key: loginsKeyIdentifier)
+    public func setLoginsKeyData(keyValue: String, canaryValue: String) {
+        addOrUpdateKeychainKey(keyValue, key: loginsKeyIdentifier)
+        addOrUpdateKeychainKey(canaryValue, key: loginsCanaryKeyIdentifier)
+    }
+
+    public func setCreditCardsKeyData(keyValue: String, canaryValue: String) {
+        addOrUpdateKeychainKey(keyValue, key: creditCardKeyIdentifier)
+        addOrUpdateKeychainKey(canaryValue, key: creditCardCanaryKeyIdentifier)
     }
 
     public func getLoginsKeyData() -> (String?, String?) {
         return getEncryptionKeyData(keyIdentifier: loginsKeyIdentifier, canaryKeyIdentifier: loginsCanaryKeyIdentifier)
     }
 
+    public func getCreditCardKeyData() -> (String?, String?) {
+        return getEncryptionKeyData(keyIdentifier: creditCardKeyIdentifier,
+                                    canaryKeyIdentifier: creditCardCanaryKeyIdentifier)
+    }
+
     public class func wipeKeychain() {
         deleteKeychainSecClass(kSecClassGenericPassword) // Generic password items
         deleteKeychainSecClass(kSecClassInternetPassword) // Internet password items
@@ -120,6 +135,26 @@ open class RustKeychain {
                    description: "\(errorDomain) - \(err.descriptionValue): \(message)")
     }
 
+    func createCreditCardsKeyData() throws -> String {
+        do {
+            return try createAndStoreKey(canaryPhrase: creditCardCanaryPhrase,
+                                         canaryIdentifier: creditCardCanaryKeyIdentifier,
+                                         keyIdentifier: creditCardKeyIdentifier)
+        } catch let err as NSError {
+            if let autofillStoreError = err as? AutofillApiError {
+                logAutofillStoreError(err: autofillStoreError,
+                                      errorDomain: err.domain,
+                                      errorMessage: "Error while creating and storing credit card key")
+            } else {
+                logger.log("Unknown error while creating and storing credit card key",
+                           level: .warning,
+                           category: .storage,
+                           description: err.localizedDescription)
+            }
+        }
+        throw LoginEncryptionKeyError.noKeyCreated
+    }
+
     func createLoginsKeyData() throws -> String {
         do {
             return try createAndStoreKey(canaryPhrase: loginsCanaryPhrase,
diff --git a/firefox-ios/firefox-ios-tests/Tests/StorageTests/RustAutofillTests.swift b/firefox-ios/firefox-ios-tests/Tests/StorageTests/RustAutofillTests.swift
