Avoid restoring several MatrixClient

libraries/matrix/api/src/main/kotlin/io/element/android/libraries/matrix/api/auth/MatrixAuthenticationService.kt

@@ -31,6 +31,7 @@ interface MatrixAuthenticationService {
     /**
      * Restore a session from a [sessionId].
      * Do not restore anything it the access token is not valid anymore.
+     * Generally this method should not be used directly, prefer using [MatrixClientProvider.getOrRestore] instead.
      */
     suspend fun restoreSession(sessionId: SessionId): Result<MatrixClient>
     fun getHomeserverDetails(): StateFlow<MatrixHomeServerDetails?>

libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseNewTokenHandler.kt

@@ -20,7 +20,7 @@ import com.squareup.anvil.annotations.ContributesBinding
 import io.element.android.libraries.core.extensions.flatMap
 import io.element.android.libraries.core.log.logger.LoggerTag
 import io.element.android.libraries.di.AppScope
-import io.element.android.libraries.matrix.api.auth.MatrixAuthenticationService
+import io.element.android.libraries.matrix.api.MatrixClientProvider
 import io.element.android.libraries.matrix.api.core.SessionId
 import io.element.android.libraries.pushproviders.api.PusherSubscriber
 import io.element.android.libraries.pushstore.api.UserPushStoreFactory
@@ -43,7 +43,7 @@ class DefaultFirebaseNewTokenHandler @Inject constructor(
     private val pusherSubscriber: PusherSubscriber,
     private val sessionStore: SessionStore,
     private val userPushStoreFactory: UserPushStoreFactory,
-    private val matrixAuthenticationService: MatrixAuthenticationService,
+    private val matrixClientProvider: MatrixClientProvider,
     private val firebaseStore: FirebaseStore,
 ) : FirebaseNewTokenHandler {
     override suspend fun handle(firebaseToken: String) {
@@ -54,8 +54,8 @@ class DefaultFirebaseNewTokenHandler @Inject constructor(
             .forEach { sessionId ->
                 val userDataStore = userPushStoreFactory.getOrCreate(sessionId)
                 if (userDataStore.getPushProviderName() == FirebaseConfig.NAME) {
-                    matrixAuthenticationService
+                    matrixClientProvider
-                        .restoreSession(sessionId)
+                        .getOrRestore(sessionId)
                         .onFailure {
                             Timber.tag(loggerTag.value).e(it, "Failed to restore session $sessionId")
                         }

libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/DefaultFirebaseNewTokenHandlerTest.kt

@@ -18,13 +18,13 @@ package io.element.android.libraries.pushproviders.firebase
 
 import com.google.common.truth.Truth.assertThat
 import io.element.android.libraries.matrix.api.MatrixClient
-import io.element.android.libraries.matrix.api.auth.MatrixAuthenticationService
+import io.element.android.libraries.matrix.api.MatrixClientProvider
 import io.element.android.libraries.matrix.test.AN_EXCEPTION
 import io.element.android.libraries.matrix.test.A_USER_ID
 import io.element.android.libraries.matrix.test.A_USER_ID_2
 import io.element.android.libraries.matrix.test.A_USER_ID_3
 import io.element.android.libraries.matrix.test.FakeMatrixClient
-import io.element.android.libraries.matrix.test.auth.FakeMatrixAuthenticationService
+import io.element.android.libraries.matrix.test.FakeMatrixClientProvider
 import io.element.android.libraries.push.test.FakePusherSubscriber
 import io.element.android.libraries.pushproviders.api.PusherSubscriber
 import io.element.android.libraries.pushstore.api.UserPushStoreFactory
@@ -64,16 +64,14 @@ class DefaultFirebaseNewTokenHandlerTest {
                 storeData(aSessionData(A_USER_ID_2))
                 storeData(aSessionData(A_USER_ID_3))
             },
-            matrixAuthenticationService = FakeMatrixAuthenticationService(
+            matrixClientProvider = FakeMatrixClientProvider { sessionId ->
-                matrixClientResult = { sessionId ->
+                when (sessionId) {
-                    when (sessionId) {
+                    A_USER_ID -> Result.success(aMatrixClient1)
-                        A_USER_ID -> Result.success(aMatrixClient1)
+                    A_USER_ID_2 -> Result.success(aMatrixClient2)
-                        A_USER_ID_2 -> Result.success(aMatrixClient2)
+                    A_USER_ID_3 -> Result.success(aMatrixClient3)
-                        A_USER_ID_3 -> Result.success(aMatrixClient3)
+                    else -> Result.failure(IllegalStateException())
-                        else -> Result.failure(IllegalStateException())
-                    }
                 }
-            ),
+            },
             userPushStoreFactory = FakeUserPushStoreFactory(
                 userPushStore = { sessionId ->
                     when (sessionId) {
@@ -103,11 +101,9 @@ class DefaultFirebaseNewTokenHandlerTest {
             sessionStore = InMemoryMultiSessionsStore().apply {
                 storeData(aSessionData(A_USER_ID))
             },
-            matrixAuthenticationService = FakeMatrixAuthenticationService(
+            matrixClientProvider = FakeMatrixClientProvider {
-                matrixClientResult = { _ ->
+                Result.failure(IllegalStateException())
-                    Result.failure(IllegalStateException())
+            },
-                }
-            ),
             userPushStoreFactory = FakeUserPushStoreFactory(
                 userPushStore = { _ ->
                     FakeUserPushStore(pushProviderName = FirebaseConfig.NAME)
@@ -129,11 +125,9 @@ class DefaultFirebaseNewTokenHandlerTest {
             sessionStore = InMemoryMultiSessionsStore().apply {
                 storeData(aSessionData(A_USER_ID))
             },
-            matrixAuthenticationService = FakeMatrixAuthenticationService(
+            matrixClientProvider = FakeMatrixClientProvider {
-                matrixClientResult = { _ ->
+                Result.success(aMatrixClient1)
-                    Result.success(aMatrixClient1)
+            },
-                }
-            ),
             userPushStoreFactory = FakeUserPushStoreFactory(
                 userPushStore = { _ ->
                     FakeUserPushStore(pushProviderName = FirebaseConfig.NAME)
@@ -152,14 +146,14 @@ class DefaultFirebaseNewTokenHandlerTest {
         pusherSubscriber: PusherSubscriber = FakePusherSubscriber(),
         sessionStore: SessionStore = InMemorySessionStore(),
         userPushStoreFactory: UserPushStoreFactory = FakeUserPushStoreFactory(),
-        matrixAuthenticationService: MatrixAuthenticationService = FakeMatrixAuthenticationService(),
+        matrixClientProvider: MatrixClientProvider = FakeMatrixClientProvider(),
         firebaseStore: FirebaseStore = InMemoryFirebaseStore(),
     ): FirebaseNewTokenHandler {
         return DefaultFirebaseNewTokenHandler(
             pusherSubscriber = pusherSubscriber,
             sessionStore = sessionStore,
             userPushStoreFactory = userPushStoreFactory,
-            matrixAuthenticationService = matrixAuthenticationService,
+            matrixClientProvider = matrixClientProvider,
             firebaseStore = firebaseStore
         )
     }

libraries/pushproviders/unifiedpush/src/main/kotlin/io/element/android/libraries/pushproviders/unifiedpush/UnifiedPushNewGatewayHandler.kt

@@ -20,7 +20,7 @@ import com.squareup.anvil.annotations.ContributesBinding
 import io.element.android.libraries.core.extensions.flatMap
 import io.element.android.libraries.core.log.logger.LoggerTag
 import io.element.android.libraries.di.AppScope
-import io.element.android.libraries.matrix.api.auth.MatrixAuthenticationService
+import io.element.android.libraries.matrix.api.MatrixClientProvider
 import io.element.android.libraries.pushproviders.api.PusherSubscriber
 import io.element.android.libraries.pushstore.api.UserPushStoreFactory
 import io.element.android.libraries.pushstore.api.clientsecret.PushClientSecret
@@ -41,7 +41,7 @@ class DefaultUnifiedPushNewGatewayHandler @Inject constructor(
     private val pusherSubscriber: PusherSubscriber,
     private val userPushStoreFactory: UserPushStoreFactory,
     private val pushClientSecret: PushClientSecret,
-    private val matrixAuthenticationService: MatrixAuthenticationService,
+    private val matrixClientProvider: MatrixClientProvider,
 ) : UnifiedPushNewGatewayHandler {
     override suspend fun handle(endpoint: String, pushGateway: String, clientSecret: String): Result<Unit> {
         // Register the pusher for the session with this client secret, if is it using UnifiedPush.
@@ -52,8 +52,8 @@ class DefaultUnifiedPushNewGatewayHandler @Inject constructor(
         }
         val userDataStore = userPushStoreFactory.getOrCreate(userId)
         return if (userDataStore.getPushProviderName() == UnifiedPushConfig.NAME) {
-            matrixAuthenticationService
+            matrixClientProvider
-                .restoreSession(userId)
+                .getOrRestore(userId)
                 .flatMap { client ->
                     pusherSubscriber.registerPusher(client, endpoint, pushGateway)
                 }

libraries/pushproviders/unifiedpush/src/test/kotlin/io/element/android/libraries/pushproviders/unifiedpush/DefaultUnifiedPushNewGatewayHandlerTest.kt

@@ -18,11 +18,11 @@ package io.element.android.libraries.pushproviders.unifiedpush
 
 import com.google.common.truth.Truth.assertThat
 import io.element.android.libraries.matrix.api.MatrixClient
-import io.element.android.libraries.matrix.api.auth.MatrixAuthenticationService
+import io.element.android.libraries.matrix.api.MatrixClientProvider
 import io.element.android.libraries.matrix.test.A_SECRET
 import io.element.android.libraries.matrix.test.A_USER_ID
 import io.element.android.libraries.matrix.test.FakeMatrixClient
-import io.element.android.libraries.matrix.test.auth.FakeMatrixAuthenticationService
+import io.element.android.libraries.matrix.test.FakeMatrixClientProvider
 import io.element.android.libraries.push.test.FakePusherSubscriber
 import io.element.android.libraries.pushproviders.api.PusherSubscriber
 import io.element.android.libraries.pushstore.api.UserPushStoreFactory
@@ -86,7 +86,7 @@ class DefaultUnifiedPushNewGatewayHandlerTest {
             pusherSubscriber = FakePusherSubscriber(
                 registerPusherResult = { _, _, _ -> Result.failure(IllegalStateException("an error")) }
             ),
-            matrixAuthenticationService = FakeMatrixAuthenticationService(matrixClientResult = { Result.success(aMatrixClient) }),
+            matrixClientProvider = FakeMatrixClientProvider { Result.success(aMatrixClient) },
         )
         val result = defaultUnifiedPushNewGatewayHandler.handle(
             endpoint = "aEndpoint",
@@ -114,7 +114,7 @@ class DefaultUnifiedPushNewGatewayHandlerTest {
             pusherSubscriber = FakePusherSubscriber(
                 registerPusherResult = lambda
             ),
-            matrixAuthenticationService = FakeMatrixAuthenticationService(matrixClientResult = { Result.success(aMatrixClient) }),
+            matrixClientProvider = FakeMatrixClientProvider { Result.success(aMatrixClient) },
         )
         val result = defaultUnifiedPushNewGatewayHandler.handle(
             endpoint = "aEndpoint",
@@ -131,13 +131,13 @@ class DefaultUnifiedPushNewGatewayHandlerTest {
         pusherSubscriber: PusherSubscriber = FakePusherSubscriber(),
         userPushStoreFactory: UserPushStoreFactory = FakeUserPushStoreFactory(),
         pushClientSecret: PushClientSecret = FakePushClientSecret(),
-        matrixAuthenticationService: MatrixAuthenticationService = FakeMatrixAuthenticationService()
+        matrixClientProvider: MatrixClientProvider = FakeMatrixClientProvider()
     ): DefaultUnifiedPushNewGatewayHandler {
         return DefaultUnifiedPushNewGatewayHandler(
             pusherSubscriber = pusherSubscriber,
             userPushStoreFactory = userPushStoreFactory,
             pushClientSecret = pushClientSecret,
-            matrixAuthenticationService = matrixAuthenticationService
+            matrixClientProvider = matrixClientProvider,
         )
     }
 }