Merge pull request #3755 from element-hq/feature/bma/rotateFirebaseToken

Rotate firebase token in case of error

libraries/push/impl/src/main/kotlin/io/element/android/libraries/push/impl/troubleshoot/PushLoopbackTest.kt

@@ -52,9 +52,10 @@ class PushLoopbackTest @Inject constructor(
         val testPushResult = try {
             pushService.testPush()
         } catch (pusherRejected: PushGatewayFailure.PusherRejected) {
+            val hasQuickFix = pushService.getCurrentPushProvider()?.canRotateToken() == true
             delegate.updateState(
                 description = stringProvider.getString(R.string.troubleshoot_notifications_test_push_loop_back_failure_1),
-                status = NotificationTroubleshootTestState.Status.Failure(false)
+                status = NotificationTroubleshootTestState.Status.Failure(hasQuickFix)
             )
             job.cancel()
             return
@@ -96,5 +97,11 @@ class PushLoopbackTest @Inject constructor(
         )
     }
 
+    override suspend fun quickFix(coroutineScope: CoroutineScope) {
+        delegate.start()
+        pushService.getCurrentPushProvider()?.rotateToken()
+        run(coroutineScope)
+    }
+
     override suspend fun reset() = delegate.reset()
 }

libraries/push/impl/src/test/kotlin/io/element/android/libraries/push/impl/troubleshoot/PushLoopbackTestTest.kt

@@ -13,9 +13,11 @@ import io.element.android.libraries.matrix.test.AN_EXCEPTION
 import io.element.android.libraries.matrix.test.A_FAILURE_REASON
 import io.element.android.libraries.push.api.gateway.PushGatewayFailure
 import io.element.android.libraries.push.test.FakePushService
+import io.element.android.libraries.pushproviders.test.FakePushProvider
 import io.element.android.libraries.troubleshoot.api.test.NotificationTroubleshootTestState
 import io.element.android.services.toolbox.test.strings.FakeStringProvider
 import io.element.android.services.toolbox.test.systemclock.FakeSystemClock
+import io.element.android.tests.testutils.lambda.lambdaRecorder
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
@@ -67,6 +69,41 @@ class PushLoopbackTestTest {
         }
     }
 
