Prevent overwriting encryption keys

Sources/StytchCore/KeychainClient/KeychainClient.swift

@@ -3,7 +3,9 @@ import Foundation
 
 protocol KeychainClient: AnyObject {
     var encryptionKey: SymmetricKey? { get }
+    var didInitializeKeychainData: Bool { get }
 
+    func getEncryptionKey() throws
     func getQueryResults(item: KeychainItem) throws -> [KeychainQueryResult]
     func valueExistsForItem(item: KeychainItem) -> Bool
     func setValueForItem(value: KeychainItem.Value, item: KeychainItem) throws

Sources/StytchCore/KeychainClient/KeychainClientImplementation.swift

@@ -4,14 +4,39 @@ import Security
 #if !os(tvOS)
 import LocalAuthentication
 #endif
+#if os(iOS)
+import UIKit
+#endif
 
 let ENCRYPTEDUSERDEFAULTSKEYNAME = "EncryptedUserDefaultsKey"
 
 final class KeychainClientImplementation: KeychainClient {
     static let shared = KeychainClientImplementation()
     private let queue: DispatchQueue
     private let queueKey = DispatchSpecificKey<Void>()
-    var encryptionKey: SymmetricKey?
+    private var cachedEncryptionKey: SymmetricKey?
+    var didInitializeKeychainData = false
+    var encryptionKey: SymmetricKey? {
+        (try? safelyEnqueue {
+            if let cachedEncryptionKey {
+                return cachedEncryptionKey
+            }
+            #if os(iOS)
+            if UIApplication.shared.isProtectedDataAvailable {
+                try? getEncryptionKey()
+                didInitializeKeychainData = true
+            } else {
+                // For some reason, we are trying to read the encryption key before protected data became available
+                // Log that this happened (which it hopefully won't?), but leave the behavior up to the caller (EncryptedUserDefaultsClient) to handle a missing key (throw an error)
+                StytchConsoleLogger.error(message: "Attempted to read encryption key but UIApplication.shared.isProtectedDataAvailable was false")
+            }
+            #else
+            try? getEncryptionKey()
+            #endif
+            return cachedEncryptionKey
+        })
+    }
+
     private var isOnQueue: Bool {
         DispatchQueue.getSpecific(key: queueKey) != nil
     }
@@ -25,18 +50,11 @@ final class KeychainClientImplementation: KeychainClient {
     private init() {
         queue = DispatchQueue(label: "StytchKeychainClientQueue")
         queue.setSpecific(key: queueKey, value: ())
-        loadEncryptionKey()
         #if !os(tvOS) && !os(watchOS)
         contextWithoutUI.interactionNotAllowed = true
         #endif
     }
 
-    func loadEncryptionKey() {
-        try? safelyEnqueue {
-            encryptionKey = try? getEncryptionKey()
-        }
-    }
-
     func safelyEnqueue<T>(_ block: () throws -> T) throws -> T {
         if isOnQueue {
             return try block()
@@ -45,18 +63,19 @@ final class KeychainClientImplementation: KeychainClient {
         }
     }
 
-    private func getEncryptionKey() throws -> SymmetricKey {
+    func getEncryptionKey() throws {
         try safelyEnqueue {
             let result = try getFirstQueryResult(KeychainItem.encryptionKey)
             guard let result else {
-                // Key doesn't exist so create it
+                // At this point, we know that protected data IS available, so if the keychain returned nil, then it means the key TRULY doesn't exist, and so we should create a new one
                 let data = SymmetricKey(size: .bits256).withUnsafeBytes {
                     Data(Array($0))
                 }
                 try setValueForItem(value: .init(data: data, account: ENCRYPTEDUSERDEFAULTSKEYNAME, label: nil, generic: nil, accessPolicy: nil), item: .encryptionKey)
-                return SymmetricKey(data: data)
+                cachedEncryptionKey = SymmetricKey(data: data)
+                return
             }
-            return SymmetricKey(data: result.data)
+            cachedEncryptionKey = SymmetricKey(data: result.data)
         }
     }
 

Sources/StytchCore/StytchClientCommon.swift

@@ -69,6 +69,13 @@ extension StytchClientCommonInternal {
                 }
             }
         }
+        NotificationCenter.default.addObserver(forName: UIApplication.protectedDataDidBecomeAvailableNotification, object: nil, queue: nil) { _ in
+            Task {
+                if !Current.keychainClient.didInitializeKeychainData {
+                    try? Current.keychainClient.getEncryptionKey()
+                }
+            }
+        }
         #endif
 
         Task {

Tests/StytchCoreTests/KeychainClient+Mock.swift

@@ -6,6 +6,12 @@ import Foundation
 public var keychainDateCreatedOffsetInMinutes = 0
 
 class KeychainClientMock: KeychainClient {
+    var didInitializeKeychainData: Bool = true
+
+    func getEncryptionKey() throws {
+        // noop
+    }
+
     var encryptionKey: SymmetricKey? {
         do {
             return SymmetricKey(data: try Current.cryptoClient.dataWithRandomBytesOfCount(256))