proper initial withheld support

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

@@ -35,6 +35,7 @@ import org.matrix.android.sdk.api.session.crypto.model.RoomKeyShareRequest
 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.events.model.content.EncryptedEventContent
+import org.matrix.android.sdk.api.session.events.model.content.RoomKeyWithHeldContent
 import org.matrix.android.sdk.api.session.events.model.toModel
 import org.matrix.android.sdk.api.util.fromBase64
 import org.matrix.android.sdk.internal.crypto.store.IMXCryptoStore
@@ -68,6 +69,7 @@ internal class OutgoingKeyRequestManager @Inject constructor(
         private val cryptoConfig: MXCryptoConfig,
         private val inboundGroupSessionStore: InboundGroupSessionStore,
         private val sendToDeviceTask: DefaultSendToDeviceTask,
+        private val deviceListManager: DeviceListManager,
         private val perSessionBackupQueryRateLimiter: PerSessionBackupQueryRateLimiter) {
 
     private val dispatcher = Executors.newSingleThreadExecutor().asCoroutineDispatcher()
@@ -215,6 +217,41 @@ internal class OutgoingKeyRequestManager @Inject constructor(
         outgoingRequestScope.launch {
             sequencer.post {
                 Timber.tag(loggerTag.value).d("Withheld received for $sessionId from ${event.senderId}|$fromDevice")
+                Timber.tag(loggerTag.value).v("Withheld content ${event.getClearContent()}")
+
+                // We want to store withheld code from the sender of the message (owner of the megolm session), not from
+                // other devices that might gossip the key. If not the initial reason might be overridden
+                // by a request to one of our session.
+                event.getClearContent().toModel<RoomKeyWithHeldContent>()?.let { withheld ->
+                    withContext(coroutineDispatchers.crypto) {
+                        tryOrNull {
+                            deviceListManager.downloadKeys(listOf(event.senderId ?: ""), false)
+                        }
+                        cryptoStore.getUserDeviceList(event.senderId ?: "")
+                                .also { devices ->
+                                    Timber.tag(loggerTag.value).v("Withheld Devices for ${event.senderId} are ${devices.orEmpty().joinToString { it.identityKey() ?: "" }}")
+                                }
+                                ?.firstOrNull {
+                                    it.identityKey() == senderKey
+                                }
+                    }.also {
+                        Timber.tag(loggerTag.value).v("Withheld device for sender key $senderKey is from ${it?.shortDebugString()}")
+                    }?.let {
+                        if (it.userId == event.senderId) {
+                            if (fromDevice != null) {
+                                if (it.deviceId == fromDevice) {
+                                    Timber.tag(loggerTag.value).v("Storing sender Withheld code ${withheld.code} for ${withheld.sessionId}")
+                                    cryptoStore.addWithHeldMegolmSession(withheld)
+                                }
+                            } else {
+                                Timber.tag(loggerTag.value).v("Storing sender Withheld code ${withheld.code} for ${withheld.sessionId}")
+                                cryptoStore.addWithHeldMegolmSession(withheld)
+                            }
+                        }
+                    }
+                }
+
+                // Here we store the replies from a given request
                 cryptoStore.updateOutgoingRoomKeyReply(
                         roomId = roomId,
                         sessionId = sessionId,

matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/algorithms/megolm/MXMegolmDecryption.kt

@@ -102,9 +102,20 @@ internal class MXMegolmDecryption(
                             if (throwable is MXCryptoError.OlmError) {
                                 // TODO Check the value of .message
                                 if (throwable.olmException.message == "UNKNOWN_MESSAGE_INDEX") {
+                                    // So we know that session, but it's ratcheted and we can't decrypt at that index
+
                                     if (requestKeysOnFail) {
                                         requestKeysForEvent(event)
                                     }
+                                    // Check if partially withheld
+                                    val withHeldInfo = cryptoStore.getWithHeldMegolmSession(event.roomId, encryptedEventContent.sessionId)
+                                    if (withHeldInfo != null) {
+                                        // Encapsulate as withHeld exception
+                                        throw MXCryptoError.Base(MXCryptoError.ErrorType.KEYS_WITHHELD,
+                                                withHeldInfo.code?.value ?: "",
+                                                withHeldInfo.reason)
+                                    }
+
                                     throw MXCryptoError.Base(
                                             MXCryptoError.ErrorType.UNKNOWN_MESSAGE_INDEX,
                                             "UNKNOWN_MESSAGE_INDEX",
@@ -121,6 +132,18 @@ internal class MXMegolmDecryption(
                             }
                             if (throwable is MXCryptoError.Base) {
                                 if (throwable.errorType == MXCryptoError.ErrorType.UNKNOWN_INBOUND_SESSION_ID) {
+                                    // Check if it was withheld by sender to enrich error code
+                                    val withHeldInfo = cryptoStore.getWithHeldMegolmSession(event.roomId, encryptedEventContent.sessionId)
+                                    if (withHeldInfo != null) {
+                                        if (requestKeysOnFail) {
+                                            requestKeysForEvent(event)
+                                        }
+                                        // Encapsulate as withHeld exception
+                                        throw MXCryptoError.Base(MXCryptoError.ErrorType.KEYS_WITHHELD,
+                                                withHeldInfo.code?.value ?: "",
+                                                withHeldInfo.reason)
+                                    }
+
                                     if (requestKeysOnFail) {
                                         requestKeysForEvent(event)
                                     }

matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/algorithms/megolm/MXMegolmEncryption.kt

@@ -328,7 +328,8 @@ internal class MXMegolmEncryption(
                 senderKey = senderKey,
                 algorithm = MXCRYPTO_ALGORITHM_MEGOLM,
                 sessionId = sessionId,
-                codeString = code.value
+                codeString = code.value,
+                fromDevice = myDeviceId
         )
         val params = SendToDeviceTask.Params(
                 EventType.ROOM_KEY_WITHHELD,

matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/store/db/RealmCryptoStore.kt

@@ -271,12 +271,13 @@ internal class RealmCryptoStore @Inject constructor(
     }
 
     override fun deviceWithIdentityKey(identityKey: String): CryptoDeviceInfo? {
-        return doWithRealm(realmConfiguration) {
-            it.where<DeviceInfoEntity>()
-                    .equalTo(DeviceInfoEntityFields.IDENTITY_KEY, identityKey)
-                    .findFirst()
-                    ?.let { deviceInfo ->
-                        CryptoMapper.mapToModel(deviceInfo)
+        return doWithRealm(realmConfiguration) { realm ->
+            realm.where<DeviceInfoEntity>()
+                    .contains(DeviceInfoEntityFields.KEYS_MAP_JSON, identityKey)
+                    .findAll()
+                    .mapNotNull { CryptoMapper.mapToModel(it)  }
+                    .firstOrNull {
+                        it.identityKey() == identityKey
                     }
         }
     }