PM-37879: feat: Add SDK state bridge (#7204)

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/datasource/disk/AuthDiskSource.kt

@@ -2,6 +2,7 @@ package com.x8bit.bitwarden.data.auth.datasource.disk
 
 import com.bitwarden.core.WrappedAccountCryptographicState
 import com.bitwarden.network.model.SyncResponseJson
+import com.bitwarden.network.model.V2UpgradeTokenJson
 import com.bitwarden.network.provider.AppIdProvider
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountTokensJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.OnboardingStatus
@@ -410,4 +411,14 @@ interface AuthDiskSource : AppIdProvider {
      * Stores the last lock timestamp for the given [userId].
      */
     fun storeLastLockTimestamp(userId: String, lastLockTimestamp: Instant?)
+
+    /**
+     * Gets the v2 upgrade token for the given [userId].
+     */
+    fun getV2UpgradeToken(userId: String): V2UpgradeTokenJson?
+
+    /**
+     * Stores the v2 upgrade token for the given [userId].
+     */
+    fun storeV2UpgradeToken(userId: String, v2UpgradeToken: V2UpgradeTokenJson?)
 }

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/datasource/disk/AuthDiskSourceImpl.kt

@@ -8,6 +8,7 @@ import com.bitwarden.core.data.util.decodeFromStringOrNull
 import com.bitwarden.data.datasource.disk.BaseEncryptedDiskSource
 import com.bitwarden.network.model.AccountKeysJson
 import com.bitwarden.network.model.SyncResponseJson
+import com.bitwarden.network.model.V2UpgradeTokenJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountTokensJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.OnboardingStatus
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.PendingAuthRequestJson
@@ -58,6 +59,7 @@ private const val SHOW_IMPORT_LOGINS_KEY = "showImportLogins"
 private const val LAST_LOCK_TIMESTAMP = "lastLockTimestamp"
 private const val PROFILE_ACCOUNT_KEYS_KEY = "profileAccountKeys"
 private const val ACCOUNT_CRYPTOGRAPHIC_STATE_KEY = "accountCryptographicState"
+private const val V2_UPGRADE_TOKEN = "v2UpgradeToken"
 
 /**
  * Primary implementation of [AuthDiskSource].
@@ -192,6 +194,7 @@ class AuthDiskSourceImpl(
             userId = userId,
             pinProtectedUserKeyEnvelope = null,
         )
+        storeV2UpgradeToken(userId = userId, v2UpgradeToken = null)
 
         // Certain values are never removed as required by the feature requirements:
         // * DeviceKey
@@ -603,6 +606,20 @@ class AuthDiskSourceImpl(
         )
     }
 
+    override fun getV2UpgradeToken(
+        userId: String,
+    ): V2UpgradeTokenJson? =
+        getString(key = V2_UPGRADE_TOKEN.appendIdentifier(identifier = userId))?.let {
+            json.decodeFromStringOrNull<V2UpgradeTokenJson>(string = it)
+        }
+
+    override fun storeV2UpgradeToken(userId: String, v2UpgradeToken: V2UpgradeTokenJson?) {
+        putString(
+            key = V2_UPGRADE_TOKEN.appendIdentifier(identifier = userId),
+            value = v2UpgradeToken?.let { json.encodeToString(value = it) },
+        )
+    }
+
     private fun generateAndStoreUniqueAppId(): String =
         UUID
             .randomUUID()

app/src/main/kotlin/com/x8bit/bitwarden/data/platform/manager/SdkClientManagerImpl.kt

@@ -16,7 +16,7 @@ import kotlinx.coroutines.runBlocking
 /**
  * Primary implementation of [SdkClientManager].
  */
-class SdkClientManagerImpl(
+internal class SdkClientManagerImpl(
     nativeLibraryManager: NativeLibraryManager,
     dispatcherManager: DispatcherManager,
     sdkRepoFactory: SdkRepositoryFactory,
@@ -42,6 +42,9 @@ class SdkClientManagerImpl(
                 platform().state().registerClientManagedRepositories(
                     repositories = sdkRepoFactory.getRepositories(userId = userId),
                 )
+                userId?.let {
+                    kmStateBridge().registerBridgeImpl(sdkRepoFactory.getStateBridge(userId = it))
+                }
             }
     },
 ) : SdkClientManager {

app/src/main/kotlin/com/x8bit/bitwarden/data/platform/manager/sdk/SdkRepositoryFactory.kt

@@ -2,6 +2,7 @@ package com.x8bit.bitwarden.data.platform.manager.sdk
 
 import com.bitwarden.core.ClientManagedTokens
 import com.bitwarden.core.ClientSettings
+import com.bitwarden.core.StateBridgeForeignImpl
 import com.bitwarden.sdk.Repositories
 import com.bitwarden.sdk.ServerCommunicationConfigRepository
 
@@ -9,6 +10,11 @@ import com.bitwarden.sdk.ServerCommunicationConfigRepository
  * Creates and manages sdk repositories.
  */
 interface SdkRepositoryFactory {
+    /**
+     * Creates a [StateBridgeForeignImpl] for use with the Bitwarden SDK.
+     */
+    fun getStateBridge(userId: String): StateBridgeForeignImpl
+
     /**
      * Retrieves or creates a [Repositories] for use with the Bitwarden SDK.
      */

app/src/main/kotlin/com/x8bit/bitwarden/data/platform/manager/sdk/SdkRepositoryFactoryImpl.kt

@@ -3,6 +3,7 @@ package com.x8bit.bitwarden.data.platform.manager.sdk
 import com.bitwarden.core.ClientManagedTokens
 import com.bitwarden.core.ClientSettings
 import com.bitwarden.core.DeviceType
+import com.bitwarden.core.StateBridgeForeignImpl
 import com.bitwarden.data.datasource.disk.ConfigDiskSource
 import com.bitwarden.network.model.BitwardenServiceClientConfig
 import com.bitwarden.sdk.Repositories
@@ -13,6 +14,7 @@ import com.x8bit.bitwarden.data.platform.manager.sdk.repository.SdkCipherReposit
 import com.x8bit.bitwarden.data.platform.manager.sdk.repository.SdkLocalUserDataKeyStateRepository
 import com.x8bit.bitwarden.data.platform.manager.sdk.repository.SdkTokenRepository
 import com.x8bit.bitwarden.data.platform.manager.sdk.repository.ServerCommunicationConfigRepositoryImpl
+import com.x8bit.bitwarden.data.platform.manager.sdk.statebridge.SdkStateBridge
 import com.x8bit.bitwarden.data.vault.datasource.disk.VaultDiskSource
 
 /**
@@ -25,6 +27,14 @@ class SdkRepositoryFactoryImpl(
     private val authDiskSource: AuthDiskSource,
     private val serviceClientConfig: BitwardenServiceClientConfig,
 ) : SdkRepositoryFactory {
+    override fun getStateBridge(
+        userId: String,
+    ): StateBridgeForeignImpl =
+        SdkStateBridge(
+            userId = userId,
+            authDiskSource = authDiskSource,
+        )
+
     override fun getRepositories(userId: String?): Repositories =
         Repositories(
             cipher = getSdkCipherRepository(userId = userId),

app/src/main/kotlin/com/x8bit/bitwarden/data/platform/manager/sdk/statebridge/SdkStateBridge.kt

@@ -0,0 +1,137 @@
+package com.x8bit.bitwarden.data.platform.manager.sdk.statebridge
+
+import com.bitwarden.core.MasterPasswordUnlockData
+import com.bitwarden.core.StateBridgeForeignImpl
+import com.bitwarden.core.V2UpgradeToken
+import com.bitwarden.core.WrappedAccountCryptographicState
+import com.bitwarden.crypto.EncString
+import com.bitwarden.crypto.PasswordProtectedKeyEnvelope
+import com.bitwarden.crypto.SymmetricCryptoKey
+import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
+import com.x8bit.bitwarden.data.auth.repository.util.updateMasterPasswordUnlock
+import com.x8bit.bitwarden.data.vault.repository.util.toSdkMasterPasswordUnlock
+import com.x8bit.bitwarden.data.vault.repository.util.toV2UpgradeToken
+import com.x8bit.bitwarden.data.vault.repository.util.toV2UpgradeTokenJson
+
+/**
+ * A user-scoped implementation of a Bitwarden SDK [StateBridgeForeignImpl].
+ */
+@Suppress("TooManyFunctions")
+internal class SdkStateBridge(
+    private val userId: String,
+    private val authDiskSource: AuthDiskSource,
+) : StateBridgeForeignImpl {
+    @Volatile
+    private var inMemoryUserKey: SymmetricCryptoKey? = null
+
+    override suspend fun setUserKey(value: SymmetricCryptoKey) {
+        inMemoryUserKey = value
+    }
+
+    override suspend fun getUserKey(): SymmetricCryptoKey? = inMemoryUserKey
+
+    override suspend fun clearUserKey() {
+        inMemoryUserKey = null
+    }
+
+    override suspend fun setPersistentPinEnvelope(value: PasswordProtectedKeyEnvelope) {
+        authDiskSource.storePersistentPinProtectedUserKeyEnvelope(
+            userId = userId,
+            pinProtectedUserKeyEnvelope = value,
+        )
+    }
+
+    override suspend fun getPersistentPinEnvelope(): PasswordProtectedKeyEnvelope? =
+        authDiskSource.getPersistentPinProtectedUserKeyEnvelope(userId = userId)
+
+    override suspend fun clearPersistentPinEnvelope() {
+        authDiskSource.storePersistentPinProtectedUserKeyEnvelope(
+            userId = userId,
+            pinProtectedUserKeyEnvelope = null,
+        )
+    }
+
+    override suspend fun setEphemeralPinEnvelope(value: PasswordProtectedKeyEnvelope) {
+        authDiskSource.storeEphemeralPinProtectedUserKeyEnvelope(
+            userId = userId,
+            pinProtectedUserKeyEnvelope = value,
+        )
+    }
+
+    override suspend fun getEphemeralPinEnvelope(): PasswordProtectedKeyEnvelope? =
+        authDiskSource.getEphemeralPinProtectedUserKeyEnvelope(userId = userId)
+
+    override suspend fun clearEphemeralPinEnvelope() {
+        authDiskSource.storeEphemeralPinProtectedUserKeyEnvelope(
+            userId = userId,
+            pinProtectedUserKeyEnvelope = null,
+        )
+    }
+
+    override suspend fun setEncryptedPin(value: EncString) {
+        authDiskSource.storeEncryptedPin(userId = userId, encryptedPin = value)
+    }
+
+    override suspend fun getEncryptedPin(): EncString? =
+        authDiskSource.getEncryptedPin(userId = userId)
+
+    override suspend fun clearEncryptedPin() {
+        authDiskSource.storeEncryptedPin(userId = userId, encryptedPin = null)
+    }
+
+    override suspend fun setV2UpgradeToken(value: V2UpgradeToken) {
+        authDiskSource.storeV2UpgradeToken(
+            userId = userId,
+            v2UpgradeToken = value.toV2UpgradeTokenJson(),
+        )
+    }
+
+    override suspend fun getV2UpgradeToken(): V2UpgradeToken? =
+        authDiskSource
+            .getV2UpgradeToken(userId = userId)
+            ?.toV2UpgradeToken()
+
+    override suspend fun clearV2UpgradeToken() {
+        authDiskSource.storeV2UpgradeToken(userId = userId, v2UpgradeToken = null)
+    }
+
+    override suspend fun setAccountCryptographicState(value: WrappedAccountCryptographicState) {
+        authDiskSource.storeAccountCryptographicState(
+            userId = userId,
+            accountCryptographicState = value,
+        )
+    }
+
+    override suspend fun getAccountCryptographicState(): WrappedAccountCryptographicState? =
+        authDiskSource.getAccountCryptographicState(userId = userId)
+
+    override suspend fun clearAccountCryptographicState() {
+        authDiskSource.storeAccountCryptographicState(
+            userId = userId,
+            accountCryptographicState = null,
+        )
+    }
+
+    override suspend fun setMasterpasswordUnlockData(value: MasterPasswordUnlockData) {
+        authDiskSource.userState = authDiskSource.userState?.updateMasterPasswordUnlock(
+            userId = userId,
+            masterPasswordUnlock = value,
+        )
+    }
+
+    override suspend fun getMasterpasswordUnlockData(): MasterPasswordUnlockData? =
+        authDiskSource
+            .userState
+            ?.accounts[userId]
+            ?.profile
+            ?.userDecryptionOptions
+            ?.masterPasswordUnlock
+            ?.toSdkMasterPasswordUnlock()
+
+    override suspend fun clearMasterpasswordUnlockData() {
+        authDiskSource.userState = authDiskSource.userState?.updateMasterPasswordUnlock(
+            userId = userId,
+            masterPasswordUnlock = null,
+        )
+    }
+}

app/src/main/kotlin/com/x8bit/bitwarden/data/vault/manager/VaultLockManagerImpl.kt

@@ -42,6 +42,7 @@ import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockData
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockResult
 import com.x8bit.bitwarden.data.vault.repository.util.logTag
 import com.x8bit.bitwarden.data.vault.repository.util.statusFor
+import com.x8bit.bitwarden.data.vault.repository.util.toV2UpgradeToken
 import com.x8bit.bitwarden.data.vault.repository.util.toVaultUnlockResult
 import com.x8bit.bitwarden.data.vault.repository.util.update
 import kotlinx.coroutines.CoroutineScope
@@ -192,7 +193,9 @@ class VaultLockManagerImpl(
                             email = email,
                             method = initUserCryptoMethod,
                             userId = userId,
-                            upgradeToken = null,
+                            upgradeToken = authDiskSource
+                                .getV2UpgradeToken(userId = userId)
+                                ?.toV2UpgradeToken(),
                         ),
                     )
                     .flatMap { result ->

app/src/main/kotlin/com/x8bit/bitwarden/data/vault/manager/VaultSyncManagerImpl.kt

@@ -328,6 +328,10 @@ class VaultSyncManagerImpl(
                     authDiskSource.userState = authDiskSource.userState?.toUpdatedUserStateJson(
                         syncResponse = syncResponse,
                     )
+                    authDiskSource.storeV2UpgradeToken(
+                        userId = userId,
+                        v2UpgradeToken = syncResponse.userDecryption?.v2UpgradeToken,
+                    )
 
                     unlockVaultForOrganizationsIfNecessary(syncResponse = syncResponse)
                     storeProfileData(syncResponse = syncResponse)

app/src/main/kotlin/com/x8bit/bitwarden/data/vault/repository/util/V2UpgradeTokenExtensions.kt

@@ -0,0 +1,22 @@
+package com.x8bit.bitwarden.data.vault.repository.util
+
+import com.bitwarden.core.V2UpgradeToken
+import com.bitwarden.network.model.V2UpgradeTokenJson
+
+/**
+ * Converts the [V2UpgradeToken] into a [V2UpgradeTokenJson].
+ */
+fun V2UpgradeToken.toV2UpgradeTokenJson(): V2UpgradeTokenJson =
+    V2UpgradeTokenJson(
+        wrappedUserKey1 = this.wrappedUserKey1,
+        wrappedUserKey2 = this.wrappedUserKey2,
+    )
+
+/**
+ * Converts the [V2UpgradeTokenJson] into a [V2UpgradeToken].
+ */
+fun V2UpgradeTokenJson.toV2UpgradeToken(): V2UpgradeToken =
+    V2UpgradeToken(
+        wrappedUserKey1 = this.wrappedUserKey1,
+        wrappedUserKey2 = this.wrappedUserKey2,
+    )

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/datasource/disk/AuthDiskSourceTest.kt

@@ -11,6 +11,7 @@ import com.bitwarden.network.model.KdfTypeJson
 import com.bitwarden.network.model.KeyConnectorUserDecryptionOptionsJson
 import com.bitwarden.network.model.TrustedDeviceUserDecryptionOptionsJson
 import com.bitwarden.network.model.UserDecryptionOptionsJson
+import com.bitwarden.network.model.V2UpgradeTokenJson
 import com.bitwarden.network.model.createMockOrganizationNetwork
 import com.bitwarden.network.model.createMockPolicy
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountJson
@@ -327,11 +328,17 @@ class AuthDiskSourceTest {
             userId = userId,
             authenticatorSyncUnlockKey = "authenticatorSyncUnlockKey",
         )
-
         authDiskSource.storeOnboardingStatus(
             userId = userId,
             onboardingStatus = OnboardingStatus.AUTOFILL_SETUP,
         )
+        authDiskSource.storeV2UpgradeToken(
+            userId = userId,
+            v2UpgradeToken = V2UpgradeTokenJson(
+                wrappedUserKey1 = "wrappedUserKey1",
+                wrappedUserKey2 = "wrappedUserKey2",
+            ),
+        )
 
         authDiskSource.clearData(userId = userId)
 
@@ -364,6 +371,7 @@ class AuthDiskSourceTest {
         assertNull(authDiskSource.getPinProtectedUserKeyEnvelope(userId = userId))
         assertNull(authDiskSource.getEphemeralPinProtectedUserKeyEnvelope(userId = userId))
         assertNull(authDiskSource.getPersistentPinProtectedUserKeyEnvelope(userId = userId))
+        assertNull(authDiskSource.getV2UpgradeToken(userId = userId))
     }
 
     @Test
@@ -492,6 +500,46 @@ class AuthDiskSourceTest {
         )
     }
 
+    @Test
+    fun `getV2UpgradeToken should pull from SharedPreferences`() {
+        val v2UpgradeTokenBaseKey = "bwPreferencesStorage:v2UpgradeToken"
+        val mockUserId = "mockUserId"
+        val mockV2UpgradeToken = V2UpgradeTokenJson(
+            wrappedUserKey1 = "wrappedUserKey1",
+            wrappedUserKey2 = "wrappedUserKey2",
+        )
+        fakeSharedPreferences.edit {
+            putString(
+                "${v2UpgradeTokenBaseKey}_$mockUserId",
+                json.encodeToString(mockV2UpgradeToken),
+            )
+        }
+        val actual = authDiskSource.getV2UpgradeToken(userId = mockUserId)
+        assertEquals(mockV2UpgradeToken, actual)
+    }
+
+    @Test
+    fun `storeV2UpgradeToken should update SharedPreferences`() {
+        val v2UpgradeTokenBaseKey = "bwPreferencesStorage:v2UpgradeToken"
+        val mockUserId = "mockUserId"
+        val mockV2UpgradeToken = V2UpgradeTokenJson(
+            wrappedUserKey1 = "wrappedUserKey1",
+            wrappedUserKey2 = "wrappedUserKey2",
+        )
+        authDiskSource.storeV2UpgradeToken(
+            userId = mockUserId,
+            v2UpgradeToken = mockV2UpgradeToken,
+        )
+        val actual = fakeSharedPreferences.getString(
+            "${v2UpgradeTokenBaseKey}_$mockUserId",
+            null,
+        )
+        assertEquals(
+            json.encodeToJsonElement(mockV2UpgradeToken),
+            json.parseToJsonElement(requireNotNull(actual)),
+        )
+    }
+
     @Test
     fun `getTwoFactorToken should pull from SharedPreferences`() {
         val twoFactorTokenBaseKey = "bwPreferencesStorage:twoFactorToken"

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/datasource/disk/util/FakeAuthDiskSource.kt

@@ -3,6 +3,7 @@ package com.x8bit.bitwarden.data.auth.datasource.disk.util
 import com.bitwarden.core.WrappedAccountCryptographicState
 import com.bitwarden.core.data.repository.util.bufferedMutableSharedFlow
 import com.bitwarden.network.model.SyncResponseJson
+import com.bitwarden.network.model.V2UpgradeTokenJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountTokensJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.OnboardingStatus
@@ -71,6 +72,7 @@ class FakeAuthDiskSource : AuthDiskSource {
         mutableMapOf<String, MutableSharedFlow<String?>>()
     private val mutablePersistentPinProtectedUserKeyEnvelopesFlowMap =
         mutableMapOf<String, MutableSharedFlow<String?>>()
+    private val storedV2UpgradeTokens = mutableMapOf<String, V2UpgradeTokenJson?>()
 
     override var userState: UserStateJson? = null
         set(value) {
@@ -97,6 +99,7 @@ class FakeAuthDiskSource : AuthDiskSource {
         storedPersistentPinProtectedUserKeyEnvelopes.remove(userId)
         storedEncryptedPins.remove(userId)
         storedPinProtectedUserKeys.remove(userId)
+        storedV2UpgradeTokens.remove(userId)
 
         mutableShouldUseKeyConnectorFlowMap.remove(userId)
         mutableOrganizationsFlowMap.remove(userId)
@@ -160,6 +163,13 @@ class FakeAuthDiskSource : AuthDiskSource {
         storedAccountCryptographicState[userId] = accountCryptographicState
     }
 
+    override fun getV2UpgradeToken(userId: String): V2UpgradeTokenJson? =
+        storedV2UpgradeTokens[userId]
+
+    override fun storeV2UpgradeToken(userId: String, v2UpgradeToken: V2UpgradeTokenJson?) {
+        storedV2UpgradeTokens[userId] = v2UpgradeToken
+    }
+
     override fun getTwoFactorToken(email: String): String? = storedTwoFactorTokens[email]
 
     override fun storeTwoFactorToken(email: String, twoFactorToken: String?) {
@@ -439,6 +449,13 @@ class FakeAuthDiskSource : AuthDiskSource {
         assertEquals(accountCryptographicState, storedAccountCryptographicState[userId])
     }
 
+    /**
+     * Assert that the [v2UpgradeToken] was stored successfully using the [userId].
+     */
+    fun assertV2UpgradeToken(userId: String, v2UpgradeToken: V2UpgradeTokenJson?) {
+        assertEquals(v2UpgradeToken, storedV2UpgradeTokens[userId])
+    }
+
     /**
      * Assert that the [twoFactorToken] was stored successfully using the [email].
      */

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/repository/util/UserStateJsonExtensionsTest.kt

@@ -2037,6 +2037,7 @@ class UserStateJsonExtensionsTest {
             ),
             userDecryption = UserDecryptionJson(
                 masterPasswordUnlock = MOCK_MASTER_PASSWORD_UNLOCK_DATA,
+                v2UpgradeToken = null,
             ),
         )
 
@@ -2128,6 +2129,7 @@ class UserStateJsonExtensionsTest {
             ),
             userDecryption = UserDecryptionJson(
                 masterPasswordUnlock = MOCK_MASTER_PASSWORD_UNLOCK_DATA,
+                v2UpgradeToken = null,
             ),
         )
 
@@ -2292,6 +2294,7 @@ class UserStateJsonExtensionsTest {
                     ),
                     masterKeyWrappedUserKey = "mockMasterKeyWrappedUserKey",
                 ),
+                v2UpgradeToken = null,
             ),
         )
 
@@ -2379,7 +2382,10 @@ class UserStateJsonExtensionsTest {
                         ),
                     ),
                 ),
-                userDecryption = UserDecryptionJson(masterPasswordUnlock = null),
+                userDecryption = UserDecryptionJson(
+                    masterPasswordUnlock = null,
+                    v2UpgradeToken = null,
+                ),
             ),
         )
 