+    @Test
+    fun `test PushLoopbackTest PusherRejected error with quick fix`() = runTest {
+        val diagnosticPushHandler = DiagnosticPushHandler()
+        val rotateTokenLambda = lambdaRecorder<Result<Unit>> { Result.success(Unit) }
+        val sut = PushLoopbackTest(
+            pushService = FakePushService(
+                testPushBlock = {
+                    throw PushGatewayFailure.PusherRejected()
+                },
+                currentPushProvider = {
+                    FakePushProvider(
+                        canRotateTokenResult = { true },
+                        rotateTokenLambda = rotateTokenLambda,
+                    )
+                }
+            ),
+            diagnosticPushHandler = diagnosticPushHandler,
+            clock = FakeSystemClock(),
+            stringProvider = FakeStringProvider(),
+        )
+        launch {
+            sut.run(this)
+        }
+        sut.state.test {
+            assertThat(awaitItem().status).isEqualTo(NotificationTroubleshootTestState.Status.Idle(true))
+            assertThat(awaitItem().status).isEqualTo(NotificationTroubleshootTestState.Status.InProgress)
+            val lastItem = awaitItem()
+            assertThat(lastItem.status).isEqualTo(NotificationTroubleshootTestState.Status.Failure(true))
+            sut.quickFix(this)
+            assertThat(awaitItem().status).isEqualTo(NotificationTroubleshootTestState.Status.InProgress)
+            assertThat(awaitItem().status).isEqualTo(NotificationTroubleshootTestState.Status.Failure(true))
+            rotateTokenLambda.assertions().isCalledOnce()
+        }
+    }
+
     @Test
     fun `test PushLoopbackTest setup error`() = runTest {
         val diagnosticPushHandler = DiagnosticPushHandler()

libraries/pushproviders/api/src/main/kotlin/io/element/android/libraries/pushproviders/api/PushProvider.kt

@@ -44,4 +44,10 @@ interface PushProvider {
     suspend fun unregister(matrixClient: MatrixClient): Result<Unit>
 
     suspend fun getCurrentUserPushConfig(): CurrentUserPushConfig?
+
+    fun canRotateToken(): Boolean
+
+    suspend fun rotateToken(): Result<Unit> {
+        error("rotateToken() not implemented, you need to override this method in your implementation")
+    }
 }

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

@@ -25,6 +25,7 @@ class FirebasePushProvider @Inject constructor(
     private val firebaseStore: FirebaseStore,
     private val pusherSubscriber: PusherSubscriber,
     private val isPlayServiceAvailable: IsPlayServiceAvailable,
+    private val firebaseTokenRotator: FirebaseTokenRotator,
 ) : PushProvider {
     override val index = FirebaseConfig.INDEX
     override val name = FirebaseConfig.NAME
@@ -71,6 +72,12 @@ class FirebasePushProvider @Inject constructor(
         }
     }
 
+    override fun canRotateToken(): Boolean = true
+
+    override suspend fun rotateToken(): Result<Unit> {
+        return firebaseTokenRotator.rotate()
+    }
+
     companion object {
         private val firebaseDistributor = Distributor("Firebase", "Firebase")
     }

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

@@ -11,6 +11,10 @@ import android.content.SharedPreferences
 import androidx.core.content.edit
 import com.squareup.anvil.annotations.ContributesBinding
 import io.element.android.libraries.di.AppScope
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.onCompletion
+import kotlinx.coroutines.flow.onStart
 import javax.inject.Inject
 
 /**
@@ -18,6 +22,7 @@ import javax.inject.Inject
  */
 interface FirebaseStore {
     fun getFcmToken(): String?
+    fun fcmTokenFlow(): Flow<String?>
     fun storeFcmToken(token: String?)
 }
 
@@ -29,6 +34,22 @@ class SharedPreferencesFirebaseStore @Inject constructor(
         return sharedPreferences.getString(PREFS_KEY_FCM_TOKEN, null)
     }
 
+    override fun fcmTokenFlow(): Flow<String?> {
+        val flow = MutableStateFlow(getFcmToken())
+        val listener = SharedPreferences.OnSharedPreferenceChangeListener { _, k ->
+            if (k == PREFS_KEY_FCM_TOKEN) {
+                try {
+                    flow.value = getFcmToken()
+                } catch (e: Exception) {
+                    flow.value = null
+                }
+            }
+        }
+        return flow
+            .onStart { sharedPreferences.registerOnSharedPreferenceChangeListener(listener) }
+            .onCompletion { sharedPreferences.unregisterOnSharedPreferenceChangeListener(listener) }
+    }
+
     override fun storeFcmToken(token: String?) {
         sharedPreferences.edit {
             putString(PREFS_KEY_FCM_TOKEN, token)

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

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2023, 2024 New Vector Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only
+ * Please see LICENSE in the repository root for full details.
+ */
+
+package io.element.android.libraries.pushproviders.firebase
+
+import com.google.firebase.messaging.FirebaseMessaging
+import com.squareup.anvil.annotations.ContributesBinding
+import io.element.android.libraries.di.AppScope
+import timber.log.Timber
+import javax.inject.Inject
+import kotlin.coroutines.resume
+import kotlin.coroutines.resumeWithException
+import kotlin.coroutines.suspendCoroutine
+
+interface FirebaseTokenDeleter {
+    /**
+     * Deletes the current Firebase token.
+     */
+    suspend fun delete()
+}
+
+@ContributesBinding(AppScope::class)
+class DefaultFirebaseTokenDeleter @Inject constructor(
+    private val isPlayServiceAvailable: IsPlayServiceAvailable,
+) : FirebaseTokenDeleter {
+    override suspend fun delete() {
+        // 'app should always check the device for a compatible Google Play services APK before accessing Google Play services features'
+        isPlayServiceAvailable.checkAvailableOrThrow()
+        suspendCoroutine { continuation ->
+            try {
+                FirebaseMessaging.getInstance().deleteToken()
+                    .addOnSuccessListener {
+                        continuation.resume(Unit)
+                    }
+                    .addOnFailureListener { e ->
+                        Timber.e(e, "## deleteFirebaseToken() : failed")
+                        continuation.resumeWithException(e)
+                    }
+            } catch (e: Throwable) {
+                Timber.e(e, "## deleteFirebaseToken() : failed")
+                continuation.resumeWithException(e)
+            }
+        }
+    }
+}

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

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2023, 2024 New Vector Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only
+ * Please see LICENSE in the repository root for full details.
+ */
+
+package io.element.android.libraries.pushproviders.firebase
+
+import com.google.firebase.messaging.FirebaseMessaging
+import com.squareup.anvil.annotations.ContributesBinding
+import io.element.android.libraries.di.AppScope
+import timber.log.Timber
+import javax.inject.Inject
+import kotlin.coroutines.resume
+import kotlin.coroutines.resumeWithException
+import kotlin.coroutines.suspendCoroutine
+
+interface FirebaseTokenGetter {
+    /**
+     * Read the current Firebase token from FirebaseMessaging.
+     * If the token does not exist, it will be generated.
+     */
+    suspend fun get(): String
+}
+
+@ContributesBinding(AppScope::class)
+class DefaultFirebaseTokenGetter @Inject constructor(
+    private val isPlayServiceAvailable: IsPlayServiceAvailable,
+) : FirebaseTokenGetter {
+    override suspend fun get(): String {
+        // 'app should always check the device for a compatible Google Play services APK before accessing Google Play services features'
+        isPlayServiceAvailable.checkAvailableOrThrow()
+        return suspendCoroutine { continuation ->
+            try {
+                FirebaseMessaging.getInstance().token
+                    .addOnSuccessListener { token ->
+                        continuation.resume(token)
+                    }
+                    .addOnFailureListener { e ->
+                        Timber.e(e, "## retrievedFirebaseToken() : failed")
+                        continuation.resumeWithException(e)
+                    }
+            } catch (e: Throwable) {
+                Timber.e(e, "## retrievedFirebaseToken() : failed")
+                continuation.resumeWithException(e)
+            }
+        }
+    }
+}

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

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2024 New Vector Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only
+ * Please see LICENSE in the repository root for full details.
+ */
+
+package io.element.android.libraries.pushproviders.firebase
+
+import com.squareup.anvil.annotations.ContributesBinding
+import io.element.android.libraries.di.AppScope
+import javax.inject.Inject
+
+interface FirebaseTokenRotator {
+    suspend fun rotate(): Result<Unit>
+}
+
+/**
+ * This class delete the Firebase token and generate a new one.
+ */
+@ContributesBinding(AppScope::class)
+class DefaultFirebaseTokenRotator @Inject constructor(
+    private val firebaseTokenDeleter: FirebaseTokenDeleter,
+    private val firebaseTokenGetter: FirebaseTokenGetter,
+) : FirebaseTokenRotator {
+    override suspend fun rotate(): Result<Unit> {
+        return runCatching {
+            firebaseTokenDeleter.delete()
+            firebaseTokenGetter.get()
+        }
+    }
+}

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

@@ -7,14 +7,9 @@
 
 package io.element.android.libraries.pushproviders.firebase
 
-import com.google.firebase.messaging.FirebaseMessaging
 import com.squareup.anvil.annotations.ContributesBinding
 import io.element.android.libraries.di.AppScope
-import timber.log.Timber
 import javax.inject.Inject
-import kotlin.coroutines.resume
-import kotlin.coroutines.resumeWithException
-import kotlin.coroutines.suspendCoroutine
 
 interface FirebaseTroubleshooter {
     suspend fun troubleshoot(): Result<Unit>
@@ -26,37 +21,12 @@ interface FirebaseTroubleshooter {
 @ContributesBinding(AppScope::class)
 class DefaultFirebaseTroubleshooter @Inject constructor(
     private val newTokenHandler: FirebaseNewTokenHandler,
-    private val isPlayServiceAvailable: IsPlayServiceAvailable,
+    private val firebaseTokenGetter: FirebaseTokenGetter,
 ) : FirebaseTroubleshooter {
     override suspend fun troubleshoot(): Result<Unit> {
         return runCatching {
-            val token = retrievedFirebaseToken()
+            val token = firebaseTokenGetter.get()
             newTokenHandler.handle(token)
         }
     }
-
-    private suspend fun retrievedFirebaseToken(): String {
-        return suspendCoroutine { continuation ->
-            // 'app should always check the device for a compatible Google Play services APK before accessing Google Play services features'
-            if (isPlayServiceAvailable.isAvailable()) {
-                try {
-                    FirebaseMessaging.getInstance().token
-                        .addOnSuccessListener { token ->
-                            continuation.resume(token)
-                        }
-                        .addOnFailureListener { e ->
-                            Timber.e(e, "## retrievedFirebaseToken() : failed")
-                            continuation.resumeWithException(e)
-                        }
-                } catch (e: Throwable) {
-                    Timber.e(e, "## retrievedFirebaseToken() : failed")
-                    continuation.resumeWithException(e)
-                }
-            } else {
-                val e = Exception("No valid Google Play Services found. Cannot use FCM.")
-                Timber.e(e)
-                continuation.resumeWithException(e)
-            }
-        }
-    }
 }

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

@@ -20,6 +20,12 @@ interface IsPlayServiceAvailable {
     fun isAvailable(): Boolean
 }
 
+fun IsPlayServiceAvailable.checkAvailableOrThrow() {
+    if (!isAvailable()) {
+        throw Exception("No valid Google Play Services found. Cannot use FCM.").also(Timber::e)
+    }
+}
+
 @ContributesBinding(AppScope::class)
 class DefaultIsPlayServiceAvailable @Inject constructor(
     @ApplicationContext private val context: Context,

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

@@ -19,7 +19,10 @@ import io.element.android.libraries.troubleshoot.api.test.NotificationTroublesho
 import io.element.android.libraries.troubleshoot.api.test.TestFilterData
 import io.element.android.services.toolbox.api.strings.StringProvider
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
 import javax.inject.Inject
 
 @ContributesMultibinding(AppScope::class)
@@ -41,23 +44,28 @@ class FirebaseTokenTest @Inject constructor(
         return data.currentPushProviderName == FirebaseConfig.NAME
     }
 
+    private var currentJob: Job? = null
     override suspend fun run(coroutineScope: CoroutineScope) {
+        currentJob?.cancel()
         delegate.start()
-        val token = firebaseStore.getFcmToken()
+        currentJob = firebaseStore.fcmTokenFlow()
-        if (token != null) {
+            .onEach { token ->
-            delegate.updateState(
+                if (token != null) {
-                description = stringProvider.getString(
+                    delegate.updateState(
-                    R.string.troubleshoot_notifications_test_firebase_token_success,
+                        description = stringProvider.getString(
-                    "${token.take(8)}*****"
+                            R.string.troubleshoot_notifications_test_firebase_token_success,
-                ),
+                            "*****${token.takeLast(8)}"
-                status = NotificationTroubleshootTestState.Status.Success
+                        ),
-            )
+                        status = NotificationTroubleshootTestState.Status.Success
-        } else {
+                    )
-            delegate.updateState(
+                } else {
-                description = stringProvider.getString(R.string.troubleshoot_notifications_test_firebase_token_failure),
+                    delegate.updateState(
-                status = NotificationTroubleshootTestState.Status.Failure(true)
+                        description = stringProvider.getString(R.string.troubleshoot_notifications_test_firebase_token_failure),
-            )
+                        status = NotificationTroubleshootTestState.Status.Failure(true)
-        }
+                    )
+                }
+            }
+            .launchIn(coroutineScope)
     }
 
     override suspend fun reset() = delegate.reset()

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

@@ -0,0 +1,18 @@
+/*
+ * Copyright 2024 New Vector Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only
+ * Please see LICENSE in the repository root for full details.
+ */
+
+package io.element.android.libraries.pushproviders.firebase
+
+import io.element.android.tests.testutils.lambda.lambdaError
+
+class FakeFirebaseTokenRotator(
+    private val rotateWithResult: () -> Result<Unit> = { lambdaError() }
+) : FirebaseTokenRotator {
+    override suspend fun rotate(): Result<Unit> {
+        return rotateWithResult()
+    }
+}

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

@@ -166,15 +166,27 @@ class FirebasePushProviderTest {
         assertThat(result).isEqualTo(CurrentUserPushConfig(FirebaseConfig.PUSHER_HTTP_URL, "aToken"))
     }
 
+    @Test
+    fun `rotateToken invokes the FirebaseTokenRotator`() = runTest {
+        val lambda = lambdaRecorder<Result<Unit>> { Result.success(Unit) }
+        val firebasePushProvider = createFirebasePushProvider(
+            firebaseTokenRotator = FakeFirebaseTokenRotator(lambda),
+        )
+        firebasePushProvider.rotateToken()
+        lambda.assertions().isCalledOnce()
+    }
+
     private fun createFirebasePushProvider(
         firebaseStore: FirebaseStore = InMemoryFirebaseStore(),
         pusherSubscriber: PusherSubscriber = FakePusherSubscriber(),
         isPlayServiceAvailable: IsPlayServiceAvailable = FakeIsPlayServiceAvailable(false),
+        firebaseTokenRotator: FirebaseTokenRotator = FakeFirebaseTokenRotator(),
     ): FirebasePushProvider {
         return FirebasePushProvider(
             firebaseStore = firebaseStore,
             pusherSubscriber = pusherSubscriber,
             isPlayServiceAvailable = isPlayServiceAvailable,
+            firebaseTokenRotator = firebaseTokenRotator,
         )
     }
 }

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

@@ -7,11 +7,16 @@
 
 package io.element.android.libraries.pushproviders.firebase
 
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.flowOf
+
 class InMemoryFirebaseStore(
     private var token: String? = null
 ) : FirebaseStore {
     override fun getFcmToken(): String? = token
 
+    override fun fcmTokenFlow(): Flow<String?> = flowOf(token)
+
     override fun storeFcmToken(token: String?) {
         this.token = token
     }

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

@@ -35,7 +35,7 @@ class FirebaseTokenTestTest {
             assertThat(awaitItem().status).isEqualTo(NotificationTroubleshootTestState.Status.InProgress)
             val lastItem = awaitItem()
             assertThat(lastItem.status).isEqualTo(NotificationTroubleshootTestState.Status.Success)
-            assertThat(lastItem.description).contains(FAKE_TOKEN.take(8))
+            assertThat(lastItem.description).contains(FAKE_TOKEN.takeLast(8))
             assertThat(lastItem.description).doesNotContain(FAKE_TOKEN)
         }
     }

libraries/pushproviders/test/src/main/kotlin/io/element/android/libraries/pushproviders/test/FakePushProvider.kt

@@ -21,6 +21,8 @@ class FakePushProvider(
     private val currentUserPushConfig: CurrentUserPushConfig? = null,
     private val registerWithResult: (MatrixClient, Distributor) -> Result<Unit> = { _, _ -> lambdaError() },
     private val unregisterWithResult: (MatrixClient) -> Result<Unit> = { lambdaError() },
+    private val canRotateTokenResult: () -> Boolean = { lambdaError() },
+    private val rotateTokenLambda: () -> Result<Unit> = { lambdaError() },
 ) : PushProvider {
     override fun getDistributors(): List<Distributor> = distributors
 
@@ -39,4 +41,12 @@ class FakePushProvider(
     override suspend fun getCurrentUserPushConfig(): CurrentUserPushConfig? {
         return currentUserPushConfig
     }
+
+    override fun canRotateToken(): Boolean {
+        return canRotateTokenResult()
+    }
+
+    override suspend fun rotateToken(): Result<Unit> {
+        return rotateTokenLambda()
+    }
 }

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

@@ -53,4 +53,6 @@ class UnifiedPushProvider @Inject constructor(
     override suspend fun getCurrentUserPushConfig(): CurrentUserPushConfig? {
         return unifiedPushCurrentUserPushConfigProvider.provide()
     }
+
+    override fun canRotateToken(): Boolean = false
 }