From 9118b26d2f1b065e02e8e491dc711ecda8c34f09 Mon Sep 17 00:00:00 2001
From: Valere <valeref@matrix.org>
Date: Tue, 23 Nov 2021 17:40:43 +0100
Subject: [PATCH] Fix DM are not e2e by default

---
 .../internal/crypto/DefaultCryptoService.kt   |  8 +----
 .../android/sdk/internal/crypto/OlmMachine.kt | 32 +++++++++++++++++++
 .../room/create/CreateRoomBodyBuilder.kt      |  2 +-
 3 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/DefaultCryptoService.kt b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/DefaultCryptoService.kt
index 70e792d4f6..cf58e11e52 100755
--- a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/DefaultCryptoService.kt
+++ b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/DefaultCryptoService.kt
@@ -1008,13 +1008,7 @@ internal class DefaultCryptoService @Inject constructor(
     override fun downloadKeys(userIds: List<String>, forceDownload: Boolean, callback: MatrixCallback<MXUsersDevicesMap<CryptoDeviceInfo>>) {
         cryptoCoroutineScope.launch(coroutineDispatchers.crypto) {
             runCatching {
-                if (forceDownload) {
-                    // TODO replicate the logic from the device list manager
-                    // where we would download the fresh info from the server.
-                    this@DefaultCryptoService.olmMachine.getUserDevicesMap(userIds) // ?: MXUsersDevicesMap()
-                } else {
-                    this@DefaultCryptoService.olmMachine.getUserDevicesMap(userIds) // ?: MXUsersDevicesMap()
-                }
+                olmMachine.ensureUserDevicesMap(userIds, forceDownload)
             }.foldToCallback(callback)
         }
     }
diff --git a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/OlmMachine.kt b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/OlmMachine.kt
index 2d503c1df0..b42f02b2e6 100644
--- a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/OlmMachine.kt
+++ b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/OlmMachine.kt
@@ -22,6 +22,7 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.withContext
 import org.matrix.android.sdk.api.auth.UserInteractiveAuthInterceptor
+import org.matrix.android.sdk.api.extensions.tryOrNull
 import org.matrix.android.sdk.api.listeners.ProgressListener
 import org.matrix.android.sdk.api.session.crypto.MXCryptoError
 import org.matrix.android.sdk.api.session.crypto.crosssigning.MXCrossSigningInfo
@@ -62,6 +63,7 @@ import uniffi.olm.RoomKeyCounts
 import uniffi.olm.setLogger
 import java.io.File
 import java.nio.charset.Charset
+import java.util.UUID
 import java.util.concurrent.ConcurrentHashMap
 import uniffi.olm.OlmMachine as InnerMachine
 import uniffi.olm.ProgressListener as RustProgressListener
@@ -644,6 +646,17 @@ internal class OlmMachine(
         return plainDevices
     }
 
+    @Throws
+    suspend fun forceKeyDownload(userIds: List<String>): MXUsersDevicesMap<CryptoDeviceInfo> {
+        withContext(Dispatchers.IO) {
+            val requestId = UUID.randomUUID().toString()
+            val response = requestSender.queryKeys(Request.KeysQuery(requestId, userIds))
+            markRequestAsSent(requestId, RequestType.KEYS_QUERY, response)
+        }
+        // TODO notify shield listener (?)
+        return getUserDevicesMap(userIds)
+    }
+
     suspend fun getUserDevicesMap(userIds: List<String>): MXUsersDevicesMap<CryptoDeviceInfo> {
         val userMap = MXUsersDevicesMap<CryptoDeviceInfo>()
 
@@ -658,6 +671,25 @@ internal class OlmMachine(
         return userMap
     }
 
+    /**
+     * If the user is untracked or forceDownload is set to true, a key query request will be made.
+     * It will suspend until query response, and the device list will be returned.
+     * As there is no API to know tracking status, we consider that if there are no known devices the user is un tracked.
+     *
+     * The key query request will be retried a few time in case of shaky connection, but could fail.
+     */
+    suspend fun ensureUserDevicesMap(userIds: List<String>, forceDownload: Boolean = false): MXUsersDevicesMap<CryptoDeviceInfo> {
+        val known = getUserDevicesMap(userIds)
+        return if (known.isEmpty /* We have no api to know if was tracked..*/ || forceDownload) {
+            updateTrackedUsers(userIds)
+             tryOrNull("Failed to download keys for $userIds") {
+                 forceKeyDownload(userIds)
+             } ?: known
+        } else {
+            known
+        }
+    }
+
     suspend fun getLiveUserIdentity(userId: String): LiveData<Optional<MXCrossSigningInfo>> {
         val identity = this.getIdentity(userId)?.toMxCrossSigningInfo().toOptional()
         val liveIdentity = LiveUserIdentity(userId, this.userIdentityUpdateObserver)
diff --git a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/session/room/create/CreateRoomBodyBuilder.kt b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/session/room/create/CreateRoomBodyBuilder.kt
index 7fcc6c3614..9b56354672 100644
--- a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/session/room/create/CreateRoomBodyBuilder.kt
+++ b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/session/room/create/CreateRoomBodyBuilder.kt
@@ -181,7 +181,7 @@ internal class CreateRoomBodyBuilder @Inject constructor(
                 params.invite3pids.isEmpty() &&
                 params.invitedUserIds.isNotEmpty() &&
                 params.invitedUserIds.let { userIds ->
-            val keys =  olmMachineProvider.olmMachine.getUserDevicesMap(userIds)
+            val keys =  olmMachineProvider.olmMachine.ensureUserDevicesMap(userIds, forceDownload = false)
             // deviceListManager.downloadKeys(userIds, forceDownload = false)
 
             userIds.all { userId ->