@@ -2447,6 +2453,7 @@ class UserStateJsonExtensionsTest {
                 ),
                 userDecryption = UserDecryptionJson(
                     masterPasswordUnlock = MOCK_MASTER_PASSWORD_UNLOCK_DATA,
+                    v2UpgradeToken = null,
                 ),
             ),
         )

app/src/test/kotlin/com/x8bit/bitwarden/data/platform/manager/sdk/SdkRepositoryFactoryTests.kt

@@ -53,6 +53,21 @@ class SdkRepositoryFactoryTests {
         serviceClientConfig = serviceClientConfig,
     )
 
+    @Test
+    fun `getStateBridge should create a new bridge`() {
+        val userId = "userId"
+        val firstBridge = sdkRepoFactory.getStateBridge(userId = userId)
+
+        // Additional calls for the same userId should create a new bridge
+        val secondBridge = sdkRepoFactory.getStateBridge(userId = userId)
+        assertNotEquals(firstBridge, secondBridge)
+
+        // Additional calls for different userIds should return a different bridge
+        val otherUserId = "otherUserId"
+        val thirdBridge = sdkRepoFactory.getStateBridge(userId = otherUserId)
+        assertNotEquals(firstBridge, thirdBridge)
+    }
+
     @Test
     fun `getRepositories should create a new client`() {
         val userId = "userId"

app/src/test/kotlin/com/x8bit/bitwarden/data/platform/manager/sdk/statebridge/SdkStateBridgeTest.kt

@@ -0,0 +1,368 @@
+package com.x8bit.bitwarden.data.platform.manager.sdk.statebridge
+
+import com.bitwarden.core.MasterPasswordUnlockData
+import com.bitwarden.core.V2UpgradeToken
+import com.bitwarden.crypto.Kdf
+import com.bitwarden.network.model.KdfJson
+import com.bitwarden.network.model.KdfTypeJson
+import com.bitwarden.network.model.MasterPasswordUnlockDataJson
+import com.bitwarden.network.model.UserDecryptionOptionsJson
+import com.bitwarden.network.model.V2UpgradeTokenJson
+import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountJson
+import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
+import com.x8bit.bitwarden.data.auth.datasource.disk.util.FakeAuthDiskSource
+import com.x8bit.bitwarden.data.auth.repository.model.createMockWrappedAccountCryptographicState
+import com.x8bit.bitwarden.data.auth.repository.util.updateMasterPasswordUnlock
+import com.x8bit.bitwarden.data.vault.repository.util.toSdkMasterPasswordUnlock
+import io.mockk.mockk
+import kotlinx.coroutines.test.runTest
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertNull
+import org.junit.jupiter.api.Test
+import java.time.Instant
+
+class SdkStateBridgeTest {
+    private val authDiskSource = FakeAuthDiskSource()
+
+    private val stateBridge = SdkStateBridge(
+        userId = USER_ID,
+        authDiskSource = authDiskSource,
+    )
+
+    @Test
+    fun `setUserKey should store the user key in memory`() = runTest {
+        stateBridge.setUserKey(value = "userKey")
+
+        assertEquals("userKey", stateBridge.getUserKey())
+    }
+
+    @Test
+    fun `getUserKey should return the in-memory user key`() = runTest {
+        assertNull(stateBridge.getUserKey())
+
+        stateBridge.setUserKey(value = "userKey")
+
+        assertEquals("userKey", stateBridge.getUserKey())
+    }
+
+    @Test
+    fun `clearUserKey should clear the in-memory user key`() = runTest {
+        stateBridge.setUserKey(value = "userKey")
+
+        stateBridge.clearUserKey()
+
+        assertNull(stateBridge.getUserKey())
+    }
+
+    @Test
+    fun `setPersistentPinEnvelope should store the persistent pin envelope`() = runTest {
+        stateBridge.setPersistentPinEnvelope(value = "pinEnvelope")
+
+        authDiskSource.assertPersistentPinProtectedUserKeyEnvelope(
+            userId = USER_ID,
+            pinProtectedUserKeyEnvelope = "pinEnvelope",
+        )
+    }
+
+    @Test
+    fun `getPersistentPinEnvelope should return the stored pin envelope`() = runTest {
+        assertNull(stateBridge.getPersistentPinEnvelope())
+
+        authDiskSource.storePersistentPinProtectedUserKeyEnvelope(
+            userId = USER_ID,
+            pinProtectedUserKeyEnvelope = "pinEnvelope",
+        )
+
+        assertEquals("pinEnvelope", stateBridge.getPersistentPinEnvelope())
+    }
+
+    @Test
+    fun `clearPersistentPinEnvelope should clear the persistent pin envelope`() = runTest {
+        authDiskSource.storePersistentPinProtectedUserKeyEnvelope(
+            userId = USER_ID,
+            pinProtectedUserKeyEnvelope = "pinEnvelope",
+        )
+
+        stateBridge.clearPersistentPinEnvelope()
+
+        authDiskSource.assertPersistentPinProtectedUserKeyEnvelope(
+            userId = USER_ID,
+            pinProtectedUserKeyEnvelope = null,
+        )
+    }
+
+    @Test
+    fun `setEphemeralPinEnvelope should store the ephemeral pin envelope`() = runTest {
+        stateBridge.setEphemeralPinEnvelope(value = "pinEnvelope")
+
+        authDiskSource.assertEphemeralPinProtectedUserKeyEnvelope(
+            userId = USER_ID,
+            pinProtectedUserKeyEnvelope = "pinEnvelope",
+        )
+    }
+
+    @Test
+    fun `getEphemeralPinEnvelope should return the stored pin envelope`() = runTest {
+        assertNull(stateBridge.getEphemeralPinEnvelope())
+
+        authDiskSource.storeEphemeralPinProtectedUserKeyEnvelope(
+            userId = USER_ID,
+            pinProtectedUserKeyEnvelope = "pinEnvelope",
+        )
+
+        assertEquals("pinEnvelope", stateBridge.getEphemeralPinEnvelope())
+    }
+
+    @Test
+    fun `clearEphemeralPinEnvelope should clear the ephemeral pin envelope`() = runTest {
+        authDiskSource.storeEphemeralPinProtectedUserKeyEnvelope(
+            userId = USER_ID,
+            pinProtectedUserKeyEnvelope = "pinEnvelope",
+        )
+
+        stateBridge.clearEphemeralPinEnvelope()
+
+        authDiskSource.assertEphemeralPinProtectedUserKeyEnvelope(
+            userId = USER_ID,
+            pinProtectedUserKeyEnvelope = null,
+        )
+    }
+
+    @Test
+    fun `setEncryptedPin should store the encrypted pin`() = runTest {
+        stateBridge.setEncryptedPin(value = "encryptedPin")
+
+        authDiskSource.assertEncryptedPin(userId = USER_ID, encryptedPin = "encryptedPin")
+    }
+
+    @Test
+    fun `getEncryptedPin should return the stored encrypted pin`() = runTest {
+        assertNull(stateBridge.getEncryptedPin())
+
+        authDiskSource.storeEncryptedPin(userId = USER_ID, encryptedPin = "encryptedPin")
+
+        assertEquals("encryptedPin", stateBridge.getEncryptedPin())
+    }
+
+    @Test
+    fun `clearEncryptedPin should clear the encrypted pin`() = runTest {
+        authDiskSource.storeEncryptedPin(userId = USER_ID, encryptedPin = "encryptedPin")
+
+        stateBridge.clearEncryptedPin()
+
+        authDiskSource.assertEncryptedPin(userId = USER_ID, encryptedPin = null)
+    }
+
+    @Test
+    fun `setV2UpgradeToken should store the token as a V2UpgradeTokenJson`() = runTest {
+        stateBridge.setV2UpgradeToken(value = V2_UPGRADE_TOKEN)
+
+        authDiskSource.assertV2UpgradeToken(
+            userId = USER_ID,
+            v2UpgradeToken = V2_UPGRADE_TOKEN_JSON,
+        )
+    }
+
+    @Test
+    fun `getV2UpgradeToken should return the stored token mapped to a V2UpgradeToken`() = runTest {
+        assertNull(stateBridge.getV2UpgradeToken())
+
+        authDiskSource.storeV2UpgradeToken(
+            userId = USER_ID,
+            v2UpgradeToken = V2_UPGRADE_TOKEN_JSON,
+        )
+
+        assertEquals(V2_UPGRADE_TOKEN, stateBridge.getV2UpgradeToken())
+    }
+
+    @Test
+    fun `clearV2UpgradeToken should clear the stored token`() = runTest {
+        authDiskSource.storeV2UpgradeToken(
+            userId = USER_ID,
+            v2UpgradeToken = V2_UPGRADE_TOKEN_JSON,
+        )
+
+        stateBridge.clearV2UpgradeToken()
+
+        authDiskSource.assertV2UpgradeToken(userId = USER_ID, v2UpgradeToken = null)
+    }
+
+    @Test
+    fun `setAccountCryptographicState should store the account cryptographic state`() = runTest {
+        val state = createMockWrappedAccountCryptographicState(number = 1)
+
+        stateBridge.setAccountCryptographicState(value = state)
+
+        authDiskSource.assertAccountCryptographicState(
+            userId = USER_ID,
+            accountCryptographicState = state,
+        )
+    }
+
+    @Test
+    fun `getAccountCryptographicState should return the stored account cryptographic state`() =
+        runTest {
+            assertNull(stateBridge.getAccountCryptographicState())
+
+            val state = createMockWrappedAccountCryptographicState(number = 1)
+            authDiskSource.storeAccountCryptographicState(
+                userId = USER_ID,
+                accountCryptographicState = state,
+            )
+
+            assertEquals(state, stateBridge.getAccountCryptographicState())
+        }
+
+    @Test
+    fun `clearAccountCryptographicState should clear the account cryptographic state`() = runTest {
+        authDiskSource.storeAccountCryptographicState(
+            userId = USER_ID,
+            accountCryptographicState = createMockWrappedAccountCryptographicState(number = 1),
+        )
+
+        stateBridge.clearAccountCryptographicState()
+
+        authDiskSource.assertAccountCryptographicState(
+            userId = USER_ID,
+            accountCryptographicState = null,
+        )
+    }
+
+    @Test
+    fun `setMasterpasswordUnlockData should update the user state with the unlock data`() =
+        runTest {
+            authDiskSource.userState = USER_STATE
+
+            stateBridge.setMasterpasswordUnlockData(value = MASTER_PASSWORD_UNLOCK_DATA)
+
+            assertEquals(
+                USER_STATE.updateMasterPasswordUnlock(
+                    userId = USER_ID,
+                    masterPasswordUnlock = MASTER_PASSWORD_UNLOCK_DATA,
+                ),
+                authDiskSource.userState,
+            )
+        }
+
+    @Test
+    fun `setMasterpasswordUnlockData should do nothing when the user state is null`() = runTest {
+        authDiskSource.userState = null
+
+        stateBridge.setMasterpasswordUnlockData(value = MASTER_PASSWORD_UNLOCK_DATA)
+
+        assertNull(authDiskSource.userState)
+    }
+
+    @Test
+    fun `getMasterpasswordUnlockData should return null when there is no unlock data`() = runTest {
+        authDiskSource.userState = null
+
+        assertNull(stateBridge.getMasterpasswordUnlockData())
+    }
+
+    @Test
+    fun `getMasterpasswordUnlockData should return the stored unlock data as the sdk model`() =
+        runTest {
+            authDiskSource.userState = USER_STATE.copy(
+                accounts = mapOf(
+                    USER_ID to ACCOUNT.copy(
+                        profile = ACCOUNT.profile.copy(
+                            userDecryptionOptions = UserDecryptionOptionsJson(
+                                hasMasterPassword = true,
+                                trustedDeviceUserDecryptionOptions = null,
+                                keyConnectorUserDecryptionOptions = null,
+                                masterPasswordUnlock = MASTER_PASSWORD_UNLOCK_DATA_JSON,
+                            ),
+                        ),
+                    ),
+                ),
+            )
+
+            assertEquals(
+                MASTER_PASSWORD_UNLOCK_DATA_JSON.toSdkMasterPasswordUnlock(),
+                stateBridge.getMasterpasswordUnlockData(),
+            )
+        }
+
+    @Test
+    fun `clearMasterpasswordUnlockData should clear the unlock data from the user state`() =
+        runTest {
+            authDiskSource.userState = USER_STATE.copy(
+                accounts = mapOf(
+                    USER_ID to ACCOUNT.copy(
+                        profile = ACCOUNT.profile.copy(
+                            userDecryptionOptions = UserDecryptionOptionsJson(
+                                hasMasterPassword = true,
+                                trustedDeviceUserDecryptionOptions = null,
+                                keyConnectorUserDecryptionOptions = null,
+                                masterPasswordUnlock = MASTER_PASSWORD_UNLOCK_DATA_JSON,
+                            ),
+                        ),
+                    ),
+                ),
+            )
+
+            stateBridge.clearMasterpasswordUnlockData()
+
+            assertNull(stateBridge.getMasterpasswordUnlockData())
+        }
+}
+
+private const val USER_ID: String = "userId"
+
+private val V2_UPGRADE_TOKEN: V2UpgradeToken = V2UpgradeToken(
+    wrappedUserKey1 = "wrappedUserKey1",
+    wrappedUserKey2 = "wrappedUserKey2",
+)
+
+private val V2_UPGRADE_TOKEN_JSON: V2UpgradeTokenJson = V2UpgradeTokenJson(
+    wrappedUserKey1 = "wrappedUserKey1",
+    wrappedUserKey2 = "wrappedUserKey2",
+)
+
+private val MASTER_PASSWORD_UNLOCK_DATA: MasterPasswordUnlockData = MasterPasswordUnlockData(
+    kdf = Kdf.Pbkdf2(iterations = 600_000u),
+    masterKeyWrappedUserKey = "masterKeyWrappedUserKey",
+    salt = "salt",
+)
+
+private val MASTER_PASSWORD_UNLOCK_DATA_JSON: MasterPasswordUnlockDataJson =
+    MasterPasswordUnlockDataJson(
+        kdf = KdfJson(
+            kdfType = KdfTypeJson.PBKDF2_SHA256,
+            iterations = 600_000,
+            memory = null,
+            parallelism = null,
+        ),
+        masterKeyWrappedUserKey = "masterKeyWrappedUserKey",
+        salt = "salt",
+    )
+
+private val ACCOUNT: AccountJson = AccountJson(
+    profile = AccountJson.Profile(
+        userId = USER_ID,
+        email = "email@bitwarden.com",
+        isEmailVerified = true,
+        name = "name",
+        stamp = null,
+        organizationId = null,
+        avatarColorHex = null,
+        hasPremiumPersonally = null,
+        hasPremiumFromOrganization = null,
+        forcePasswordResetReason = null,
+        kdfType = KdfTypeJson.PBKDF2_SHA256,
+        kdfIterations = 600_000,
+        kdfMemory = null,
+        kdfParallelism = null,
+        userDecryptionOptions = null,
+        isTwoFactorEnabled = false,
+        creationDate = Instant.parse("2024-09-13T01:00:00.00Z"),
+    ),
+    tokens = mockk(),
+    settings = mockk(),
+)
+
+private val USER_STATE: UserStateJson = UserStateJson(
+    activeUserId = USER_ID,
+    accounts = mapOf(USER_ID to ACCOUNT),
+)

app/src/test/kotlin/com/x8bit/bitwarden/data/vault/manager/VaultLockManagerTest.kt

@@ -9,6 +9,7 @@ import com.bitwarden.core.InitOrgCryptoRequest
 import com.bitwarden.core.InitUserCryptoMethod
 import com.bitwarden.core.InitUserCryptoRequest
 import com.bitwarden.core.MasterPasswordUnlockData
+import com.bitwarden.core.V2UpgradeToken
 import com.bitwarden.core.WrappedAccountCryptographicState
 import com.bitwarden.core.data.manager.dispatcher.FakeDispatcherManager
 import com.bitwarden.core.data.manager.realtime.RealtimeManager
@@ -38,6 +39,7 @@ import com.x8bit.bitwarden.data.vault.datasource.sdk.model.InitializeCryptoResul
 import com.x8bit.bitwarden.data.vault.manager.model.VaultStateEvent
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockData
 import com.x8bit.bitwarden.data.vault.repository.model.VaultUnlockResult
+import com.x8bit.bitwarden.data.vault.repository.util.toV2UpgradeTokenJson
 import io.mockk.clearMocks
 import io.mockk.coEvery
 import io.mockk.coVerify
@@ -1001,6 +1003,73 @@ class VaultLockManagerTest {
             }
         }
 
