quick incremental backup support

matrix-sdk-android/src/main/java/org/matrix/android/sdk/api/session/crypto/model/MXUsersDevicesMap.kt

@@ -19,7 +19,7 @@ package org.matrix.android.sdk.api.session.crypto.model
 class MXUsersDevicesMap<E> {
 
     // A map of maps (userId -> (deviceId -> Object)).
-    val map = HashMap<String /* userId */, HashMap<String /* deviceId */, E>>()
+    val map = HashMap<String /* userId */, MutableMap<String /* deviceId */, E>>()
 
     /**
      * @return the user Ids
@@ -104,8 +104,8 @@ class MXUsersDevicesMap<E> {
         map.clear()
     }
 
-    fun join(other: Map<out String, HashMap<String, E>>) {
+    fun join(other: Map<out String, Map<String, E>>) {
-        map.putAll(other)
+        map.putAll(other.map { it.key to it.value.toMutableMap() })
     }
 
     /**

matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/PerSessionBackupQueryRateLimiter.kt

@@ -23,7 +23,6 @@ import org.matrix.android.sdk.api.logger.LoggerTag
 import org.matrix.android.sdk.api.session.crypto.keysbackup.KeysBackupService
 import org.matrix.android.sdk.api.session.crypto.keysbackup.KeysVersionResult
 import org.matrix.android.sdk.api.session.crypto.keysbackup.SavedKeyBackupKeyInfo
-import org.matrix.android.sdk.internal.crypto.store.IMXCryptoStore
 import org.matrix.android.sdk.internal.util.time.Clock
 import timber.log.Timber
 import javax.inject.Inject
@@ -39,7 +38,6 @@ private val loggerTag = LoggerTag("OutgoingGossipingRequestManager", LoggerTag.C
 internal class PerSessionBackupQueryRateLimiter @Inject constructor(
         private val coroutineDispatchers: MatrixCoroutineDispatchers,
         private val keysBackupService: Lazy<KeysBackupService>,
-        private val cryptoStore: IMXCryptoStore,
         private val clock: Clock,
 ) {
 
@@ -68,11 +66,11 @@ internal class PerSessionBackupQueryRateLimiter @Inject constructor(
     var backupWasCheckedFromServer: Boolean = false
     val now = clock.epochMillis()
 
-    fun refreshBackupInfoIfNeeded(force: Boolean = false) {
+    suspend fun refreshBackupInfoIfNeeded(force: Boolean = false) {
         if (backupWasCheckedFromServer && !force) return
         Timber.tag(loggerTag.value).v("Checking if can access a backup")
         backupWasCheckedFromServer = true
-        val knownBackupSecret = cryptoStore.getKeyBackupRecoveryKeyInfo()
+        val knownBackupSecret = keysBackupService.get().getKeyBackupRecoveryKeyInfo()
                 ?: return Unit.also {
                     Timber.tag(loggerTag.value).v("We don't have the backup secret!")
                 }

matrix-sdk-android/src/rustCrypto/java/org/matrix/android/sdk/internal/crypto/CryptoModule.kt

@@ -25,8 +25,10 @@ import kotlinx.coroutines.SupervisorJob
 import org.matrix.android.sdk.api.MatrixCoroutineDispatchers
 import org.matrix.android.sdk.api.session.crypto.CryptoService
 import org.matrix.android.sdk.api.session.crypto.crosssigning.CrossSigningService
+import org.matrix.android.sdk.api.session.crypto.keysbackup.KeysBackupService
 import org.matrix.android.sdk.api.session.crypto.verification.VerificationService
 import org.matrix.android.sdk.internal.crypto.api.CryptoApi
+import org.matrix.android.sdk.internal.crypto.keysbackup.RustKeyBackupService
 import org.matrix.android.sdk.internal.crypto.keysbackup.api.RoomKeysApi
 import org.matrix.android.sdk.internal.crypto.keysbackup.tasks.CreateKeysBackupVersionTask
 import org.matrix.android.sdk.internal.crypto.keysbackup.tasks.DefaultCreateKeysBackupVersionTask
@@ -248,4 +250,7 @@ internal abstract class CryptoModule {
 
     @Binds
     abstract fun bindSendEventTask(task: DefaultSendEventTask): SendEventTask
+
+    @Binds
+    abstract fun bindKeysBackupService(service: RustKeyBackupService): KeysBackupService
 }

matrix-sdk-android/src/rustCrypto/java/org/matrix/android/sdk/internal/crypto/RustCryptoService.kt

@@ -603,7 +603,7 @@ internal class RustCryptoService @Inject constructor(
                     val sessionId = content.sessionId
 
                     notifyRoomKeyReceived(roomId, sessionId)
-                    matrixConfiguration.cryptoAnalyticsPlugin?.onRoomKeyImported(sessionId, EventType.FORWARDED_ROOM_KEY)
+                    matrixConfiguration.cryptoAnalyticsPlugin?.onRoomKeyImported(sessionId, EventType.ROOM_KEY)
                 }
                 EventType.FORWARDED_ROOM_KEY -> {
                     val content = event.getClearContent().toModel<ForwardedRoomKeyContent>() ?: return@forEach

matrix-sdk-android/src/rustCrypto/java/org/matrix/android/sdk/internal/crypto/keysbackup/RustKeyBackupService.kt

@@ -20,6 +20,7 @@ import android.os.Handler
 import android.os.Looper
 import androidx.annotation.VisibleForTesting
 import androidx.annotation.WorkerThread
+import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.async
@@ -27,6 +28,7 @@ import kotlinx.coroutines.awaitAll
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
+import org.matrix.android.sdk.api.MatrixConfiguration
 import org.matrix.android.sdk.api.MatrixCoroutineDispatchers
 import org.matrix.android.sdk.api.crypto.MXCRYPTO_ALGORITHM_MEGOLM_BACKUP
 import org.matrix.android.sdk.api.extensions.tryOrNull
@@ -51,6 +53,7 @@ import org.matrix.android.sdk.api.session.crypto.model.ImportRoomKeysResult
 import org.matrix.android.sdk.internal.crypto.MegolmSessionData
 import org.matrix.android.sdk.internal.crypto.MegolmSessionImportManager
 import org.matrix.android.sdk.internal.crypto.OlmMachine
+import org.matrix.android.sdk.internal.crypto.PerSessionBackupQueryRateLimiter
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.SignalableMegolmBackupAuthData
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.CreateKeysBackupVersionBody
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.KeyBackupData
@@ -80,6 +83,8 @@ internal class RustKeyBackupService @Inject constructor(
         private val coroutineDispatchers: MatrixCoroutineDispatchers,
         private val megolmSessionImportManager: MegolmSessionImportManager,
         private val cryptoCoroutineScope: CoroutineScope,
+        private val matrixConfiguration: MatrixConfiguration,
+        private val backupQueryRateLimiter: dagger.Lazy<PerSessionBackupQueryRateLimiter>,
 ) : KeysBackupService {
     companion object {
         // Maximum delay in ms in {@link maybeBackupKeys}
@@ -94,9 +99,7 @@ internal class RustKeyBackupService @Inject constructor(
     override var keysBackupVersion: KeysVersionResult? = null
         private set
 
-//    private var backupAllGroupSessionsCallback: MatrixCallback<Unit>? = null
+    private val importScope = CoroutineScope(cryptoCoroutineScope.coroutineContext + SupervisorJob() + CoroutineName("backupImport"))
-
-    private val importScope = CoroutineScope(SupervisorJob() + coroutineDispatchers.main)
 
     private var keysBackupStateListener: KeysBackupStateListener? = null
 
@@ -381,20 +384,22 @@ internal class RustKeyBackupService @Inject constructor(
                 if (version != null) {
                     val key = BackupRecoveryKey.fromBase64(secret)
                     if (isValidRecoveryKey(key, version)) {
-                        trustKeysBackupVersion(version, true)
-                        // we don't want to wait for that
-                        importScope.launch {
-                            try {
-                                val importResult = restoreBackup(version, key, null, null, null)
-                                val recoveredKeys = importResult.successfullyNumberOfImportedKeys
-                                Timber.i("onSecretKeyGossip: Recovered keys $recoveredKeys out of ${importResult.totalNumberOfKeys}")
-                            } catch (failure: Throwable) {
-                                // fail silently..
-                                Timber.e(failure, "onSecretKeyGossip: Failed to import keys from backup")
-                            }
-                        }
                         // we can save, it's valid
                         saveBackupRecoveryKey(key, version.version)
+                        importScope.launch {
+                            backupQueryRateLimiter.get().refreshBackupInfoIfNeeded(true)
+                        }
+                        // we don't want to wait for that
+//                        importScope.launch {
+//                            try {
+//                                val importResult = restoreBackup(version, key, null, null, null)
+//                                val recoveredKeys = importResult.successfullyNumberOfImportedKeys
+//                                Timber.i("onSecretKeyGossip: Recovered keys $recoveredKeys out of ${importResult.totalNumberOfKeys}")
+//                            } catch (failure: Throwable) {
+//                                // fail silently..
+//                                Timber.e(failure, "onSecretKeyGossip: Failed to import keys from backup")
+//                            }
+//                        }
                     } else {
                         Timber.d("Invalid recovery key")
                     }
@@ -549,6 +554,10 @@ internal class RustKeyBackupService @Inject constructor(
             }
 
             val result = olmMachine.importDecryptedKeys(sessionsData, progressListener).also {
+                sessionsData.onEach { sessionData ->
+                    matrixConfiguration.cryptoAnalyticsPlugin
+                            ?.onRoomKeyImported(sessionData.sessionId.orEmpty(), keysVersionResult.algorithm)
+                }
                 megolmSessionImportManager.dispatchKeyImportResults(it)
             }
 

matrix-sdk-android/src/rustCrypto/java/org/matrix/android/sdk/internal/crypto/network/RequestSender.kt

@@ -19,17 +19,24 @@ package org.matrix.android.sdk.internal.crypto.network
 import com.squareup.moshi.Moshi
 import com.squareup.moshi.Types
 import dagger.Lazy
+import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.launch
 import org.matrix.android.sdk.api.auth.UserInteractiveAuthInterceptor
 import org.matrix.android.sdk.api.failure.Failure
 import org.matrix.android.sdk.api.failure.MatrixError
 import org.matrix.android.sdk.api.session.crypto.keysbackup.KeysBackupLastVersionResult
 import org.matrix.android.sdk.api.session.crypto.keysbackup.KeysVersion
 import org.matrix.android.sdk.api.session.crypto.keysbackup.KeysVersionResult
+import org.matrix.android.sdk.api.session.crypto.model.GossipingToDeviceObject
 import org.matrix.android.sdk.api.session.crypto.model.MXUsersDevicesMap
 import org.matrix.android.sdk.api.session.events.model.Content
 import org.matrix.android.sdk.api.session.events.model.Event
+import org.matrix.android.sdk.api.session.events.model.EventType
 import org.matrix.android.sdk.api.session.uia.UiaResult
 import org.matrix.android.sdk.internal.auth.registration.handleUIA
+import org.matrix.android.sdk.internal.crypto.PerSessionBackupQueryRateLimiter
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.BackupKeysResult
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.CreateKeysBackupVersionBody
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.KeysBackupData
@@ -87,8 +94,14 @@ internal class RequestSender @Inject constructor(
         private val getRoomSessionsDataTask: GetRoomSessionsDataTask,
         private val getRoomSessionDataTask: GetRoomSessionDataTask,
         private val moshi: Moshi,
+        private val cryptoCoroutineScope: CoroutineScope,
+        private val rateLimiter: PerSessionBackupQueryRateLimiter,
 ) {
 
+    private val scope = CoroutineScope(
+            cryptoCoroutineScope.coroutineContext + SupervisorJob() + CoroutineName("backupRequest")
+    )
+
     suspend fun claimKeys(request: Request.KeysClaim): String {
         val claimParams = ClaimOneTimeKeysForUsersDeviceTask.Params(request.oneTimeKeys)
         val response = oneTimeKeysForUsersDeviceTask.execute(claimParams)
@@ -211,9 +224,37 @@ internal class RequestSender @Inject constructor(
                 .newBuilder()
                 .add(CheckNumberType.JSON_ADAPTER_FACTORY)
                 .build()
-                .adapter<Map<String, HashMap<String, Any>>>(Map::class.java)
+                .adapter<Map<String, Map<String, Any>>>(Map::class.java)
         val jsonBody = adapter.fromJson(body)!!
 
+        if (eventType == EventType.ROOM_KEY_REQUEST) {
+            scope.launch {
+                Timber.v("Intercepting key request, try backup")
+                /**
+                 * It's a bit hacky, check how this can be better integrated with rust?
+                 */
+                try {
+                    jsonBody.forEach { (_, deviceToContent) ->
+                        deviceToContent.forEach { (_, content) ->
+                            val hashMap = content as? Map<*, *>
+                            val action = hashMap?.get("action")?.toString()
+                            if (GossipingToDeviceObject.ACTION_SHARE_REQUEST == action) {
+                                val body = hashMap.get("body") as? Map<*, *>
+                                val roomId = body?.get("room_id") as? String
+                                val sessionId = body?.get("session_id") as? String
+                                if (roomId != null && sessionId != null) {
+                                    rateLimiter.tryFromBackupIfPossible(sessionId, roomId)
+                                }
+                            }
+                        }
+                    }
+                    Timber.v("Intercepting key request, try backup")
+                } catch (failure: Throwable) {
+                    Timber.v(failure, "Failed to use backup")
+                }
+            }
+        }
+
         val userMap = MXUsersDevicesMap<Any>()
         userMap.join(jsonBody)