[PM-41330] feat: Add custom header support for self-hosted environments

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

@@ -11,6 +11,7 @@ import com.x8bit.bitwarden.data.auth.repository.model.LogoutReason
 import com.x8bit.bitwarden.data.platform.datasource.disk.PushDiskSource
 import com.x8bit.bitwarden.data.platform.datasource.disk.SettingsDiskSource
 import com.x8bit.bitwarden.data.platform.manager.CredentialExchangeRegistryManager
+import com.x8bit.bitwarden.data.platform.manager.CustomHeadersManager
 import com.x8bit.bitwarden.data.tools.generator.datasource.disk.GeneratorDiskSource
 import com.x8bit.bitwarden.data.tools.generator.datasource.disk.PasswordHistoryDiskSource
 import com.x8bit.bitwarden.data.vault.datasource.disk.VaultDiskSource
@@ -28,6 +29,7 @@ import timber.log.Timber
 @Suppress("LongParameterList")
 class UserLogoutManagerImpl(
     private val authDiskSource: AuthDiskSource,
+    private val customHeadersManager: CustomHeadersManager,
     private val generatorDiskSource: GeneratorDiskSource,
     private val passwordHistoryDiskSource: PasswordHistoryDiskSource,
     private val pushDiskSource: PushDiskSource,
@@ -49,6 +51,11 @@ class UserLogoutManagerImpl(
     override fun logout(userId: String, reason: LogoutReason) {
         authDiskSource.userState ?: return
         Timber.d("logout reason=$reason")
+
+        // Clean up the account's stored custom headers while its environment data is still
+        // available. The removal is reference-counted, so headers still used by another account
+        // or the pre-auth environment remain.
+        customHeadersManager.removeCustomHeadersForUser(userId = userId)
         val isSecurityStamp = reason == LogoutReason.SecurityStamp
         if (isSecurityStamp) {
             showToast(message = BitwardenString.login_expired)

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

@@ -28,6 +28,7 @@ import com.x8bit.bitwarden.data.auth.manager.UserLogoutManagerImpl
 import com.x8bit.bitwarden.data.platform.datasource.disk.PushDiskSource
 import com.x8bit.bitwarden.data.platform.datasource.disk.SettingsDiskSource
 import com.x8bit.bitwarden.data.platform.manager.CredentialExchangeRegistryManager
+import com.x8bit.bitwarden.data.platform.manager.CustomHeadersManager
 import com.x8bit.bitwarden.data.platform.manager.FeatureFlagManager
 import com.x8bit.bitwarden.data.platform.manager.PushManager
 import com.x8bit.bitwarden.data.tools.generator.datasource.disk.GeneratorDiskSource
@@ -115,6 +116,7 @@ object AuthManagerModule {
     @Singleton
     fun provideUserLogoutManager(
         authDiskSource: AuthDiskSource,
+        customHeadersManager: CustomHeadersManager,
         generatorDiskSource: GeneratorDiskSource,
         passwordHistoryDiskSource: PasswordHistoryDiskSource,
         pushDiskSource: PushDiskSource,
@@ -127,6 +129,7 @@ object AuthManagerModule {
     ): UserLogoutManager =
         UserLogoutManagerImpl(
             authDiskSource = authDiskSource,
+            customHeadersManager = customHeadersManager,
             generatorDiskSource = generatorDiskSource,
             passwordHistoryDiskSource = passwordHistoryDiskSource,
             pushDiskSource = pushDiskSource,

app/src/main/kotlin/com/x8bit/bitwarden/data/platform/datasource/disk/CustomHeadersDiskSource.kt

@@ -0,0 +1,23 @@
+package com.x8bit.bitwarden.data.platform.datasource.disk
+
+/**
+ * Disk source for persisting the custom headers sent with requests to a self-hosted environment.
+ */
+interface CustomHeadersDiskSource {
+
+    /**
+     * Gets the custom headers stored under the given [id].
+     *
+     * @param id The identifier of the custom headers.
+     * @return The custom headers, or null if none are stored.
+     */
+    fun getCustomHeaders(id: String): Map<String, String>?
+
+    /**
+     * Stores the custom [headers] under the given [id]. Pass `null` to delete the stored headers.
+     *
+     * @param id The identifier to store the custom headers under.
+     * @param headers The custom headers to persist, or `null` to delete.
+     */
+    fun storeCustomHeaders(id: String, headers: Map<String, String>?)
+}

app/src/main/kotlin/com/x8bit/bitwarden/data/platform/datasource/disk/CustomHeadersDiskSourceImpl.kt

@@ -0,0 +1,37 @@
+package com.x8bit.bitwarden.data.platform.datasource.disk
+
+import android.content.SharedPreferences
+import com.bitwarden.core.data.util.decodeFromStringOrNull
+import com.bitwarden.data.datasource.disk.BaseEncryptedDiskSource
+import kotlinx.serialization.json.Json
+
+private const val CUSTOM_HEADERS_PREFIX = "customHeaders"
+
+/**
+ * Implementation of [CustomHeadersDiskSource] using encrypted SharedPreferences.
+ *
+ * The header values may contain credentials, such as Cloudflare Access service tokens, so they
+ * are only ever written to encrypted storage. The environment data referencing them stores only
+ * the opaque identifier.
+ */
+class CustomHeadersDiskSourceImpl(
+    sharedPreferences: SharedPreferences,
+    encryptedSharedPreferences: SharedPreferences,
+    private val json: Json,
+) : CustomHeadersDiskSource,
+    BaseEncryptedDiskSource(
+        sharedPreferences = sharedPreferences,
+        encryptedSharedPreferences = encryptedSharedPreferences,
+    ) {
+
+    override fun getCustomHeaders(id: String): Map<String, String>? =
+        getEncryptedString(key = CUSTOM_HEADERS_PREFIX.appendIdentifier(id))
+            ?.let { json.decodeFromStringOrNull<Map<String, String>>(it) }
+
+    override fun storeCustomHeaders(id: String, headers: Map<String, String>?) {
+        putEncryptedString(
+            key = CUSTOM_HEADERS_PREFIX.appendIdentifier(id),
+            value = headers?.let { json.encodeToString(it) },
+        )
+    }
+}

app/src/main/kotlin/com/x8bit/bitwarden/data/platform/datasource/disk/di/PlatformDiskModule.kt

@@ -10,6 +10,8 @@ import com.bitwarden.data.datasource.disk.di.EncryptedPreferences
 import com.bitwarden.data.datasource.disk.di.UnencryptedPreferences
 import com.x8bit.bitwarden.data.platform.datasource.disk.CookieDiskSource
 import com.x8bit.bitwarden.data.platform.datasource.disk.CookieDiskSourceImpl
+import com.x8bit.bitwarden.data.platform.datasource.disk.CustomHeadersDiskSource
+import com.x8bit.bitwarden.data.platform.datasource.disk.CustomHeadersDiskSourceImpl
 import com.x8bit.bitwarden.data.platform.datasource.disk.EnvironmentDiskSource
 import com.x8bit.bitwarden.data.platform.datasource.disk.EnvironmentDiskSourceImpl
 import com.x8bit.bitwarden.data.platform.datasource.disk.EventDiskSource
@@ -158,6 +160,18 @@ object PlatformDiskModule {
         sharedPreferences = sharedPreferences,
     )
 
+    @Provides
+    @Singleton
+    fun provideCustomHeadersDiskSource(
+        @UnencryptedPreferences sharedPreferences: SharedPreferences,
+        @EncryptedPreferences encryptedSharedPreferences: SharedPreferences,
+        json: Json,
+    ): CustomHeadersDiskSource = CustomHeadersDiskSourceImpl(
+        sharedPreferences = sharedPreferences,
+        encryptedSharedPreferences = encryptedSharedPreferences,
+        json = json,
+    )
+
     @Provides
     @Singleton
     fun provideCookieDiskSource(

app/src/main/kotlin/com/x8bit/bitwarden/data/platform/datasource/network/di/PlatformNetworkModule.kt

@@ -15,6 +15,7 @@ import com.x8bit.bitwarden.data.platform.datasource.network.util.HEADER_VALUE_CL
 import com.x8bit.bitwarden.data.platform.datasource.network.util.HEADER_VALUE_CLIENT_VERSION
 import com.x8bit.bitwarden.data.platform.datasource.network.util.HEADER_VALUE_USER_AGENT
 import com.x8bit.bitwarden.data.platform.manager.CertificateManager
+import com.x8bit.bitwarden.data.platform.manager.CustomHeadersManager
 import com.x8bit.bitwarden.data.platform.manager.network.NetworkCookieManager
 import com.x8bit.bitwarden.data.platform.manager.network.NetworkPermissionManager
 import dagger.Module
@@ -64,6 +65,7 @@ object PlatformNetworkModule {
         baseUrlsProvider: BaseUrlsProvider,
         authDiskSource: AuthDiskSource,
         certificateManager: CertificateManager,
+        customHeadersManager: CustomHeadersManager,
         buildInfoManager: BuildInfoManager,
         networkCookieManager: NetworkCookieManager,
         networkPermissionManager: NetworkPermissionManager,
@@ -79,6 +81,7 @@ object PlatformNetworkModule {
         authTokenProvider = authTokenManager,
         baseUrlsProvider = baseUrlsProvider,
         certificateProvider = certificateManager,
+        customHeadersProvider = customHeadersManager,
         enableHttpBodyLogging = buildInfoManager.isDevBuild,
         cookieProvider = networkCookieManager,
         permissionProvider = networkPermissionManager,

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

@@ -0,0 +1,45 @@
+package com.x8bit.bitwarden.data.platform.manager
+
+import com.bitwarden.network.provider.CustomHeadersProvider
+
+/**
+ * Responsible for managing the custom headers sent with requests to a self-hosted environment.
+ *
+ * The header values are kept in encrypted storage under an opaque identifier; only that
+ * identifier is persisted with the environment data.
+ */
+interface CustomHeadersManager : CustomHeadersProvider {
+
+    /**
+     * Gets the custom headers stored under the given [id].
+     *
+     * @param id The identifier of the custom headers.
+     * @return The custom headers, or null if none are stored.
+     */
+    fun getStoredCustomHeaders(id: String): Map<String, String>?
+
+    /**
+     * Stores the custom [headers] and returns the identifier under which they were saved.
+     *
+     * @param headers The custom headers to store.
+     * @return The identifier of the stored custom headers.
+     */
+    fun saveCustomHeaders(headers: Map<String, String>): String
+
+    /**
+     * Removes the custom headers with the given [id] from storage if no environment still
+     * references them.
+     *
+     * @param id The identifier of the custom headers to remove.
+     */
+    fun removeCustomHeaders(id: String)
+
+    /**
+     * Removes the custom headers referenced by the environment of the account with the given
+     * [userId] if no other environment still references them. This must be called while the
+     * account's data is still available.
+     *
+     * @param userId The user ID of the account being removed.
+     */
+    fun removeCustomHeadersForUser(userId: String)
+}

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

@@ -0,0 +1,90 @@
+package com.x8bit.bitwarden.data.platform.manager
+
+import com.bitwarden.data.repository.util.baseApiUrl
+import com.bitwarden.data.repository.util.baseEventsUrl
+import com.bitwarden.data.repository.util.baseIconUrl
+import com.bitwarden.data.repository.util.baseIdentityUrl
+import com.bitwarden.data.repository.util.baseWebVaultUrlOrNull
+import com.bitwarden.data.repository.util.toEnvironmentUrls
+import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
+import com.x8bit.bitwarden.data.platform.datasource.disk.CustomHeadersDiskSource
+import com.x8bit.bitwarden.data.platform.datasource.disk.EnvironmentDiskSource
+import com.x8bit.bitwarden.data.platform.util.toUriOrNull
+import java.util.UUID
+
+/**
+ * Default implementation of [CustomHeadersManager].
+ */
+class CustomHeadersManagerImpl(
+    private val authDiskSource: AuthDiskSource,
+    private val customHeadersDiskSource: CustomHeadersDiskSource,
+    private val environmentDiskSource: EnvironmentDiskSource,
+) : CustomHeadersManager {
+
+    override fun getCustomHeaders(url: String): Map<String, String> {
+        val environmentUrlData = environmentDiskSource.preAuthEnvironmentUrlData
+        val id = environmentUrlData?.customHeadersId ?: return emptyMap()
+        val requestUri = url.toUriOrNull() ?: return emptyMap()
+        val requestHost = requestUri.host ?: return emptyMap()
+
+        // Scope the headers to the environment's URLs, matching on both scheme and host, so
+        // their values, which may contain credentials, are never sent to third parties or
+        // downgraded to a cleartext connection.
+        val environment = environmentUrlData.toEnvironmentUrls()
+        val isEnvironmentUrl = listOfNotNull(
+            environment.baseApiUrl,
+            environment.baseEventsUrl,
+            environment.baseIconUrl,
+            environment.baseIdentityUrl,
+            environment.baseWebVaultUrlOrNull,
+        )
+            .mapNotNull { it.toUriOrNull() }
+            .any { it.scheme == requestUri.scheme && it.host == requestHost }
+        if (!isEnvironmentUrl) return emptyMap()
+
+        return customHeadersDiskSource.getCustomHeaders(id = id).orEmpty()
+    }
+
+    override fun getStoredCustomHeaders(id: String): Map<String, String>? =
+        customHeadersDiskSource.getCustomHeaders(id = id)
+
+    override fun saveCustomHeaders(headers: Map<String, String>): String {
+        val id = UUID.randomUUID().toString()
+        customHeadersDiskSource.storeCustomHeaders(id = id, headers = headers)
+        return id
+    }
+
+    override fun removeCustomHeaders(id: String) {
+        if (isIdInUse(id = id)) return
+        customHeadersDiskSource.storeCustomHeaders(id = id, headers = null)
+    }
+
+    override fun removeCustomHeadersForUser(userId: String) {
+        val id = authDiskSource
+            .userState
+            ?.accounts
+            ?.get(userId)
+            ?.settings
+            ?.environmentUrlData
+            ?.customHeadersId
+            ?: return
+        if (isIdInUse(id = id, excludedUserId = userId)) return
+        customHeadersDiskSource.storeCustomHeaders(id = id, headers = null)
+    }
+
+    /**
+     * Returns whether the pre-auth environment or any account's environment, other than the
+     * [excludedUserId] account's, still references the given custom headers [id].
+     */
+    private fun isIdInUse(id: String, excludedUserId: String? = null): Boolean {
+        if (environmentDiskSource.preAuthEnvironmentUrlData?.customHeadersId == id) return true
+        return authDiskSource
+            .userState
+            ?.accounts
+            .orEmpty()
+            .any { (userId, account) ->
+                userId != excludedUserId &&
+                    account.settings.environmentUrlData?.customHeadersId == id
+            }
+    }
+}

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

@@ -26,6 +26,8 @@ import com.x8bit.bitwarden.data.autofill.accessibility.manager.AccessibilityEnab
 import com.x8bit.bitwarden.data.autofill.manager.AutofillEnabledManager
 import com.x8bit.bitwarden.data.autofill.manager.browser.BrowserThirdPartyAutofillEnabledManager
 import com.x8bit.bitwarden.data.platform.datasource.disk.CookieDiskSource
+import com.x8bit.bitwarden.data.platform.datasource.disk.CustomHeadersDiskSource
+import com.x8bit.bitwarden.data.platform.datasource.disk.EnvironmentDiskSource
 import com.x8bit.bitwarden.data.platform.datasource.disk.EventDiskSource
 import com.x8bit.bitwarden.data.platform.datasource.disk.PushDiskSource
 import com.x8bit.bitwarden.data.platform.datasource.disk.SettingsDiskSource
@@ -42,6 +44,8 @@ import com.x8bit.bitwarden.data.platform.manager.CookieAcquisitionRequestManager
 import com.x8bit.bitwarden.data.platform.manager.CookieAcquisitionRequestManagerImpl
 import com.x8bit.bitwarden.data.platform.manager.CredentialExchangeRegistryManager
 import com.x8bit.bitwarden.data.platform.manager.CredentialExchangeRegistryManagerImpl
+import com.x8bit.bitwarden.data.platform.manager.CustomHeadersManager
+import com.x8bit.bitwarden.data.platform.manager.CustomHeadersManagerImpl
 import com.x8bit.bitwarden.data.platform.manager.DatabaseSchemeManager
 import com.x8bit.bitwarden.data.platform.manager.DatabaseSchemeManagerImpl
 import com.x8bit.bitwarden.data.platform.manager.DebugMenuFeatureFlagManagerImpl
@@ -405,6 +409,18 @@ object PlatformManagerModule {
         environmentRepository = environmentRepository,
     )
 
+    @Provides
+    @Singleton
+    fun provideCustomHeadersManager(
+        authDiskSource: AuthDiskSource,
+        customHeadersDiskSource: CustomHeadersDiskSource,
+        environmentDiskSource: EnvironmentDiskSource,
+    ): CustomHeadersManager = CustomHeadersManagerImpl(
+        authDiskSource = authDiskSource,
+        customHeadersDiskSource = customHeadersDiskSource,
+        environmentDiskSource = environmentDiskSource,
+    )
+
     @Provides
     @Singleton
     fun provideAppResumeManager(

app/src/main/kotlin/com/x8bit/bitwarden/data/platform/repository/EnvironmentRepositoryImpl.kt

@@ -6,6 +6,7 @@ import com.bitwarden.data.repository.util.toEnvironmentUrls
 import com.bitwarden.data.repository.util.toEnvironmentUrlsOrDefault
 import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
 import com.x8bit.bitwarden.data.platform.datasource.disk.EnvironmentDiskSource
+import com.x8bit.bitwarden.data.platform.manager.CustomHeadersManager
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
@@ -22,6 +23,7 @@ import timber.log.Timber
 class EnvironmentRepositoryImpl(
     private val environmentDiskSource: EnvironmentDiskSource,
     private val authDiskSource: AuthDiskSource,
+    private val customHeadersManager: CustomHeadersManager,
     dispatcherManager: DispatcherManager,
 ) : EnvironmentRepository {
 
@@ -32,7 +34,14 @@ class EnvironmentRepositoryImpl(
             .preAuthEnvironmentUrlData
             .toEnvironmentUrlsOrDefault()
         set(value) {
+            val droppedCustomHeadersId = environmentDiskSource
+                .preAuthEnvironmentUrlData
+                ?.customHeadersId
+                ?.takeUnless { it == value.environmentUrlData.customHeadersId }
             environmentDiskSource.preAuthEnvironmentUrlData = value.environmentUrlData
+            // Replacing the environment drops its reference to any stored custom headers. The
+            // removal is reference-counted, so headers still used by an account remain.
+            droppedCustomHeadersId?.let { customHeadersManager.removeCustomHeaders(id = it) }
         }
 
     override val environmentStateFlow: StateFlow<Environment> = environmentDiskSource

app/src/main/kotlin/com/x8bit/bitwarden/data/platform/repository/di/PlatformRepositoryModule.kt

@@ -12,6 +12,7 @@ import com.x8bit.bitwarden.data.platform.datasource.disk.CookieDiskSource
 import com.x8bit.bitwarden.data.platform.datasource.disk.EnvironmentDiskSource
 import com.x8bit.bitwarden.data.platform.datasource.disk.FeatureFlagOverrideDiskSource
 import com.x8bit.bitwarden.data.platform.datasource.disk.SettingsDiskSource
+import com.x8bit.bitwarden.data.platform.manager.CustomHeadersManager
 import com.x8bit.bitwarden.data.platform.manager.PolicyManager
 import com.x8bit.bitwarden.data.platform.repository.AuthenticatorBridgeRepository
 import com.x8bit.bitwarden.data.platform.repository.AuthenticatorBridgeRepositoryImpl
@@ -54,11 +55,13 @@ object PlatformRepositoryModule {
     fun provideEnvironmentRepository(
         environmentDiskSource: EnvironmentDiskSource,
         authDiskSource: AuthDiskSource,
+        customHeadersManager: CustomHeadersManager,
         dispatcherManager: DispatcherManager,
     ): EnvironmentRepository =
         EnvironmentRepositoryImpl(
             environmentDiskSource = environmentDiskSource,
             authDiskSource = authDiskSource,
+            customHeadersManager = customHeadersManager,
             dispatcherManager = dispatcherManager,
         )