+    @Suppress("MaxLineLength")
+    @Test
+    fun `unlockVault with a stored V2 upgrade token should pass the mapped token to initializeCrypto`() =
+        runTest {
+            val kdf = MOCK_PROFILE.toSdkParams()
+            val email = MOCK_PROFILE.email
+            val masterPassword = "mockValue"
+            val v2UpgradeToken = V2UpgradeToken(
+                wrappedUserKey1 = "wrappedUserKey1",
+                wrappedUserKey2 = "wrappedUserKey2",
+            )
+            fakeAuthDiskSource.storeV2UpgradeToken(
+                userId = USER_ID,
+                v2UpgradeToken = v2UpgradeToken.toV2UpgradeTokenJson(),
+            )
+            coEvery {
+                vaultSdkSource.initializeCrypto(
+                    userId = USER_ID,
+                    request = InitUserCryptoRequest(
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
+                        userId = USER_ID,
+                        kdfParams = kdf,
+                        email = email,
+                        method = InitUserCryptoMethod.MasterPasswordUnlock(
+                            password = masterPassword,
+                            masterPasswordUnlock = MOCK_MASTER_PASSWORD_UNLOCK_DATA,
+                        ),
+                        upgradeToken = v2UpgradeToken,
+                    ),
+                )
+            } returns InitializeCryptoResult.Success.asSuccess()
+            coEvery {
+                trustedDeviceManager.trustThisDeviceIfNecessary(userId = USER_ID)
+            } returns false.asSuccess()
+            mutableVaultTimeoutStateFlow.value = VaultTimeout.ThirtyMinutes
+
+            val result = vaultLockManager.unlockVault(
+                accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
+                userId = USER_ID,
+                email = email,
+                kdf = kdf,
+                initUserCryptoMethod = InitUserCryptoMethod.MasterPasswordUnlock(
+                    password = masterPassword,
+                    masterPasswordUnlock = MOCK_MASTER_PASSWORD_UNLOCK_DATA,
+                ),
+                organizationKeys = null,
+            )
+
+            assertEquals(VaultUnlockResult.Success, result)
+            coVerify(exactly = 1) {
+                vaultSdkSource.initializeCrypto(
+                    userId = USER_ID,
+                    request = InitUserCryptoRequest(
+                        accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE,
+                        userId = USER_ID,
+                        kdfParams = kdf,
+                        email = email,
+                        method = InitUserCryptoMethod.MasterPasswordUnlock(
+                            password = masterPassword,
+                            masterPasswordUnlock = MOCK_MASTER_PASSWORD_UNLOCK_DATA,
+                        ),
+                        upgradeToken = v2UpgradeToken,
+                    ),
+                )
+            }
+        }
+
     @Suppress("MaxLineLength")
     @Test
     fun `unlockVault with initializeCrypto success for a Never VaultTimeout should return Success, save the auto-unlock key, and clear invalid unlock attempts`() =

