Iterate on client creation

features/enterprise/api/src/main/kotlin/io/element/android/features/enterprise/api/ClientBuilderEnterpriseHook.kt

@@ -13,6 +13,23 @@ import io.element.android.libraries.matrix.api.core.SessionId
 /**
  * A hook that can be used to customize the [MatrixClientBuilder] for enterprise features.
  */
-fun interface ClientBuilderEnterpriseHook {
-    suspend operator fun invoke(clientBuilder: MatrixClientBuilder, sessionId: SessionId): MatrixClientBuilder
+interface ClientBuilderEnterpriseHook {
+    /**
+     * Customize the [MatrixClientBuilder] for enterprise features.
+     * This method is invoked everytime a new [MatrixClientBuilder] is created.
+     *
+     * @param clientBuilder The [MatrixClientBuilder] to customize.
+     * @return The customized [MatrixClientBuilder].
+     */
+    suspend fun beforeClientCreation(clientBuilder: MatrixClientBuilder): MatrixClientBuilder
+
+    /**
+     * Customize the [MatrixClientBuilder] for enterprise features.
+     * This method is invoked as well as the other method when a new [MatrixClientBuilder] is created to build a client for a specific session.
+     *
+     * @param clientBuilder The [MatrixClientBuilder] to customize.
+     * @param sessionId The [SessionId] for which the [MatrixClientBuilder] is being created.
+     * @return The customized [MatrixClientBuilder].
+     */
+    suspend fun beforeClientCreationWithSession(clientBuilder: MatrixClientBuilder, sessionId: SessionId): MatrixClientBuilder
 }

features/enterprise/impl-foss/src/main/kotlin/io/element/android/features/enterprise/impl/DefaultClientBuilderEnterpriseHook.kt

@@ -18,7 +18,12 @@ import io.element.android.libraries.matrix.api.core.SessionId
  */
 @ContributesBinding(AppScope::class)
 class DefaultClientBuilderEnterpriseHook : ClientBuilderEnterpriseHook {
-    override suspend fun invoke(clientBuilder: MatrixClientBuilder, sessionId: SessionId): MatrixClientBuilder {
+    override suspend fun beforeClientCreation(clientBuilder: MatrixClientBuilder): MatrixClientBuilder {
+        // No modification
+        return clientBuilder
+    }
+
+    override suspend fun beforeClientCreationWithSession(clientBuilder: MatrixClientBuilder, sessionId: SessionId): MatrixClientBuilder {
         // No modification
         return clientBuilder
     }

features/enterprise/test/src/main/kotlin/io/element/android/features/enterprise/test/FakeClientBuilderEnterpriseHook.kt

@@ -0,0 +1,25 @@
+/*
+ * Copyright (c) 2026 Element Creations Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial.
+ * Please see LICENSE files in the repository root for full details.
+ */
+
+package io.element.android.features.enterprise.test
+
+import io.element.android.features.enterprise.api.ClientBuilderEnterpriseHook
+import io.element.android.libraries.matrix.api.MatrixClientBuilder
+import io.element.android.libraries.matrix.api.core.SessionId
+
+class FakeClientBuilderEnterpriseHook(
+    private val beforeClientCreationResult: (MatrixClientBuilder) -> MatrixClientBuilder = { it },
+    private val beforeClientCreationWithSessionResult: (MatrixClientBuilder, SessionId) -> MatrixClientBuilder = { clientBuilder, _ -> clientBuilder },
+) : ClientBuilderEnterpriseHook {
+    override suspend fun beforeClientCreation(clientBuilder: MatrixClientBuilder): MatrixClientBuilder {
+        return beforeClientCreationResult(clientBuilder)
+    }
+
+    override suspend fun beforeClientCreationWithSession(clientBuilder: MatrixClientBuilder, sessionId: SessionId): MatrixClientBuilder {
+        return beforeClientCreationWithSessionResult(clientBuilder, sessionId)
+    }
+}

libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/RustMatrixClientFactory.kt

@@ -35,7 +35,6 @@ import io.element.android.services.analytics.api.AnalyticsService
 import io.element.android.services.toolbox.api.systemclock.SystemClock
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.withContext
-import org.matrix.rustcomponents.sdk.Client
 import org.matrix.rustcomponents.sdk.ClientBuilder
 import org.matrix.rustcomponents.sdk.CrossProcessLockConfig
 import org.matrix.rustcomponents.sdk.RequestConfig
@@ -100,7 +99,12 @@ class RustMatrixClientFactory(
         )
             .homeserverUrl(sessionData.homeserverUrl)
             .enableAutomaticBackPagination(featureFlagService.isFeatureEnabled(FeatureFlags.AutomaticBackPagination))
-            .let { (clientBuilderEnterpriseHook(RustMatrixClientBuilder(it), SessionId(sessionData.userId)) as RustMatrixClientBuilder).inner }
+            .let {
+                (clientBuilderEnterpriseHook.beforeClientCreationWithSession(
+                    clientBuilder = RustMatrixClientBuilder(it),
+                    sessionId = SessionId(sessionData.userId),
+                ) as RustMatrixClientBuilder).inner
+            }
             .use { it.build() }
 
         client.setMediaRetentionPolicy(
@@ -118,14 +122,6 @@ class RustMatrixClientFactory(
 
         client.restoreSession(sessionData.toSession())
 
-        create(client, sessionData, isMessageSearchAvailable)
-    }
-
-    suspend fun create(
-        client: Client,
-        sessionData: SessionData,
-        isMessageSearchAvailable: Boolean,
-    ): RustMatrixClient {
         val (anonymizedAccessToken, anonymizedRefreshToken) = client.session().anonymizedTokens()
 
         client.setUtdDelegate(UtdTracker(analyticsService))
@@ -136,7 +132,7 @@ class RustMatrixClientFactory(
             .withProfilesExtension()
             .finish()
 
-        return RustMatrixClient(
+        RustMatrixClient(
             sessionPaths = sessionData.getSessionPaths(),
             innerClient = client,
             sessionStore = sessionStore,
@@ -228,6 +224,9 @@ class RustMatrixClientFactory(
                 // Workaround for non-nullable proxy parameter in the SDK, since each call to the ClientBuilder returns a new reference we need to keep
                 proxyProvider.provides()?.let { proxy(it) } ?: this
             }
+            .let {
+                (clientBuilderEnterpriseHook.beforeClientCreation(RustMatrixClientBuilder(it)) as RustMatrixClientBuilder).inner
+            }
     }
 }
 

libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/auth/RustMatrixAuthenticationService.kt

@@ -44,6 +44,7 @@ import io.element.android.libraries.matrix.impl.keys.SecretGenerator
 import io.element.android.libraries.matrix.impl.mapper.toSessionData
 import io.element.android.libraries.matrix.impl.paths.SessionPathsFactory
 import io.element.android.libraries.sessionstorage.api.LoginType
+import io.element.android.libraries.sessionstorage.api.SessionData
 import io.element.android.libraries.sessionstorage.api.SessionStore
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.flow.first
@@ -101,17 +102,7 @@ class RustMatrixAuthenticationService(
         runCatchingExceptions {
             val sessionData = sessionStore.getSession(sessionId.value)
             if (sessionData != null) {
-                if (sessionData.isTokenValid) {
-                    // Use the sessionData.passphrase, which can be null for a previously created session
-                    if (sessionData.passphrase == null) {
-                        Timber.w("Restoring a session without a passphrase")
-                    } else {
-                        Timber.w("Restoring a session with a passphrase")
-                    }
-                    rustMatrixClientFactory.create(sessionData)
-                } else {
-                    throw SessionRestorationException.InvalidToken()
-                }
+                restoreSession(sessionData)
             } else {
                 throw SessionRestorationException.MissingSession(sessionId)
             }
@@ -120,6 +111,19 @@ class RustMatrixAuthenticationService(
         }
     }
 
+    private suspend fun restoreSession(sessionData: SessionData): MatrixClient {
+        if (!sessionData.isTokenValid) {
+            throw SessionRestorationException.InvalidToken()
+        }
+        // Use the sessionData.passphrase, which can be null for a previously created session
+        if (sessionData.passphrase == null) {
+            Timber.w("Restoring a session without a passphrase")
+        } else {
+            Timber.w("Restoring a session with a passphrase")
+        }
+        return rustMatrixClientFactory.create(sessionData)
+    }
+
     private fun getDatabaseKey(): ClientSecret {
         Timber.d("New sessions will be encrypted with a raw key")
         return secretGenerator.generateKey()
@@ -137,7 +141,7 @@ class RustMatrixAuthenticationService(
 
                 client.homeserverLoginDetails().map()
             }.onFailure {
-                clear(destroyClient = true)
+                clear()
             }.mapFailure { failure ->
                 Timber.e(failure, "Failed to set homeserver to $homeserver")
                 failure.mapAuthenticationException()
@@ -165,18 +169,7 @@ class RustMatrixAuthenticationService(
                         passphrase = pendingKey.formattedAsString(),
                         sessionPaths = currentSessionPaths,
                     )
-                val matrixClient = rustMatrixClientFactory.create(client, sessionData, isMessageSearchAvailable())
-
-                // Apply enterprise hooks to the newly created client as soon as possible
-                clientEnterpriseHook(matrixClient)
-
-                newMatrixClientObservers.forEach { it.invoke(matrixClient) }
-                sessionStore.addSession(sessionData)
-
-                // Clean up the strong reference held here since it's no longer necessary
-                clear(destroyClient = false)
-
-                SessionId(sessionData.userId)
+                finalizeClientCreation(sessionData)
             }.mapFailure { failure ->
                 Timber.e(failure, "Failed to login")
                 failure.mapAuthenticationException()
@@ -305,20 +298,7 @@ class RustMatrixAuthenticationService(
                     passphrase = pendingKey.formattedAsString(),
                     sessionPaths = currentSessionPaths,
                 )
-                val matrixClient = rustMatrixClientFactory.create(client, sessionData, isMessageSearchAvailable())
-
-                // Apply enterprise hooks to the newly created client as soon as possible
-                clientEnterpriseHook(matrixClient)
-
-                matrixClient.waitForKnownVerificationState()
-
-                newMatrixClientObservers.forEach { it.invoke(matrixClient) }
-                sessionStore.addSession(sessionData)
-
-                // Clean up the strong reference held here since it's no longer necessary
-                clear(destroyClient = false)
-
-                SessionId(sessionData.userId)
+                finalizeClientCreation(sessionData)
             }.mapFailure { failure ->
                 Timber.e(failure, "Failed to login with OAuth")
                 failure.mapAuthenticationException()
@@ -357,6 +337,7 @@ class RustMatrixAuthenticationService(
                     sessionPaths = emptySessionPaths,
                     qrCodeData = sdkQrCodeLoginData,
                 )
+                currentClient = client
                 client.newLoginWithQrCodeHandler(
                     oauthConfiguration = oAuthConfiguration,
                 ).use {
@@ -374,32 +355,44 @@ class RustMatrixAuthenticationService(
                         passphrase = pendingKey.formattedAsString(),
                         sessionPaths = emptySessionPaths,
                     )
-                val matrixClient = rustMatrixClientFactory.create(client, sessionData, isMessageSearchAvailable())
-
-                // Apply enterprise hooks to the newly created client as soon as possible
-                clientEnterpriseHook(matrixClient)
-
-                newMatrixClientObservers.forEach { it.invoke(matrixClient) }
-                sessionStore.addSession(sessionData)
-
-                // Clean up the strong reference held here since it's no longer necessary
-                clear(destroyClient = false)
-
-                SessionId(sessionData.userId)
+                finalizeClientCreation(sessionData)
             }.mapFailure {
                 when (it) {
                     is QrCodeDecodeException -> QrErrorMapper.map(it)
                     is HumanQrLoginException -> QrErrorMapper.map(it)
                     else -> it
                 }
             }.onFailure { throwable ->
+                // A QR code login always builds its own client, so it can be disposed of on failure.
+                clear()
                 if (throwable is CancellationException) {
                     throw throwable
                 }
                 Timber.e(throwable, "Failed to login with QR code")
             }
         }
 
+    private suspend fun finalizeClientCreation(sessionData: SessionData): SessionId {
+        // Close the client which was used to perform the login before creating the final client.
+        // Both use the same session paths, so their SQLite stores must never be opened at the same time.
+        clear()
+
+        val matrixClient = restoreSession(sessionData)
+        // Apply enterprise hooks to the newly created client as soon as possible
+        clientEnterpriseHook(matrixClient)
+
+        matrixClient.waitForKnownVerificationState()
+
+        newMatrixClientObservers.forEach { it.invoke(matrixClient) }
+        sessionStore.addSession(sessionData)
+
+        // The session paths now hold the data of the account which has just been logged in, so forget
+        // them: they must not be deleted by the rotateSessionPath() of the next login attempt.
+        sessionPaths = null
+
+        return SessionId(sessionData.userId)
+    }
+
     private suspend fun makeClient(
         sessionPaths: SessionPaths,
         config: suspend ClientBuilder.() -> ClientBuilder,
@@ -449,10 +442,8 @@ class RustMatrixAuthenticationService(
             .build()
     }
 
-    private fun clear(destroyClient: Boolean) {
-        if (destroyClient) {
-            currentClient?.close()
-        }
+    private fun clear() {
+        currentClient?.close()
         currentClient = null
     }
 

libraries/matrix/impl/src/test/kotlin/io/element/android/libraries/matrix/impl/RustMatrixClientFactoryTest.kt

@@ -9,11 +9,13 @@
 package io.element.android.libraries.matrix.impl
 
 import com.google.common.truth.Truth.assertThat
+import io.element.android.features.enterprise.test.FakeClientBuilderEnterpriseHook
 import io.element.android.libraries.featureflag.test.FakeFeatureFlagService
 import io.element.android.libraries.matrix.api.core.SessionId
 import io.element.android.libraries.matrix.impl.auth.FakeProxyProvider
 import io.element.android.libraries.matrix.impl.room.FakeTimelineEventFilterFactory
 import io.element.android.libraries.matrix.impl.storage.FakeSqliteStoreBuilderProvider
+import io.element.android.libraries.matrix.impl.storage.SqliteStoreBuilderProvider
 import io.element.android.libraries.network.useragent.SimpleUserAgentProvider
 import io.element.android.libraries.sessionstorage.api.SessionStore
 import io.element.android.libraries.sessionstorage.test.InMemorySessionStore
@@ -51,6 +53,7 @@ fun TestScope.createRustMatrixClientFactory(
     ),
     clientBuilderProvider: ClientBuilderProvider = FakeClientBuilderProvider(),
     workManagerScheduler: FakeWorkManagerScheduler = FakeWorkManagerScheduler(),
+    sqliteStoreBuilderProvider: SqliteStoreBuilderProvider = FakeSqliteStoreBuilderProvider(),
 ) = RustMatrixClientFactory(
     cacheDirectory = cacheDirectory,
     appCoroutineScope = backgroundScope,
@@ -63,7 +66,7 @@ fun TestScope.createRustMatrixClientFactory(
     featureFlagService = FakeFeatureFlagService(),
     timelineEventFilterFactory = FakeTimelineEventFilterFactory(),
     clientBuilderProvider = clientBuilderProvider,
-    sqliteStoreBuilderProvider = FakeSqliteStoreBuilderProvider(),
+    sqliteStoreBuilderProvider = sqliteStoreBuilderProvider,
     workManagerScheduler = workManagerScheduler,
-    clientBuilderEnterpriseHook = { builder, _ -> builder },
+    clientBuilderEnterpriseHook = FakeClientBuilderEnterpriseHook(),
 )

libraries/matrix/impl/src/test/kotlin/io/element/android/libraries/matrix/impl/RustTemporaryMatrixClientFactoryTest.kt

@@ -8,6 +8,7 @@
 package io.element.android.libraries.matrix.impl
 
 import com.google.common.truth.Truth.assertThat
+import io.element.android.features.enterprise.test.FakeClientBuilderEnterpriseHook
 import io.element.android.libraries.featureflag.test.FakeFeatureFlagService
 import io.element.android.libraries.matrix.impl.auth.FakeProxyProvider
 import io.element.android.libraries.matrix.impl.paths.SessionPathsFactory
@@ -77,6 +78,6 @@ class RustTemporaryMatrixClientFactoryTest {
         clientBuilderProvider = clientBuilderProvider,
         sqliteStoreBuilderProvider = FakeSqliteStoreBuilderProvider(),
         workManagerScheduler = workManagerScheduler,
-        clientBuilderEnterpriseHook = { builder, _ -> builder },
+        clientBuilderEnterpriseHook = FakeClientBuilderEnterpriseHook(),
     )
 }