Chore: Move dispatcher for sdk functions into SDK sources (#6995)

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/datasource/sdk/AuthSdkSourceImpl.kt

@@ -13,6 +13,7 @@ import com.bitwarden.core.KeyConnectorResponse
 import com.bitwarden.core.MasterPasswordPolicyOptions
 import com.bitwarden.core.RegisterKeyResponse
 import com.bitwarden.core.RegisterTdeKeyResponse
+import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
 import com.bitwarden.crypto.HashPurpose
 import com.bitwarden.crypto.Kdf
 import com.bitwarden.policies.OrganizationUserPolicyContext
@@ -24,6 +25,7 @@ import com.x8bit.bitwarden.data.auth.datasource.sdk.util.toPasswordStrengthOrNul
 import com.x8bit.bitwarden.data.auth.datasource.sdk.util.toUByte
 import com.x8bit.bitwarden.data.platform.datasource.sdk.BaseSdkSource
 import com.x8bit.bitwarden.data.platform.manager.SdkClientManager
+import kotlinx.coroutines.withContext
 
 /**
  * Primary implementation of [AuthSdkSource] that serves as a convenience wrapper around a
@@ -31,6 +33,7 @@ import com.x8bit.bitwarden.data.platform.manager.SdkClientManager
  */
 @Suppress("TooManyFunctions")
 class AuthSdkSourceImpl(
+    private val dispatcherManager: DispatcherManager,
     sdkClientManager: SdkClientManager,
 ) : BaseSdkSource(sdkClientManager = sdkClientManager),
     AuthSdkSource {
@@ -45,10 +48,8 @@ class AuthSdkSourceImpl(
         masterPasswordHint: String?,
         shouldResetPasswordEnroll: Boolean,
     ): Result<JitMasterPasswordRegistrationResponse> = runCatchingWithLogs {
-        getClient(userId = userId)
-            .auth()
-            .registration()
-            .postKeysForJitPasswordRegistration(
+        withContext(context = dispatcherManager.io) {
+            getClient(userId = userId).auth().registration().postKeysForJitPasswordRegistration(
                 request = JitMasterPasswordRegistrationRequest(
                     orgId = organizationId,
                     orgPublicKey = organizationPublicKey,
@@ -60,6 +61,7 @@ class AuthSdkSourceImpl(
                     resetPasswordEnroll = shouldResetPasswordEnroll,
                 ),
             )
+        }
     }
 
     override suspend fun postKeysForKeyConnectorRegistration(
@@ -68,11 +70,13 @@ class AuthSdkSourceImpl(
         keyConnectorUrl: String,
         ssoOrganizationIdentifier: String,
     ): Result<KeyConnectorRegistrationResult> = runCatchingWithLogs {
-        useClient(userId = userId, accessToken = accessToken) {
-            auth().registration().postKeysForKeyConnectorRegistration(
-                keyConnectorUrl = keyConnectorUrl,
-                ssoOrgIdentifier = ssoOrganizationIdentifier,
-            )
+        withContext(context = dispatcherManager.io) {
+            useClient(userId = userId, accessToken = accessToken) {
+                auth().registration().postKeysForKeyConnectorRegistration(
+                    keyConnectorUrl = keyConnectorUrl,
+                    ssoOrgIdentifier = ssoOrganizationIdentifier,
+                )
+            }
         }
     }
 
@@ -83,10 +87,8 @@ class AuthSdkSourceImpl(
         deviceIdentifier: String,
         shouldTrustDevice: Boolean,
     ): Result<TdeRegistrationResponse> = runCatchingWithLogs {
-        getClient(userId = userId)
-            .auth()
-            .registration()
-            .postKeysForTdeRegistration(
+        withContext(context = dispatcherManager.io) {
+            getClient(userId = userId).auth().registration().postKeysForTdeRegistration(
                 request = TdeRegistrationRequest(
                     orgId = organizationId,
                     orgPublicKey = organizationPublicKey,
@@ -95,6 +97,7 @@ class AuthSdkSourceImpl(
                     trustDevice = shouldTrustDevice,
                 ),
             )
+        }
     }
 
     override suspend fun postKeysForUserPasswordRegistration(
@@ -104,23 +107,25 @@ class AuthSdkSourceImpl(
         masterPasswordHint: String?,
         emailVerificationToken: String,
     ): Result<UserMasterPasswordRegistrationResponse> = runCatchingWithLogs {
-        useClient {
-            auth().registration().postKeysForUserPasswordRegistration(
-                request = UserMasterPasswordRegistrationRequest(
-                    email = email,
-                    salt = salt,
-                    masterPassword = masterPassword,
-                    masterPasswordHint = masterPasswordHint,
-                    emailVerificationToken = emailVerificationToken,
-                    organizationUserId = null,
-                    orgInviteToken = null,
-                    orgSponsoredFreeFamilyPlanToken = null,
-                    acceptEmergencyAccessInviteToken = null,
-                    acceptEmergencyAccessId = null,
-                    providerInviteToken = null,
-                    providerUserId = null,
-                ),
-            )
+        withContext(context = dispatcherManager.io) {
+            useClient {
+                auth().registration().postKeysForUserPasswordRegistration(
+                    request = UserMasterPasswordRegistrationRequest(
+                        email = email,
+                        salt = salt,
+                        masterPassword = masterPassword,
+                        masterPasswordHint = masterPasswordHint,
+                        emailVerificationToken = emailVerificationToken,
+                        organizationUserId = null,
+                        orgInviteToken = null,
+                        orgSponsoredFreeFamilyPlanToken = null,
+                        acceptEmergencyAccessInviteToken = null,
+                        acceptEmergencyAccessId = null,
+                        providerInviteToken = null,
+                        providerUserId = null,
+                    ),
+                )
+            }
         }
     }
 

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/datasource/sdk/di/AuthSdkModule.kt

@@ -1,5 +1,6 @@
 package com.x8bit.bitwarden.data.auth.datasource.sdk.di
 
+import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
 import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
 import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSourceImpl
 import com.x8bit.bitwarden.data.platform.manager.SdkClientManager
@@ -19,8 +20,10 @@ object AuthSdkModule {
     @Provides
     @Singleton
     fun provideAuthSdkSource(
+        dispatcherManager: DispatcherManager,
         sdkClientManager: SdkClientManager,
     ): AuthSdkSource = AuthSdkSourceImpl(
+        dispatcherManager = dispatcherManager,
         sdkClientManager = sdkClientManager,
     )
 }

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/manager/KeyConnectorManagerImpl.kt

@@ -1,7 +1,6 @@
 package com.x8bit.bitwarden.data.auth.manager
 
 import com.bitwarden.core.WrappedAccountCryptographicState
-import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
 import com.bitwarden.core.data.manager.model.FlagKey
 import com.bitwarden.core.data.util.flatMap
 import com.bitwarden.crypto.Kdf
@@ -17,7 +16,6 @@ import com.x8bit.bitwarden.data.auth.repository.util.toAccountCryptographicState
 import com.x8bit.bitwarden.data.platform.manager.FeatureFlagManager
 import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
 import com.x8bit.bitwarden.data.vault.datasource.sdk.model.DeriveKeyConnectorResult
-import kotlinx.coroutines.withContext
 
 /**
  * The default implementation of the [KeyConnectorManager].
@@ -27,7 +25,6 @@ class KeyConnectorManagerImpl(
     private val authSdkSource: AuthSdkSource,
     private val vaultSdkSource: VaultSdkSource,
     private val featureFlagManager: FeatureFlagManager,
-    private val dispatcherManager: DispatcherManager,
 ) : KeyConnectorManager {
     override suspend fun getMasterKeyFromKeyConnector(
         url: String,
@@ -97,26 +94,24 @@ class KeyConnectorManagerImpl(
         organizationIdentifier: String,
     ): Result<MigrateNewUserToKeyConnectorResult> =
         if (featureFlagManager.getFeatureFlag(FlagKey.V2EncryptionKeyConnector)) {
-            withContext(dispatcherManager.io) {
-                authSdkSource
-                    .postKeysForKeyConnectorRegistration(
-                        userId = userId,
-                        accessToken = accessToken,
-                        keyConnectorUrl = url,
-                        ssoOrganizationIdentifier = organizationIdentifier,
+            authSdkSource
+                .postKeysForKeyConnectorRegistration(
+                    userId = userId,
+                    accessToken = accessToken,
+                    keyConnectorUrl = url,
+                    ssoOrganizationIdentifier = organizationIdentifier,
+                )
+                .map {
+                    MigrateNewUserToKeyConnectorResult(
+                        masterKey = it.keyConnectorKey,
+                        encryptedUserKey = it.keyConnectorKeyWrappedUserKey,
+                        privateKey = when (val state = it.accountCryptographicState) {
+                            is WrappedAccountCryptographicState.V1 -> state.privateKey
+                            is WrappedAccountCryptographicState.V2 -> state.privateKey
+                        },
+                        accountCryptographicState = it.accountCryptographicState,
                     )
-                    .map {
-                        MigrateNewUserToKeyConnectorResult(
-                            masterKey = it.keyConnectorKey,
-                            encryptedUserKey = it.keyConnectorKeyWrappedUserKey,
-                            privateKey = when (val state = it.accountCryptographicState) {
-                                is WrappedAccountCryptographicState.V1 -> state.privateKey
-                                is WrappedAccountCryptographicState.V2 -> state.privateKey
-                            },
-                            accountCryptographicState = it.accountCryptographicState,
-                        )
-                    }
-            }
+                }
         } else {
             legacyMigrateNewUserToKeyConnector(
                 accountKeys = accountKeys,

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/manager/di/AuthManagerModule.kt

@@ -90,14 +90,12 @@ object AuthManagerModule {
         authSdkSource: AuthSdkSource,
         vaultSdkSource: VaultSdkSource,
         featureFlagManager: FeatureFlagManager,
-        dispatcherManager: DispatcherManager,
     ): KeyConnectorManager =
         KeyConnectorManagerImpl(
             accountsService = accountsService,
             authSdkSource = authSdkSource,
             vaultSdkSource = vaultSdkSource,
             featureFlagManager = featureFlagManager,
-            dispatcherManager = dispatcherManager,
         )
 
     @Provides

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt

@@ -153,7 +153,6 @@ import kotlinx.coroutines.flow.merge
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.receiveAsFlow
 import kotlinx.coroutines.flow.stateIn
-import kotlinx.coroutines.withContext
 import timber.log.Timber
 import java.time.Clock
 import javax.inject.Singleton
@@ -190,7 +189,7 @@ class AuthRepositoryImpl(
     private val featureFlagManager: FeatureFlagManager,
     logsManager: LogsManager,
     pushManager: PushManager,
-    private val dispatcherManager: DispatcherManager,
+    dispatcherManager: DispatcherManager,
 ) : AuthRepository,
     AuthRequestManager by authRequestManager,
     BiometricsEncryptionManager by biometricsEncryptionManager,
@@ -565,15 +564,14 @@ class AuthRepositoryImpl(
     ): Result<VaultUnlockResult> {
         val userId = profile.userId
         val shouldTrustDevice = authDiskSource.getShouldTrustDevice(userId = userId) == true
-        return withContext(dispatcherManager.io) {
-            authSdkSource.postKeysForTdeRegistration(
+        return authSdkSource
+            .postKeysForTdeRegistration(
                 userId = userId,
                 organizationId = orgAutoEnrollStatus.organizationId,
                 organizationPublicKey = orgKeys.publicKey,
                 deviceIdentifier = authDiskSource.uniqueAppId,
                 shouldTrustDevice = shouldTrustDevice,
             )
-        }
             .map { response ->
                 // Clear the 'should trust device' flag, since the SDK trusted the device above.
                 authDiskSource.storeShouldTrustDevice(userId = userId, shouldTrustDevice = null)
@@ -976,15 +974,14 @@ class AuthRepositoryImpl(
             return RegisterResult.WeakPassword
         }
         if (featureFlagManager.getFeatureFlag(key = FlagKey.V2EncryptionPassword)) {
-            return withContext(dispatcherManager.io) {
-                authSdkSource.postKeysForUserPasswordRegistration(
+            return authSdkSource
+                .postKeysForUserPasswordRegistration(
                     email = email,
                     salt = email,
                     masterPassword = masterPassword,
                     masterPasswordHint = masterPasswordHint,
                     emailVerificationToken = emailVerificationToken,
                 )
-            }
                 .fold(
                     onSuccess = { RegisterResult.Success },
                     onFailure = { RegisterResult.Error(errorMessage = null, error = it) },
@@ -1281,18 +1278,16 @@ class AuthRepositoryImpl(
                     .map { orgKeys -> enrollStatus to orgKeys }
             }
             .flatMap { (enrollStatus, orgKeys) ->
-                withContext(dispatcherManager.io) {
-                    authSdkSource.postKeysForJitPasswordRegistration(
-                        userId = userId,
-                        organizationId = enrollStatus.organizationId,
-                        organizationPublicKey = orgKeys.publicKey,
-                        organizationSsoIdentifier = organizationIdentifier,
-                        salt = profile.email,
-                        masterPassword = password,
-                        masterPasswordHint = passwordHint,
-                        shouldResetPasswordEnroll = enrollStatus.isResetPasswordEnabled,
-                    )
-                }
+                authSdkSource.postKeysForJitPasswordRegistration(
+                    userId = userId,
+                    organizationId = enrollStatus.organizationId,
+                    organizationPublicKey = orgKeys.publicKey,
+                    organizationSsoIdentifier = organizationIdentifier,
+                    salt = profile.email,
+                    masterPassword = password,
+                    masterPasswordHint = passwordHint,
+                    shouldResetPasswordEnroll = enrollStatus.isResetPasswordEnabled,
+                )
             }
             .onSuccess { response ->
                 authDiskSource.storeAccountKeys(

app/src/main/kotlin/com/x8bit/bitwarden/data/tools/generator/datasource/sdk/GeneratorSdkSourceImpl.kt

@@ -1,11 +1,13 @@
 package com.x8bit.bitwarden.data.tools.generator.datasource.sdk
 
+import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
 import com.bitwarden.generators.PassphraseGeneratorRequest
 import com.bitwarden.generators.PasswordGeneratorRequest
 import com.bitwarden.generators.UsernameGeneratorRequest
 import com.bitwarden.sdk.GeneratorClients
 import com.x8bit.bitwarden.data.platform.datasource.sdk.BaseSdkSource
 import com.x8bit.bitwarden.data.platform.manager.SdkClientManager
+import kotlinx.coroutines.withContext
 
 /**
  * Implementation of [GeneratorSdkSource] that delegates password generation.
@@ -14,6 +16,7 @@ import com.x8bit.bitwarden.data.platform.manager.SdkClientManager
  * [GeneratorClients] provided by the Bitwarden SDK.
  */
 class GeneratorSdkSourceImpl(
+    private val dispatcherManager: DispatcherManager,
     sdkClientManager: SdkClientManager,
 ) : BaseSdkSource(sdkClientManager = sdkClientManager),
     GeneratorSdkSource {
@@ -51,6 +54,8 @@ class GeneratorSdkSourceImpl(
     override suspend fun generateForwardedServiceEmail(
         request: UsernameGeneratorRequest.Forwarded,
     ): Result<String> = runCatchingWithLogs {
-        useClient { generators().username(request) }
+        withContext(context = dispatcherManager.io) {
+            useClient { generators().username(request) }
+        }
     }
 }

app/src/main/kotlin/com/x8bit/bitwarden/data/tools/generator/datasource/sdk/di/GeneratorSdkModule.kt

@@ -1,5 +1,6 @@
 package com.x8bit.bitwarden.data.tools.generator.datasource.sdk.di
 
+import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
 import com.x8bit.bitwarden.data.platform.manager.SdkClientManager
 import com.x8bit.bitwarden.data.tools.generator.datasource.sdk.GeneratorSdkSource
 import com.x8bit.bitwarden.data.tools.generator.datasource.sdk.GeneratorSdkSourceImpl
@@ -19,6 +20,10 @@ object GeneratorSdkModule {
     @Provides
     @Singleton
     fun provideGeneratorSdkSource(
+        dispatcherManager: DispatcherManager,
         sdkClientManager: SdkClientManager,
-    ): GeneratorSdkSource = GeneratorSdkSourceImpl(sdkClientManager = sdkClientManager)
+    ): GeneratorSdkSource = GeneratorSdkSourceImpl(
+        dispatcherManager = dispatcherManager,
+        sdkClientManager = sdkClientManager,
+    )
 }

app/src/main/kotlin/com/x8bit/bitwarden/data/tools/generator/repository/GeneratorRepositoryImpl.kt

@@ -38,7 +38,6 @@ import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.receiveAsFlow
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.withContext
 import java.time.Clock
 import javax.inject.Singleton
 
@@ -55,7 +54,7 @@ class GeneratorRepositoryImpl(
     private val vaultSdkSource: VaultSdkSource,
     private val passwordHistoryDiskSource: PasswordHistoryDiskSource,
     private val reviewPromptManager: ReviewPromptManager,
-    private val dispatcherManager: DispatcherManager,
+    dispatcherManager: DispatcherManager,
 ) : GeneratorRepository {
 
     private val scope = CoroutineScope(dispatcherManager.io)
@@ -193,8 +192,9 @@ class GeneratorRepositoryImpl(
 
     override suspend fun generateForwardedServiceUsername(
         forwardedServiceGeneratorRequest: UsernameGeneratorRequest.Forwarded,
-    ): GeneratedForwardedServiceUsernameResult = withContext(dispatcherManager.io) {
-        generatorSdkSource.generateForwardedServiceEmail(forwardedServiceGeneratorRequest)
+    ): GeneratedForwardedServiceUsernameResult =
+        generatorSdkSource
+            .generateForwardedServiceEmail(forwardedServiceGeneratorRequest)
             .fold(
                 onSuccess = { generatedEmail ->
                     GeneratedForwardedServiceUsernameResult.Success(generatedEmail)
@@ -203,7 +203,6 @@ class GeneratorRepositoryImpl(
                     GeneratedForwardedServiceUsernameResult.InvalidRequest(it.message, error = it)
                 },
             )
-    }
 
     override fun getPasscodeGenerationOptions(): PasscodeGenerationOptions? {
         val userId = authDiskSource.userState?.activeUserId

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/datasource/sdk/AuthSdkSourceTest.kt

@@ -13,6 +13,7 @@ import com.bitwarden.core.KeyConnectorResponse
 import com.bitwarden.core.MasterPasswordPolicyOptions
 import com.bitwarden.core.RegisterKeyResponse
 import com.bitwarden.core.RegisterTdeKeyResponse
+import com.bitwarden.core.data.manager.dispatcher.FakeDispatcherManager
 import com.bitwarden.core.data.util.asSuccess
 import com.bitwarden.crypto.HashPurpose
 import com.bitwarden.crypto.Kdf
@@ -57,6 +58,7 @@ class AuthSdkSourceTest {
     }
 
     private val authSkdSource: AuthSdkSource = AuthSdkSourceImpl(
+        dispatcherManager = FakeDispatcherManager(),
         sdkClientManager = sdkClientManager,
     )
 

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/manager/KeyConnectorManagerTest.kt

@@ -3,7 +3,6 @@ package com.x8bit.bitwarden.data.auth.manager
 import com.bitwarden.auth.KeyConnectorRegistrationResult
 import com.bitwarden.core.KeyConnectorResponse
 import com.bitwarden.core.WrappedAccountCryptographicState
-import com.bitwarden.core.data.manager.dispatcher.FakeDispatcherManager
 import com.bitwarden.core.data.manager.model.FlagKey
 import com.bitwarden.core.data.util.asFailure
 import com.bitwarden.core.data.util.asSuccess
@@ -40,7 +39,6 @@ class KeyConnectorManagerTest {
         authSdkSource = authSdkSource,
         vaultSdkSource = vaultSdkSource,
         featureFlagManager = featureFlagManager,
-        dispatcherManager = FakeDispatcherManager(),
     )
 
     @Test

app/src/test/kotlin/com/x8bit/bitwarden/data/tools/generator/datasource/sdk/GeneratorSdkSourceTest.kt

@@ -1,5 +1,6 @@
 package com.x8bit.bitwarden.data.tools.generator.datasource.sdk
 
+import com.bitwarden.core.data.manager.dispatcher.FakeDispatcherManager
 import com.bitwarden.core.data.util.asSuccess
 import com.bitwarden.generators.AppendType
 import com.bitwarden.generators.ForwarderServiceType
@@ -27,7 +28,10 @@ class GeneratorSdkSourceTest {
         val slot = slot<suspend Client.() -> String>()
         coEvery { singleUseClient(block = capture(slot)) } coAnswers { slot.captured(client) }
     }
-    private val generatorSdkSource: GeneratorSdkSource = GeneratorSdkSourceImpl(sdkClientManager)
+    private val generatorSdkSource: GeneratorSdkSource = GeneratorSdkSourceImpl(
+        dispatcherManager = FakeDispatcherManager(),
+        sdkClientManager = sdkClientManager,
+    )
 
     @Test
     fun `generatePassword should call SDK and return a Result with the generated password`() =