app/src/test/kotlin/com/x8bit/bitwarden/data/vault/manager/VaultSyncManagerTest.kt

@@ -24,6 +24,7 @@ import com.bitwarden.network.model.createMockPolicy
 import com.bitwarden.network.model.createMockProfile
 import com.bitwarden.network.model.createMockSend
 import com.bitwarden.network.model.createMockSyncResponse
+import com.bitwarden.network.model.createMockV2UpgradeToken
 import com.bitwarden.network.service.SyncService
 import com.bitwarden.send.SendView
 import com.bitwarden.vault.DecryptCipherListResult
@@ -34,9 +35,9 @@ import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.util.FakeAuthDiskSource
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
 import com.x8bit.bitwarden.data.auth.manager.UserStateManager
-import com.x8bit.bitwarden.data.autofill.manager.FillAssistManager
 import com.x8bit.bitwarden.data.auth.repository.model.LogoutReason
 import com.x8bit.bitwarden.data.auth.repository.model.createMockWrappedAccountCryptographicState
+import com.x8bit.bitwarden.data.autofill.manager.FillAssistManager
 import com.x8bit.bitwarden.data.platform.datasource.disk.SettingsDiskSource
 import com.x8bit.bitwarden.data.platform.error.NoActiveUserException
 import com.x8bit.bitwarden.data.platform.manager.DatabaseSchemeManager
