bitwarden · fedemkr · Nov 4, 2024 · Oct 17, 2024 · Oct 18, 2024 · Oct 18, 2024
@@ -109,6 +109,9 @@ public class ServiceContainer: Services { // swiftlint:disable:this type_body_le
     /// The service for managing the polices for the user.
     let policyService: PolicyService
 
+    /// The helper used for app rehydration.
+    let rehydrationHelper: RehydrationHelper
+
     /// The repository used by the application to manage send data for the UI layer.
     public let sendRepository: SendRepository
 
@@ -185,6 +188,7 @@ public class ServiceContainer: Services { // swiftlint:disable:this type_body_le
     ///   - notificationCenterService: The service used by the application to access the system's notification center.
     ///   - notificationService: The service used by the application to handle notifications.
     ///   - pasteboardService: The service used by the application for sharing data with other apps.
+    ///   - rehydrationHelper: The helper used for app rehydration.
     ///   - policyService: The service for managing the polices for the user.
     ///   - sendRepository: The repository used by the application to manage send data for the UI layer.
     ///   - settingsRepository: The repository used by the application to manage data for the UI layer.
@@ -231,6 +235,7 @@ public class ServiceContainer: Services { // swiftlint:disable:this type_body_le
         notificationService: NotificationService,
         pasteboardService: PasteboardService,
         policyService: PolicyService,
+        rehydrationHelper: RehydrationHelper,
         sendRepository: SendRepository,
         settingsRepository: SettingsRepository,
         stateService: StateService,
@@ -275,6 +280,7 @@ public class ServiceContainer: Services { // swiftlint:disable:this type_body_le
         self.notificationService = notificationService
         self.pasteboardService = pasteboardService
         self.policyService = policyService
+        self.rehydrationHelper = rehydrationHelper
         self.sendRepository = sendRepository
         self.settingsRepository = settingsRepository
         self.stateService = stateService
@@ -324,6 +330,12 @@ public class ServiceContainer: Services { // swiftlint:disable:this type_body_le
             keychainRepository: keychainRepository
         )
 
+        let rehydrationHelper = DefaultRehydrationHelper(
+            errorReporter: errorReporter,
+            stateService: stateService,
+            timeProvider: timeProvider
+        )
+
         let environmentService = DefaultEnvironmentService(errorReporter: errorReporter, stateService: stateService)
         let collectionService = DefaultCollectionService(collectionDataStore: dataStore, stateService: stateService)
         let settingsService = DefaultSettingsService(settingsDataStore: dataStore, stateService: stateService)
@@ -675,6 +687,7 @@ public class ServiceContainer: Services { // swiftlint:disable:this type_body_le
             notificationService: notificationService,
             pasteboardService: pasteboardService,
             policyService: policyService,
+            rehydrationHelper: rehydrationHelper,
             sendRepository: sendRepository,
             settingsRepository: settingsRepository,
             stateService: stateService,

@@ -31,6 +31,7 @@ typealias Services = HasAPIService
     & HasOrganizationAPIService
     & HasPasteboardService
     & HasPolicyService
+    & HasRehydrationHelper
     & HasSendRepository
     & HasSettingsRepository
     & HasStateService
@@ -254,6 +255,12 @@ protocol HasPolicyService {
     var policyService: PolicyService { get }
 }
 
+/// Protocol for an object that provides a `RehydrationHelper`.
+protocol HasRehydrationHelper {
+    /// The helper for app rehydration.
+    var rehydrationHelper: RehydrationHelper { get }
+}
+
 /// Protocol for an object that provides a `SendRepository`.
 ///
 public protocol HasSendRepository {

@@ -112,6 +112,11 @@ protocol StateService: AnyObject {
     ///
     func getAllowSyncOnRefresh(userId: String?) async throws -> Bool
 
+    /// Gets the app rehydration state.
+    /// - Parameter userId: The user ID associated with this state.
+    /// - Returns: The rehydration state.
+    func getAppRehydrationState(userId: String?) async throws -> AppRehydrationState?
+
     /// Get the app theme.
     ///
     /// - Returns: The app theme.
@@ -564,6 +569,12 @@ protocol StateService: AnyObject {
     /// - Parameter config: The server config to use prior to user authentication.
     func setPreAuthServerConfig(config: ServerConfig) async
 
+    /// Sets the app rehydration state for the active account.
+    /// - Parameters:
+    ///   - rehydrationState: The app rehydration state.
+    ///   - userId: The user ID of the rehydration state.
+    func setAppRehydrationState(_ rehydrationState: AppRehydrationState?, userId: String?) async throws
+
     /// Sets the server configuration as provided by a server for a user ID.
     ///
     /// - Parameters:
@@ -775,6 +786,12 @@ extension StateService {
         try await getAllowSyncOnRefresh(userId: nil)
     }
 
+    /// Gets the app rehydration state for the active account.
+    /// - Returns: The rehydration state.
+    func getAppRehydrationState() async throws -> AppRehydrationState? {
+        try await getAppRehydrationState(userId: nil)
+    }
+
     /// Gets the clear clipboard value for the active account.
     ///
     /// - Returns: The clear clipboard value.
@@ -1083,6 +1100,14 @@ extension StateService {
         try await setPasswordGenerationOptions(options, userId: nil)
     }
 
+    /// Sets the app rehydration state for the active account.
+    ///
+    /// - Parameter rehydrationState: The app rehydration state.
+    ///
+    func setAppRehydrationState(_ rehydrationState: AppRehydrationState?) async throws {
+        try await setAppRehydrationState(rehydrationState, userId: nil)
+    }
+
     /// Sets the server config for the active account.
     ///
     /// - Parameter config: The server config.
@@ -1348,6 +1373,11 @@ actor DefaultStateService: StateService { // swiftlint:disable:this type_body_le
         return appSettingsStore.allowSyncOnRefresh(userId: userId)
     }
 
+    func getAppRehydrationState(userId: String?) async throws -> AppRehydrationState? {
+        let userId = try userId ?? getActiveAccountUserId()
+        return appSettingsStore.appRehydrationState(userId: userId)
+    }
+
     func getAppTheme() async -> AppTheme {
         AppTheme(appSettingsStore.appTheme)
     }
@@ -1714,6 +1744,11 @@ actor DefaultStateService: StateService { // swiftlint:disable:this type_body_le
         appSettingsStore.preAuthServerConfig = config
     }
 
+    func setAppRehydrationState(_ rehydrationState: AppRehydrationState?, userId: String?) async throws {
+        let userId = try userId ?? getActiveAccountUserId()
+        appSettingsStore.setAppRehydrationState(rehydrationState, userId: userId)
+    }
+
     func setServerConfig(_ config: ServerConfig?, userId: String?) async throws {
         let userId = try userId ?? getActiveAccountUserId()
         appSettingsStore.setServerConfig(config, userId: userId)

@@ -481,6 +481,24 @@ class StateServiceTests: BitwardenTestCase { // swiftlint:disable:this type_body
         XCTAssertFalse(value)
     }
 
+    /// `getAppRehydrationState(userId:)` returns the app rehydration state for the active account.
+    func test_getAppRehydrationState() async throws {
+        await subject.addAccount(.fixture())
+        appSettingsStore.appRehydrationState["1"] = AppRehydrationState(
+            target: .viewCipher(cipherId: "1"),
+            expirationTime: .now
+        )
+        let value = try await subject.getAppRehydrationState()
+        XCTAssertEqual(value?.target, .viewCipher(cipherId: "1"))
+    }
+
+    /// `getAppRehydrationState(userId:)` throws when there's no active account.
+    func test_getAppRehydrationState_throws() async throws {
+        await assertAsyncThrows(error: StateServiceError.noActiveAccount) {
+            _ = try await subject.getAppRehydrationState()
+        }
+    }
+
     /// `getClearClipboardValue()` returns the clear clipboard value for the active account.
     func test_getClearClipboardValue() async throws {
         await subject.addAccount(.fixture())
@@ -1408,6 +1426,27 @@ class StateServiceTests: BitwardenTestCase { // swiftlint:disable:this type_body
         XCTAssertEqual(appSettingsStore.allowSyncOnRefreshes["1"], true)
     }
 
+    /// `setAppRehydrationState(_:userId:)` sets the app rehydration state for the given account.
+    func test_setAppRehydrationState() async throws {
+        await subject.addAccount(.fixture())
+        try await subject.setAppRehydrationState(
+            AppRehydrationState(
+                target: .viewCipher(cipherId: "1"),
+                expirationTime: .now
+            ),
+            userId: "1"
+        )
+        let value = appSettingsStore.appRehydrationState["1"]
+        XCTAssertEqual(value?.target, .viewCipher(cipherId: "1"))
+    }
+
+    /// `setAppRehydrationState(_:userId:)` throws when there's no active account.
+    func test_setAppRehydrationState_throws() async throws {
+        await assertAsyncThrows(error: StateServiceError.noActiveAccount) {
+            _ = try await subject.setAppRehydrationState(nil)
+        }
+    }
+
     /// `setBiometricAuthenticationEnabled(isEnabled:)` sets biometric unlock preference for the default user.
     func test_setBiometricAuthenticationEnabled_default() async throws {
         await subject.addAccount(.fixture())

@@ -81,6 +81,11 @@ protocol AppSettingsStore: AnyObject {
     ///
     func allowSyncOnRefresh(userId: String) -> Bool
 
+    /// Gets the app rehydration state.
+    /// - Parameter userId: The user ID associated with this state.
+    /// - Returns: The rehydration state.
+    func appRehydrationState(userId: String) -> AppRehydrationState?
+
     /// Gets the time after which the clipboard should be cleared.
     ///
     /// - Parameter userId: The user ID associated with the clipboard clearing time.
@@ -252,6 +257,12 @@ protocol AppSettingsStore: AnyObject {
     ///
     func setAccountSetupVaultUnlock(_ vaultUnlockSetup: AccountSetupProgress?, userId: String)
 
+    /// Sets the app rehydration state to be used after timeout lock and user unlock.
+    /// - Parameters:
+    ///   - state: The state to save.
+    ///   - userId: The user ID the state belongs to.
+    func setAppRehydrationState(_ state: AppRehydrationState?, userId: String)
+
     /// Whether the vault should sync on refreshing.
     ///
     /// - Parameters:
@@ -649,6 +660,7 @@ extension DefaultAppSettingsStore: AppSettingsStore {
         case allowSyncOnRefresh(userId: String)
         case appId
         case appLocale
+        case appRehydrationState(userId: String)
         case appTheme
         case biometricAuthEnabled(userId: String)
         case clearClipboardValue(userId: String)
@@ -705,6 +717,8 @@ extension DefaultAppSettingsStore: AppSettingsStore {
                 key = "appId"
             case .appLocale:
                 key = "appLocale"
+            case let .appRehydrationState(userId):
+                key = "appRehydrationState_\(userId)"
             case .appTheme:
                 key = "theme"
             case let .biometricAuthEnabled(userId):
@@ -873,6 +887,10 @@ extension DefaultAppSettingsStore: AppSettingsStore {
         fetch(for: .allowSyncOnRefresh(userId: userId))
     }
 
+    func appRehydrationState(userId: String) -> AppRehydrationState? {
+        fetch(for: .appRehydrationState(userId: userId))
+    }
+
     func clearClipboardValue(userId: String) -> ClearClipboardValue {
         if let rawValue: Int = fetch(for: .clearClipboardValue(userId: userId)),
            let value = ClearClipboardValue(rawValue: rawValue) {
@@ -971,6 +989,10 @@ extension DefaultAppSettingsStore: AppSettingsStore {
         store(allowSyncOnRefresh, for: .allowSyncOnRefresh(userId: userId))
     }
 
+    func setAppRehydrationState(_ state: AppRehydrationState?, userId: String) {
+        store(state, for: .appRehydrationState(userId: userId))
+    }
+
     func setBiometricAuthenticationEnabled(_ isEnabled: Bool?, for userId: String) {
         store(isEnabled, for: .biometricAuthEnabled(userId: userId))
     }

@@ -154,6 +154,49 @@ class AppSettingsStoreTests: BitwardenTestCase { // swiftlint:disable:this type_
         XCTAssertNil(userDefaults.string(forKey: "bwPreferencesStorage:appLocale"))
     }
 
+    /// `appRehydrationState(userId:)` is initially `nil`
+    func test_appRehydrationState_isInitiallyNil() {
+        XCTAssertNil(subject.appRehydrationState(userId: "-1"))
+    }
+
+    /// `appRehydrationState(userId:)` is initially `nil`
+    func test_appRehydrationState_withValue() {
+        subject.setAppRehydrationState(
+            AppRehydrationState(target: .viewCipher(cipherId: "1"), expirationTime: .now),
+            userId: "1"
+        )
+        subject.setAppRehydrationState(
+            AppRehydrationState(target: .viewCipher(cipherId: "2"), expirationTime: .now),
+            userId: "2"
+        )
+
+        XCTAssertEqual(subject.appRehydrationState(userId: "1")?.target, .viewCipher(cipherId: "1"))
+        XCTAssertEqual(subject.appRehydrationState(userId: "2")?.target, .viewCipher(cipherId: "2"))
+
+        try XCTAssertEqual(
+            JSONDecoder().decode(
+                AppRehydrationState.self,
+                from: XCTUnwrap(
+                    userDefaults
+                        .string(forKey: "bwPreferencesStorage:appRehydrationState_1")?
+                        .data(using: .utf8)
+                )
+            ).target,
+            .viewCipher(cipherId: "1")
+        )
+        try XCTAssertEqual(
+            JSONDecoder().decode(
+                AppRehydrationState.self,
+                from: XCTUnwrap(
+                    userDefaults
+                        .string(forKey: "bwPreferencesStorage:appRehydrationState_2")?
+                        .data(using: .utf8)
+                )
+            ).target,
+            .viewCipher(cipherId: "2")
+        )
+    }
+
     /// `appTheme` returns `nil` if there isn't a previously stored value.
     func test_appTheme_isInitiallyNil() {
         XCTAssertNil(subject.appTheme)

@@ -3,14 +3,15 @@ import Foundation
 
 @testable import BitwardenShared
 
-class MockAppSettingsStore: AppSettingsStore {
+class MockAppSettingsStore: AppSettingsStore { // swiftlint:disable:this type_body_length
     var accountSetupAutofill = [String: AccountSetupProgress]()
     var accountSetupImportLogins = [String: AccountSetupProgress]()
     var accountSetupVaultUnlock = [String: AccountSetupProgress]()
     var addSitePromptShown = false
     var allowSyncOnRefreshes = [String: Bool]()
     var appId: String?
     var appLocale: String?
+    var appRehydrationState = [String: AppRehydrationState]()
     var appTheme: String?
     var disableWebIcons = false
     var introCarouselShown = false
@@ -74,6 +75,10 @@ class MockAppSettingsStore: AppSettingsStore {
         allowSyncOnRefreshes[userId] ?? false
     }
 
+    func appRehydrationState(userId: String) -> BitwardenShared.AppRehydrationState? {
+        appRehydrationState[userId]
+    }
+
     func clearClipboardValue(userId: String) -> ClearClipboardValue {
         clearClipboardValues[userId] ?? .never
     }
@@ -167,6 +172,14 @@ class MockAppSettingsStore: AppSettingsStore {
         allowSyncOnRefreshes[userId] = allowSyncOnRefresh
     }
 
+    func setAppRehydrationState(_ state: BitwardenShared.AppRehydrationState?, userId: String) {
+        guard state != nil else {
+            appRehydrationState.removeValue(forKey: userId)
+            return
+        }
+        appRehydrationState[userId] = state
+    }
+
     func setClearClipboardValue(_ clearClipboardValue: ClearClipboardValue?, userId: String) {
         clearClipboardValues[userId] = clearClipboardValue
     }