Migration for big accounts + update sdk

library/rustCrypto/matrix-rust-sdk-crypto.aar

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:755bd2766aa1f317722c60259faf9759b9049885f20b91ed2a4dcfbcba861bfc
+oid sha256:8472e9bbcf8be3c6dc1b5ea3841c29098f5144d4f2975e843103318dde4cdf40
-size 17402892
+size 17414209

matrix-sdk-android/src/kotlinCrypto/java/org/matrix/android/sdk/internal/crypto/MXOlmDevice.kt

@@ -864,7 +864,8 @@ internal class MXOlmDevice @Inject constructor(
             } else {
                 val isDeviceOwnerOfSession = sessionHolder.wrapper.sessionData.keysClaimed?.get("ed25519") == sendingDevice.fingerprint()
                 if (!isDeviceOwnerOfSession) {
-                    MessageVerificationState.MISMATCH
+                    // should it fail to decrypt here?
+                    MessageVerificationState.UNSAFE_SOURCE
                 } else if (sendingDevice.isVerified) {
                     MessageVerificationState.VERIFIED
                 } else {

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

@@ -25,7 +25,6 @@ enum class MessageVerificationState {
     UN_SIGNED_DEVICE,
     UNKNOWN_DEVICE,
     UNSAFE_SOURCE,
-    MISMATCH,
 }
 
 /**

matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/session/SessionModule.kt

@@ -99,9 +99,11 @@ import org.matrix.android.sdk.internal.session.typing.DefaultTypingUsersTracker
 import org.matrix.android.sdk.internal.session.user.accountdata.DefaultSessionAccountDataService
 import org.matrix.android.sdk.internal.session.widgets.DefaultWidgetURLFormatter
 import retrofit2.Retrofit
+import timber.log.Timber
 import java.io.File
 import javax.inject.Provider
 import javax.inject.Qualifier
+import kotlin.system.measureTimeMillis
 
 @Qualifier
 @Retention(AnnotationRetention.RUNTIME)
@@ -189,7 +191,13 @@ internal abstract class SessionModule {
                 @CryptoDatabase realmConfiguration: RealmConfiguration,
         ): File {
             val target = File(parent, "rustFlavor")
-            return MigrateEAtoEROperation().execute(realmConfiguration, target)
+            val file: File
+            measureTimeMillis {
+                file = MigrateEAtoEROperation().execute(realmConfiguration, target)
+            }.let { duration ->
+                Timber.v("Migrating to ER in $duration ms")
+            }
+            return file
         }
 
         @JvmStatic

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

@@ -30,7 +30,7 @@ internal class DecryptRoomEventUseCase @Inject constructor(private val olmMachin
 
     suspend fun decryptAndSaveResult(event: Event) {
         tryOrNull(message = "Unable to decrypt the event") {
-            olmMachine.decryptRoomEvent(event)
+            invoke(event)
         }
                 ?.let { result ->
                     event.mxDecryptionResult = OlmDecryptionResult(

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

@@ -129,9 +129,14 @@ internal class OlmMachine @Inject constructor(
         private val getUserIdentity: GetUserIdentityUseCase,
         private val ensureUsersKeys: EnsureUsersKeysUseCase,
         private val matrixConfiguration: MatrixConfiguration,
+        private val megolmSessionImportManager: MegolmSessionImportManager,
 ) {
 
-    private val inner: InnerMachine = InnerMachine(userId, deviceId, path.toString(), null)
+    private val inner: InnerMachine
+
+    init {
+        inner = InnerMachine(userId, deviceId, path.toString(), null)
+    }
 
     private val flowCollectors = FlowCollectors()
 
@@ -312,6 +317,22 @@ internal class OlmMachine @Inject constructor(
         inner.receiveVerificationEvent(serializedEvent, roomId)
     }
 
+    /**
+     * Used for lazy migration of inboundGroupSession from EA to ER
+     */
+    suspend fun importRoomKey(inbound: InboundGroupSessionHolder): Result<Unit> {
+        Timber.v("Migration:: Tentative lazy migration")
+        return withContext(coroutineDispatchers.io) {
+            val export = inbound.wrapper.exportKeys()
+                    ?: return@withContext Result.failure(Exception("Failed to export key"))
+            val result = importDecryptedKeys(listOf(export), null).also {
+                Timber.v("Migration:: Tentative lazy migration result: ${it.totalNumberOfKeys}")
+            }
+            if (result.totalNumberOfKeys == 1) return@withContext Result.success(Unit)
+            return@withContext Result.failure(Exception("Import failed"))
+        }
+    }
+
     /**
      * Mark the given list of users to be tracked, triggering a key query request for them.
      *
@@ -567,7 +588,9 @@ internal class OlmMachine @Inject constructor(
                                 details.putAll(it.keys)
                             }
                         }
-                ImportRoomKeysResult(totalImported.toInt(), accTotal.toInt(), details)
+                ImportRoomKeysResult(totalImported.toInt(), accTotal.toInt(), details).also {
+                    megolmSessionImportManager.dispatchKeyImportResults(it)
+                }
             }
 
     @Throws(CryptoStoreException::class)

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

@@ -24,9 +24,8 @@ fun InnerVerificationState.fromInner(): MessageVerificationState {
         InnerVerificationState.VERIFIED -> MessageVerificationState.VERIFIED
         InnerVerificationState.SIGNED_DEVICE_OF_UNVERIFIED_USER -> MessageVerificationState.SIGNED_DEVICE_OF_UNVERIFIED_USER
         InnerVerificationState.UN_SIGNED_DEVICE_OF_VERIFIED_USER -> MessageVerificationState.UN_SIGNED_DEVICE_OF_VERIFIED_USER
-        InnerVerificationState.UN_SIGNED_DEVICE -> MessageVerificationState.UN_SIGNED_DEVICE
+        InnerVerificationState.UN_SIGNED_DEVICE_OF_UNVERIFIED_USER -> MessageVerificationState.UN_SIGNED_DEVICE
         InnerVerificationState.UNKNOWN_DEVICE -> MessageVerificationState.UNKNOWN_DEVICE
         InnerVerificationState.UNSAFE_SOURCE -> MessageVerificationState.UNSAFE_SOURCE
-        InnerVerificationState.MISMATCH -> MessageVerificationState.MISMATCH
     }
 }

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

@@ -36,6 +36,8 @@ 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.InboundGroupSessionStore
+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.rest.BackupKeysResult
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.CreateKeysBackupVersionBody
@@ -100,7 +102,9 @@ internal class RequestSender @Inject constructor(
         private val moshi: Moshi,
         cryptoCoroutineScope: CoroutineScope,
         private val rateLimiter: PerSessionBackupQueryRateLimiter,
-        private val localEchoRepository: LocalEchoRepository
+        private val inboundGroupSessionStore: InboundGroupSessionStore,
+        private val localEchoRepository: LocalEchoRepository,
+        private val olmMachine: Lazy<OlmMachine>,
 ) {
 
     private val scope = CoroutineScope(
@@ -259,12 +263,17 @@ internal class RequestSender @Inject constructor(
                                 val requestBody = hashMap["body"] as? Map<*, *>
                                 val roomId = requestBody?.get("room_id") as? String
                                 val sessionId = requestBody?.get("session_id") as? String
+                                val senderKey = requestBody?.get("sender_key") as? String
                                 if (roomId != null && sessionId != null) {
+                                    // try to perform a lazy migration from legacy store
+                                    val legacy = inboundGroupSessionStore.getInboundGroupSession(sessionId, senderKey.orEmpty())
+                                    if (legacy == null || olmMachine.get().importRoomKey(legacy).isFailure) {
                                         rateLimiter.tryFromBackupIfPossible(sessionId, roomId)
                                     }
                                 }
                             }
                         }
+                    }
                     Timber.v("Intercepting key request, try backup")
                 } catch (failure: Throwable) {
                     Timber.v(failure, "Failed to use backup")

matrix-sdk-android/src/rustCrypto/java/org/matrix/android/sdk/internal/crypto/store/db/migration/rust/ExtractMigrationDataUseCase.kt

@@ -19,6 +19,7 @@ package org.matrix.android.sdk.internal.crypto.store.db.migration.rust
 import io.realm.Realm
 import io.realm.RealmConfiguration
 import io.realm.kotlin.where
+import org.matrix.android.sdk.api.extensions.orFalse
 import org.matrix.android.sdk.internal.crypto.store.db.deserializeFromRealm
 import org.matrix.android.sdk.internal.crypto.store.db.model.CryptoMetadataEntity
 import org.matrix.android.sdk.internal.crypto.store.db.model.OlmInboundGroupSessionEntity
@@ -33,14 +34,15 @@ import org.matrix.rustcomponents.sdk.crypto.PickledInboundGroupSession
 import org.matrix.rustcomponents.sdk.crypto.PickledSession
 import timber.log.Timber
 import java.nio.charset.Charset
+import kotlin.system.measureTimeMillis
 
 private val charset = Charset.forName("UTF-8")
 
 internal class ExtractMigrationDataUseCase {
 
-    fun extractData(realm: Realm): MigrationData {
+    fun extractData(realm: Realm, importPartial: ((MigrationData) -> Unit)): MigrationData {
         return try {
-            extract(realm) ?: throw ExtractMigrationDataFailure
+            extract(realm, importPartial) ?: throw ExtractMigrationDataFailure
         } catch (failure: Throwable) {
             throw ExtractMigrationDataFailure
         }
@@ -55,17 +57,12 @@ internal class ExtractMigrationDataUseCase {
         }
     }
 
-    private fun extract(realm: Realm): MigrationData? {
+    private fun extract(realm: Realm, importPartial: ((MigrationData) -> Unit)): MigrationData? {
         val metadataEntity = realm.where<CryptoMetadataEntity>().findFirst() ?: return null.also {
             Timber.w("Rust db migration: No existing metadataEntity")
         }
 
         val pickleKey = OlmUtility.getRandomKey()
-        val olmSessionEntities = realm.where<OlmSessionEntity>().findAll()
-        val pickledSessions = olmSessionEntities.map { it.toPickledSession(pickleKey) }
-
-        val inboundGroupSessionEntities = realm.where<OlmInboundGroupSessionEntity>().findAll()
-        val pickledInboundGroupSessions = inboundGroupSessionEntities.mapNotNull { it.toPickledInboundGroupSession(pickleKey) }
 
         val masterKey = metadataEntity.xSignMasterPrivateKey
         val userKey = metadataEntity.xSignUserPrivateKey
@@ -76,14 +73,11 @@ internal class ExtractMigrationDataUseCase {
         val backupVersion = metadataEntity.backupVersion
         val backupRecoveryKey = metadataEntity.keyBackupRecoveryKey
 
-        val trackedUserEntities = realm.where<UserEntity>().findAll()
-        val trackedUserIds = trackedUserEntities.mapNotNull {
-            it.userId
-        }
         val isOlmAccountShared = metadataEntity.deviceKeysSentToServer
 
         val olmAccount = metadataEntity.getOlmAccount() ?: return null
         val pickledOlmAccount = olmAccount.pickle(pickleKey, StringBuffer()).asString()
+        olmAccount.oneTimeKeys()
         val pickledAccount = PickledAccount(
                 userId = userId,
                 deviceId = deviceId,
@@ -91,20 +85,90 @@ internal class ExtractMigrationDataUseCase {
                 shared = isOlmAccountShared,
                 uploadedSignedKeyCount = 50
         )
-        return MigrationData(
+
+        val baseExtract = MigrationData(
                 account = pickledAccount,
-                sessions = pickledSessions,
-                inboundGroupSessions = pickledInboundGroupSessions,
                 pickleKey = pickleKey.map { it.toUByte() },
-                backupVersion = backupVersion,
-                backupRecoveryKey = backupRecoveryKey,
                 crossSigning = CrossSigningKeyExport(
                         masterKey = masterKey,
                         selfSigningKey = selfSignedKey,
                         userSigningKey = userKey
                 ),
-                trackedUsers = trackedUserIds
+                sessions = emptyList(),
+                backupRecoveryKey = backupRecoveryKey,
+                trackedUsers = emptyList(),
+                inboundGroupSessions = emptyList(),
+                backupVersion = backupVersion,
         )
+        // import the account asap
+        importPartial(baseExtract)
+
+        val chunkSize = 500
+        realm.where<UserEntity>()
+                .findAll()
+                .chunked(chunkSize) { chunk ->
+                    val trackedUserIds = chunk.mapNotNull { it.userId }
+                    importPartial(
+                            baseExtract.copy(trackedUsers = trackedUserIds)
+                    )
+                }
+
+        var migratedOlmSessionCount = 0
+        var readTime = 0L
+        var writeTime = 0L
+        measureTimeMillis {
+            realm.where<OlmSessionEntity>().findAll()
+                    .chunked(chunkSize) { chunk ->
+                        migratedOlmSessionCount += chunk.size
+                        val export: List<PickledSession>
+                        measureTimeMillis {
+                            export = chunk.map { it.toPickledSession(pickleKey) }
+                        }.also {
+                            readTime += it
+                        }
+                        measureTimeMillis {
+                            importPartial(
+                                    baseExtract.copy(sessions = export)
+                            )
+                        }.also { writeTime += it }
+                    }
+        }.also {
+            Timber.i("Migration: took $it ms to migrate $migratedOlmSessionCount olm sessions")
+            Timber.i("Migration: extract time $readTime")
+            Timber.i("Migration: rust import time $writeTime")
+        }
+
+        // We don't migrate outbound session directly after migration
+        // We are going to do it lazyly when decryption fails
+//        var migratedInboundGroupSessionCount = 0
+//        readTime = 0
+//        writeTime = 0
+//        measureTimeMillis {
+//        realm.where<OlmInboundGroupSessionEntity>()
+//                .findAll()
+//                .chunked(chunkSize) { chunk ->
+//                    val export: List<PickledInboundGroupSession>
+//                    measureTimeMillis {
+//                        export = chunk.mapNotNull { it.toPickledInboundGroupSession(pickleKey) }
+//                    }.also {
+//                        readTime += it
+//                    }
+//                    migratedInboundGroupSessionCount+=export.size
+//                    measureTimeMillis {
+//                        importPartial(
+//                                baseExtract.copy(inboundGroupSessions = export)
+//                        )
+//                    }.also {
+//                        writeTime += it
+//                    }
+//                }
+//        }.also {
+//            Timber.i("Migration: took $it ms to migrate $migratedInboundGroupSessionCount group sessions")
+//            Timber.i("Migration: extract time $readTime")
+//            Timber.i("Migration: rust import time $writeTime")
+//        }
+
+        return baseExtract
     }
 
     private fun OlmInboundGroupSessionEntity.toPickledInboundGroupSession(pickleKey: ByteArray): PickledInboundGroupSession? {
@@ -126,7 +190,7 @@ internal class ExtractMigrationDataUseCase {
                 signingKey = data.keysClaimed.orEmpty(),
                 roomId = roomId,
                 forwardingChains = data.forwardingCurve25519KeyChain.orEmpty(),
-                imported = true,
+                imported = data.trusted.orFalse().not(),
                 backedUp = backedUp
         )
     }

matrix-sdk-android/src/rustCrypto/java/org/matrix/android/sdk/internal/session/MigrateEAtoEROperation.kt

@@ -25,14 +25,14 @@ import java.io.File
 
 class MigrateEAtoEROperation {
 
-    fun execute(cryptoRealm: RealmConfiguration, sessionFilesDir: File): File {
+    fun execute(cryptoRealm: RealmConfiguration, rustFilesDir: File): File {
         // Temporary code for migration
-        if (!sessionFilesDir.exists()) {
+        if (!rustFilesDir.exists()) {
-            sessionFilesDir.mkdir()
+            rustFilesDir.mkdir()
             // perform a migration?
             val extractMigrationData = ExtractMigrationDataUseCase()
             val hasExitingData = extractMigrationData.hasExistingData(cryptoRealm)
-            if (!hasExitingData) return sessionFilesDir
+            if (!hasExitingData) return rustFilesDir
 
             try {
                 val progressListener = object : ProgressListener {
@@ -42,14 +42,15 @@ class MigrateEAtoEROperation {
                 }
 
                 Realm.getInstance(cryptoRealm).use { realm ->
-                    val migrationData = extractMigrationData.extractData(realm)
+                    extractMigrationData.extractData(realm) {
-                    org.matrix.rustcomponents.sdk.crypto.migrate(migrationData, sessionFilesDir.path, null, progressListener)
+                        org.matrix.rustcomponents.sdk.crypto.migrate(it, rustFilesDir.path, null, progressListener)
+                    }
                 }
             } catch (failure: Throwable) {
                 Timber.e(failure, "Failure while calling rust migration method")
                 throw failure
             }
         }
-        return sessionFilesDir
+        return rustFilesDir
     }
 }

vector/src/main/java/im/vector/app/core/di/ActiveSessionHolder.kt

@@ -30,7 +30,6 @@ import im.vector.app.features.session.SessionListener
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.withContext
 import org.matrix.android.sdk.api.auth.AuthenticationService
 import org.matrix.android.sdk.api.session.Session
 import org.matrix.android.sdk.api.util.Optional

vector/src/main/java/im/vector/app/features/home/room/detail/timeline/helper/MessageInformationDataFactory.kt

@@ -26,7 +26,6 @@ import im.vector.app.features.home.room.detail.timeline.item.MessageInformationD
 import im.vector.app.features.home.room.detail.timeline.item.ReferencesInfoData
 import im.vector.app.features.home.room.detail.timeline.item.SendStateDecoration
 import im.vector.app.features.home.room.detail.timeline.style.TimelineMessageLayoutFactory
-import kotlinx.coroutines.runBlocking
 import org.matrix.android.sdk.api.extensions.orFalse
 import org.matrix.android.sdk.api.session.Session
 import org.matrix.android.sdk.api.session.crypto.model.MessageVerificationState
@@ -75,11 +74,9 @@ class MessageInformationDataFactory @Inject constructor(
                 prevDisplayableEvent?.root?.localDateTime()?.toLocalDate() != date.toLocalDate()
 
         val time = dateFormatter.format(event.root.originServerTs, DateFormatKind.MESSAGE_SIMPLE)
-//        val e2eDecoration = runBlocking {
-//            getE2EDecoration(roomSummary, params.lastEdit ?: event.root)
-//        }
         val e2eDecoration = getE2EDecorationV2(roomSummary, params.lastEdit ?: event.root)
-        val senderId = runBlocking { getSenderId(event) }
+        // this is claimed data or not depending on the e2e decoration
+        val senderId = event.senderInfo.userId
         // SendState Decoration
         val sendStateDecoration = if (isSentByMe) {
             getSendStateDecoration(
@@ -168,7 +165,6 @@ class MessageInformationDataFactory @Inject constructor(
             MessageVerificationState.UN_SIGNED_DEVICE ->  E2EDecoration.NONE
             MessageVerificationState.UNKNOWN_DEVICE -> E2EDecoration.WARN_SENT_BY_DELETED_SESSION
             MessageVerificationState.UNSAFE_SOURCE -> E2EDecoration.WARN_UNSAFE_KEY
-            MessageVerificationState.MISMATCH -> E2EDecoration.WARN_UNSAFE_KEY
             null -> {
                 // No verification state.
                 // So could be a clear event, or a legacy decryption, or an UTD event