@@ -758,6 +759,10 @@ class VaultSyncManagerTest {
                 userId = userId,
                 accountCryptographicState = createMockWrappedAccountCryptographicState(number = 1),
             )
+            fakeAuthDiskSource.assertV2UpgradeToken(
+                userId = userId,
+                v2UpgradeToken = createMockV2UpgradeToken(number = 1),
+            )
             fakeAuthDiskSource.assertOrganizationKeys(
                 userId = userId,
                 organizationKeys = mapOf(userId to "mockKey-1"),

app/src/test/kotlin/com/x8bit/bitwarden/data/vault/repository/util/V2UpgradeTokenExtensionsTest.kt

@@ -0,0 +1,36 @@
+package com.x8bit.bitwarden.data.vault.repository.util
+
+import com.bitwarden.core.V2UpgradeToken
+import com.bitwarden.network.model.V2UpgradeTokenJson
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+
+class V2UpgradeTokenExtensionsTest {
+    @Test
+    fun `toV2UpgradeTokenJson maps all fields to a V2UpgradeTokenJson`() {
+        val token = V2UpgradeToken(
+            wrappedUserKey1 = "wrappedUserKey1",
+            wrappedUserKey2 = "wrappedUserKey2",
+        )
+        val expected = V2UpgradeTokenJson(
+            wrappedUserKey1 = "wrappedUserKey1",
+            wrappedUserKey2 = "wrappedUserKey2",
+        )
+
+        assertEquals(expected, token.toV2UpgradeTokenJson())
+    }
+
+    @Test
+    fun `toV2UpgradeToken maps all fields to a V2UpgradeToken`() {
+        val json = V2UpgradeTokenJson(
+            wrappedUserKey1 = "wrappedUserKey1",
+            wrappedUserKey2 = "wrappedUserKey2",
+        )
+        val expected = V2UpgradeToken(
+            wrappedUserKey1 = "wrappedUserKey1",
+            wrappedUserKey2 = "wrappedUserKey2",
+        )
+
+        assertEquals(expected, json.toV2UpgradeToken())
+    }
+}

network/src/main/kotlin/com/bitwarden/network/model/UserDecryptionJson.kt

@@ -5,9 +5,17 @@ import kotlinx.serialization.Serializable
 
 /**
  * Represents the user decryption options received on sync.
+ *
+ * @property masterPasswordUnlock The unlock data when the user has a master password that can be
+ * used to decrypt their vault.
+ * @property v2UpgradeToken The V2 upgrade token returned when available, allowing vault unlock
+ * after V1 → V2 upgrade.
  */
 @Serializable
 data class UserDecryptionJson(
     @SerialName("masterPasswordUnlock")
     val masterPasswordUnlock: MasterPasswordUnlockDataJson?,
+
+    @SerialName("v2UpgradeToken")
+    val v2UpgradeToken: V2UpgradeTokenJson?,
 )

network/src/main/kotlin/com/bitwarden/network/model/V2UpgradeTokenJson.kt

@@ -0,0 +1,16 @@
+package com.bitwarden.network.model
+
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+
+/**
+ * Represents the V2 upgrade token, allowing vault unlock after V1 → V2 upgrade.
+ */
+@Serializable
+data class V2UpgradeTokenJson(
+    @SerialName("wrappedUserKey1")
+    val wrappedUserKey1: String,
+
+    @SerialName("wrappedUserKey2")
+    val wrappedUserKey2: String,
+)

network/src/test/kotlin/com/bitwarden/network/service/SyncServiceTest.kt

@@ -597,6 +597,10 @@ private const val SYNC_SUCCESS_JSON = """
       },
       "masterKeyWrappedUserKey": "mockMasterKeyWrappedUserKey-1",
       "salt": "mockSalt-1"
+    },
+    "v2UpgradeToken": {
+      "wrappedUserKey1": "mockWrappedUserKey1-1",
+      "wrappedUserKey2": "mockWrappedUserKey2-1"
     }
   }
 }

network/src/testFixtures/kotlin/com/bitwarden/network/model/SyncResponseUtil.kt

@@ -35,9 +35,24 @@ fun createMockUserDecryption(
     masterPasswordUnlock: MasterPasswordUnlockDataJson? = createMockMasterPasswordUnlock(
         number = number,
     ),
+    v2UpgradeToken: V2UpgradeTokenJson? = createMockV2UpgradeToken(number = number),
 ): UserDecryptionJson =
     UserDecryptionJson(
         masterPasswordUnlock = masterPasswordUnlock,
+        v2UpgradeToken = v2UpgradeToken,
+    )
+
+/**
+ * Create a mock [V2UpgradeTokenJson] with a given [number].
+ */
+fun createMockV2UpgradeToken(
+    number: Int,
+    wrappedUserKey1: String = "mockWrappedUserKey1-$number",
+    wrappedUserKey2: String = "mockWrappedUserKey2-$number",
+): V2UpgradeTokenJson =
+    V2UpgradeTokenJson(
+        wrappedUserKey1 = wrappedUserKey1,
+        wrappedUserKey2 = wrappedUserKey2,
     )
 
 /**