BIT-1197 Add token refresh handling (#274)

2026-06-08 23:16:33 -05:00 · 2023-11-27 13:40:45 -06:00
parent acfc39ae3c
commit b914f52d0f
16 changed files with 742 additions and 18 deletions
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepository.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepository.kt
@@ -1,18 +1,19 @@
 package com.x8bit.bitwarden.data.auth.repository

 import com.x8bit.bitwarden.data.auth.datasource.sdk.model.PasswordStrength
-import com.x8bit.bitwarden.data.auth.repository.model.UserState
 import com.x8bit.bitwarden.data.auth.repository.model.AuthState
 import com.x8bit.bitwarden.data.auth.repository.model.LoginResult
 import com.x8bit.bitwarden.data.auth.repository.model.RegisterResult
+import com.x8bit.bitwarden.data.auth.repository.model.UserState
 import com.x8bit.bitwarden.data.auth.repository.util.CaptchaCallbackTokenResult
+import com.x8bit.bitwarden.data.platform.datasource.network.authenticator.AuthenticatorProvider
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow

 /**
 * Provides an API for observing an modifying authentication state.
 */
-interface AuthRepository {
+interface AuthRepository : AuthenticatorProvider {
    /**
     * Models the current auth state.
     */
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt
@@ -5,6 +5,7 @@ import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
 import com.x8bit.bitwarden.data.auth.datasource.network.model.GetTokenResponseJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.GetTokenResponseJson.CaptchaRequired
 import com.x8bit.bitwarden.data.auth.datasource.network.model.GetTokenResponseJson.Success
+import com.x8bit.bitwarden.data.auth.datasource.network.model.RefreshTokenResponseJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.RegisterRequestJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.RegisterResponseJson
 import com.x8bit.bitwarden.data.auth.datasource.network.service.AccountsService
@@ -20,6 +21,7 @@ import com.x8bit.bitwarden.data.auth.repository.model.UserState
 import com.x8bit.bitwarden.data.auth.repository.util.CaptchaCallbackTokenResult
 import com.x8bit.bitwarden.data.auth.repository.util.toSdkParams
 import com.x8bit.bitwarden.data.auth.repository.util.toUserState
+import com.x8bit.bitwarden.data.auth.repository.util.toUserStateJson
 import com.x8bit.bitwarden.data.auth.util.KdfParamsConstants.DEFAULT_PBKDF2_ITERATIONS
 import com.x8bit.bitwarden.data.auth.util.toSdkParams
 import com.x8bit.bitwarden.data.platform.manager.dispatcher.DispatcherManager
@@ -56,12 +58,13 @@ class AuthRepositoryImpl constructor(
 ) : AuthRepository {
    private val scope = CoroutineScope(dispatcherManager.io)

+    override val activeUserId: String? get() = authDiskSource.userState?.activeUserId
+
    override val authStateFlow: StateFlow<AuthState> = authDiskSource
        .userStateFlow
        .map { userState ->
            userState
                ?.let {
-                    @Suppress("UnsafeCallOnNullableType")
                    AuthState.Authenticated(
                        userState
                            .activeAccount
@@ -179,21 +182,42 @@ class AuthRepositoryImpl constructor(
            },
        )

-    override fun logout() {
-        val currentUserState = authDiskSource.userState ?: return
+    override fun refreshAccessTokenSynchronously(userId: String): Result<RefreshTokenResponseJson> {
+        val refreshAccount = authDiskSource.userState?.accounts?.get(userId)
+            ?: return IllegalStateException("Must be logged in.").asFailure()
+        return identityService
+            .refreshTokenSynchronously(refreshAccount.tokens.refreshToken)
+            .onSuccess {
+                // Update the existing UserState with updated token information
+                authDiskSource.userState = it.toUserStateJson(
+                    userId = userId,
+                    previousUserState = requireNotNull(authDiskSource.userState),
+                )
+            }
+    }

-        val activeUserId = currentUserState.activeUserId
+    override fun logout() {
+        activeUserId?.let { userId -> logout(userId) }
+    }
+
+    override fun logout(userId: String) {
+        val currentUserState = authDiskSource.userState ?: return

        // Remove the active user from the accounts map
        val updatedAccounts = currentUserState
            .accounts
-            .filterKeys { it != activeUserId }
-        authDiskSource.storeUserKey(userId = activeUserId, userKey = null)
-        authDiskSource.storePrivateKey(userId = activeUserId, privateKey = null)
+            .filterKeys { it != userId }
+        authDiskSource.storeUserKey(userId = userId, userKey = null)
+        authDiskSource.storePrivateKey(userId = userId, privateKey = null)
        // Check if there is a new active user
        if (updatedAccounts.isNotEmpty()) {
-            val (updatedActiveUserId, updatedActiveAccount) =
-                updatedAccounts.entries.first()
+            // If we logged out a non-active user, we want to leave the active user unchanged.
+            // If we logged out the active user, we want to set the active user to the first one
+            // in the list.
+            val updatedActiveUserId = currentUserState
+                .activeUserId
+                .takeUnless { it == userId }
+                ?: updatedAccounts.entries.first().key

            // Update the user information and emit an updated token
            authDiskSource.userState = currentUserState.copy(
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/util/RefreshTokenResponseExtensions.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/util/RefreshTokenResponseExtensions.kt
@@ -0,0 +1,45 @@
+package com.x8bit.bitwarden.data.auth.repository.util
+
+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.network.model.RefreshTokenResponseJson
+
+/**
+ * Converts the given [RefreshTokenResponseJson] to a [UserStateJson], given the following
+ * additional information:
+ *
+ * - the [userId]
+ * - the [previousUserState]
+ */
+fun RefreshTokenResponseJson.toUserStateJson(
+    userId: String,
+    previousUserState: UserStateJson,
+): UserStateJson {
+    val refreshedAccount = requireNotNull(previousUserState.accounts[userId])
+    val accessToken = this.accessToken
+    val jwtTokenData = requireNotNull(parseJwtTokenDataOrNull(jwtToken = accessToken))
+
+    val account = refreshedAccount.copy(
+        profile = refreshedAccount.profile.copy(
+            userId = jwtTokenData.userId,
+            email = jwtTokenData.email,
+            isEmailVerified = jwtTokenData.isEmailVerified,
+            name = jwtTokenData.name,
+            hasPremium = jwtTokenData.hasPremium,
+        ),
+        tokens = AccountJson.Tokens(
+            accessToken = accessToken,
+            refreshToken = this.refreshToken,
+        ),
+    )
+
+    // Update the existing UserState.
+    return previousUserState.copy(
+        accounts = previousUserState
+            .accounts
+            .toMutableMap()
+            .apply {
+                put(userId, account)
+            },
+    )
+}
--- a/app/src/main/java/com/x8bit/bitwarden/data/platform/datasource/network/authenticator/AuthenticatorProvider.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/platform/datasource/network/authenticator/AuthenticatorProvider.kt
@@ -0,0 +1,27 @@
+package com.x8bit.bitwarden.data.platform.datasource.network.authenticator
+
+import com.x8bit.bitwarden.data.auth.datasource.network.model.RefreshTokenResponseJson
+
+/**
+ * A provider for all the functionality needed to properly refresh the users access token.
+ */
+interface AuthenticatorProvider {
+
+    /**
+     * The currently active user's ID.
+     */
+    val activeUserId: String?
+
+    /**
+     * Attempts to logout the user based on the [userId].
+     */
+    fun logout(userId: String)
+
+    /**
+     * Attempt to refresh the user's access token based on the [userId].
+     *
+     * This call is both synchronous and performs a network request. Make sure that you are calling
+     * from an appropriate thread.
+     */
+    fun refreshAccessTokenSynchronously(userId: String): Result<RefreshTokenResponseJson>
+}
--- a/app/src/main/java/com/x8bit/bitwarden/data/platform/datasource/network/authenticator/RefreshAuthenticator.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/platform/datasource/network/authenticator/RefreshAuthenticator.kt
@@ -0,0 +1,69 @@
+package com.x8bit.bitwarden.data.platform.datasource.network.authenticator
+
+import com.x8bit.bitwarden.data.auth.repository.util.parseJwtTokenDataOrNull
+import com.x8bit.bitwarden.data.platform.datasource.network.util.HEADER_KEY_AUTHORIZATION
+import com.x8bit.bitwarden.data.platform.datasource.network.util.HEADER_VALUE_BEARER_PREFIX
+import okhttp3.Authenticator
+import okhttp3.Request
+import okhttp3.Response
+import okhttp3.Route
+import javax.inject.Singleton
+
+/**
+ * An authenticator used to refresh the access token when a 401 is returned from an API. Upon
+ * successfully getting a new access token, the original request is retried.
+ */
+@Singleton
+class RefreshAuthenticator : Authenticator {
+
+    /**
+     * A provider required to update tokens.
+     */
+    var authenticatorProvider: AuthenticatorProvider? = null
+
+    override fun authenticate(
+        route: Route?,
+        response: Response,
+    ): Request? {
+        val accessToken = requireNotNull(
+            response
+                .request
+                .header(name = HEADER_KEY_AUTHORIZATION)
+                ?.substringAfter(delimiter = HEADER_VALUE_BEARER_PREFIX),
+        )
+        return when (val userId = parseJwtTokenDataOrNull(accessToken)?.userId) {
+            null -> {
+                // We unable to get the user ID, let's just let the 401 pass through.
+                null
+            }
+
+            authenticatorProvider?.activeUserId -> {
+                // In order to prevent potential deadlocks or thread starvation we want the call
+                // to refresh the access token to be strictly synchronous with no internal thread
+                // hopping.
+                authenticatorProvider
+                    ?.refreshAccessTokenSynchronously(userId)
+                    ?.fold(
+                        onFailure = {
+                            authenticatorProvider?.logout(userId)
+                            null
+                        },
+                        onSuccess = {
+                            response.request
+                                .newBuilder()
+                                .header(
+                                    name = HEADER_KEY_AUTHORIZATION,
+                                    value = "$HEADER_VALUE_BEARER_PREFIX${it.accessToken}",
+                                )
+                                .build()
+                        },
+                    )
+            }
+
+            else -> {
+                // We are no longer the active user, let's just cancel.
+                null
+            }
+        }
+    }
+}
--- a/app/src/main/java/com/x8bit/bitwarden/data/platform/datasource/network/di/PlatformNetworkModule.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/platform/datasource/network/di/PlatformNetworkModule.kt
@@ -1,5 +1,6 @@
 package com.x8bit.bitwarden.data.platform.datasource.network.di

+import com.x8bit.bitwarden.data.platform.datasource.network.authenticator.RefreshAuthenticator
 import com.x8bit.bitwarden.data.platform.datasource.network.interceptor.AuthTokenInterceptor
 import com.x8bit.bitwarden.data.platform.datasource.network.interceptor.BaseUrlInterceptors
 import com.x8bit.bitwarden.data.platform.datasource.network.retrofit.Retrofits
@@ -36,16 +37,22 @@ object PlatformNetworkModule {
    @Singleton
    fun providesAuthTokenInterceptor(): AuthTokenInterceptor = AuthTokenInterceptor()

+    @Provides
+    @Singleton
+    fun providesRefreshAuthenticator(): RefreshAuthenticator = RefreshAuthenticator()
+
    @Provides
    @Singleton
    fun provideRetrofits(
        authTokenInterceptor: AuthTokenInterceptor,
        baseUrlInterceptors: BaseUrlInterceptors,
+        refreshAuthenticator: RefreshAuthenticator,
        json: Json,
    ): Retrofits =
        RetrofitsImpl(
            authTokenInterceptor = authTokenInterceptor,
            baseUrlInterceptors = baseUrlInterceptors,
+            refreshAuthenticator = refreshAuthenticator,
            json = json,
        )

--- a/app/src/main/java/com/x8bit/bitwarden/data/platform/datasource/network/interceptor/AuthTokenInterceptor.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/platform/datasource/network/interceptor/AuthTokenInterceptor.kt
@@ -1,5 +1,7 @@
 package com.x8bit.bitwarden.data.platform.datasource.network.interceptor

+import com.x8bit.bitwarden.data.platform.datasource.network.util.HEADER_KEY_AUTHORIZATION
+import com.x8bit.bitwarden.data.platform.datasource.network.util.HEADER_VALUE_BEARER_PREFIX
 import okhttp3.Interceptor
 import okhttp3.Response
 import java.io.IOException
@@ -22,7 +24,10 @@ class AuthTokenInterceptor : Interceptor {
        val request = chain
            .request()
            .newBuilder()
-            .addHeader("Authorization", "Bearer $token")
+            .addHeader(
+                name = HEADER_KEY_AUTHORIZATION,
+                value = "$HEADER_VALUE_BEARER_PREFIX$token",
+            )
            .build()
        return chain
            .proceed(request)
--- a/app/src/main/java/com/x8bit/bitwarden/data/platform/datasource/network/retrofit/RetrofitsImpl.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/platform/datasource/network/retrofit/RetrofitsImpl.kt
@@ -1,6 +1,7 @@
 package com.x8bit.bitwarden.data.platform.datasource.network.retrofit

 import com.jakewharton.retrofit2.converter.kotlinx.serialization.asConverterFactory
+import com.x8bit.bitwarden.data.platform.datasource.network.authenticator.RefreshAuthenticator
 import com.x8bit.bitwarden.data.platform.datasource.network.core.ResultCallAdapterFactory
 import com.x8bit.bitwarden.data.platform.datasource.network.interceptor.AuthTokenInterceptor
 import com.x8bit.bitwarden.data.platform.datasource.network.interceptor.BaseUrlInterceptor
@@ -17,6 +18,7 @@ import retrofit2.Retrofit
 class RetrofitsImpl(
    authTokenInterceptor: AuthTokenInterceptor,
    baseUrlInterceptors: BaseUrlInterceptors,
+    refreshAuthenticator: RefreshAuthenticator,
    json: Json,
 ) : Retrofits {
    //region Authenticated Retrofits
@@ -73,6 +75,7 @@ class RetrofitsImpl(
    private val authenticatedOkHttpClient: OkHttpClient by lazy {
        baseOkHttpClient
            .newBuilder()
+            .authenticator(refreshAuthenticator)
            .addInterceptor(authTokenInterceptor)
            .build()
    }
--- a/app/src/main/java/com/x8bit/bitwarden/data/platform/datasource/network/util/HeaderUtils.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/platform/datasource/network/util/HeaderUtils.kt
@@ -0,0 +1,11 @@
+package com.x8bit.bitwarden.data.platform.datasource.network.util
+
+/**
+ * The bearer prefix used for the 'authorization' headers value.
+ */
+const val HEADER_VALUE_BEARER_PREFIX: String = "Bearer "
+
+/**
+ * The key used for the 'authorization' headers.
+ */
+const val HEADER_KEY_AUTHORIZATION: String = "Authorization"
--- a/app/src/main/java/com/x8bit/bitwarden/data/platform/manager/NetworkConfigManagerImpl.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/platform/manager/NetworkConfigManagerImpl.kt
@@ -2,6 +2,7 @@ package com.x8bit.bitwarden.data.platform.manager

 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
 import com.x8bit.bitwarden.data.auth.repository.model.AuthState
+import com.x8bit.bitwarden.data.platform.datasource.network.authenticator.RefreshAuthenticator
 import com.x8bit.bitwarden.data.platform.datasource.network.interceptor.AuthTokenInterceptor
 import com.x8bit.bitwarden.data.platform.datasource.network.interceptor.BaseUrlInterceptors
 import com.x8bit.bitwarden.data.platform.manager.dispatcher.DispatcherManager
@@ -18,6 +19,7 @@ class NetworkConfigManagerImpl(
    private val authTokenInterceptor: AuthTokenInterceptor,
    private val environmentRepository: EnvironmentRepository,
    private val baseUrlInterceptors: BaseUrlInterceptors,
+    refreshAuthenticator: RefreshAuthenticator,
    dispatcherManager: DispatcherManager,
 ) : NetworkConfigManager {

@@ -41,5 +43,7 @@ class NetworkConfigManagerImpl(
                baseUrlInterceptors.environment = environment
            }
            .launchIn(scope)
+
+        refreshAuthenticator.authenticatorProvider = authRepository
    }
 }
--- a/app/src/main/java/com/x8bit/bitwarden/data/platform/manager/di/PlatformManagerModule.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/platform/manager/di/PlatformManagerModule.kt
@@ -1,6 +1,7 @@
 package com.x8bit.bitwarden.data.platform.manager.di

 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
+import com.x8bit.bitwarden.data.platform.datasource.network.authenticator.RefreshAuthenticator
 import com.x8bit.bitwarden.data.platform.datasource.network.interceptor.AuthTokenInterceptor
 import com.x8bit.bitwarden.data.platform.datasource.network.interceptor.BaseUrlInterceptors
 import com.x8bit.bitwarden.data.platform.manager.NetworkConfigManager
@@ -25,6 +26,7 @@ object PlatformManagerModule {
    @Singleton
    fun provideBitwardenDispatchers(): DispatcherManager = DispatcherManagerImpl()

+    @Suppress("LongParameterList")
    @Provides
    @Singleton
    fun provideNetworkConfigManager(
@@ -32,6 +34,7 @@ object PlatformManagerModule {
        authTokenInterceptor: AuthTokenInterceptor,
        environmentRepository: EnvironmentRepository,
        baseUrlInterceptors: BaseUrlInterceptors,
+        refreshAuthenticator: RefreshAuthenticator,
        dispatcherManager: DispatcherManager,
    ): NetworkConfigManager =
        NetworkConfigManagerImpl(
@@ -39,6 +42,7 @@ object PlatformManagerModule {
            authTokenInterceptor = authTokenInterceptor,
            environmentRepository = environmentRepository,
            baseUrlInterceptors = baseUrlInterceptors,
+            refreshAuthenticator = refreshAuthenticator,
            dispatcherManager = dispatcherManager,
        )
 }