From e8bd7ea967b1379908132de531a17dfd8be5f4d5 Mon Sep 17 00:00:00 2001
From: Valere <valeref@matrix.org>
Date: Wed, 13 Jul 2022 11:47:48 +0200
Subject: [PATCH] fix olm session proliferation

---
 changelog.d/6534.bugfix                       |  1 +
 .../sdk/internal/crypto/UnwedgingTest.kt      |  1 -
 .../sdk/internal/crypto/EventDecryptor.kt     | 42 +++++++++++++------
 3 files changed, 30 insertions(+), 14 deletions(-)
 create mode 100644 changelog.d/6534.bugfix

diff --git a/changelog.d/6534.bugfix b/changelog.d/6534.bugfix
new file mode 100644
index 0000000000..721b61a2d5
--- /dev/null
+++ b/changelog.d/6534.bugfix
@@ -0,0 +1 @@
+Unwedging could cause the SDK to force creating a new olm session every hour
diff --git a/matrix-sdk-android/src/androidTest/java/org/matrix/android/sdk/internal/crypto/UnwedgingTest.kt b/matrix-sdk-android/src/androidTest/java/org/matrix/android/sdk/internal/crypto/UnwedgingTest.kt
index e8a474a54a..03e19f0bbb 100644
--- a/matrix-sdk-android/src/androidTest/java/org/matrix/android/sdk/internal/crypto/UnwedgingTest.kt
+++ b/matrix-sdk-android/src/androidTest/java/org/matrix/android/sdk/internal/crypto/UnwedgingTest.kt
@@ -60,7 +60,6 @@ import kotlin.coroutines.resume
  */
 @RunWith(AndroidJUnit4::class)
 @FixMethodOrder(MethodSorters.JVM)
-@Ignore
 class UnwedgingTest : InstrumentedTest {
 
     private lateinit var messagesReceivedByBob: List<TimelineEvent>
diff --git a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/EventDecryptor.kt b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/EventDecryptor.kt
index c1d04eb22b..bc3309132a 100644
--- a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/EventDecryptor.kt
+++ b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/EventDecryptor.kt
@@ -18,6 +18,8 @@ package org.matrix.android.sdk.internal.crypto
 
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
 import kotlinx.coroutines.withContext
 import org.matrix.android.sdk.api.MatrixCallback
 import org.matrix.android.sdk.api.MatrixCoroutineDispatchers
@@ -68,6 +70,7 @@ internal class EventDecryptor @Inject constructor(
             val senderKey: String?
     )
 
+    private val wedgedMutex = Mutex()
     private val wedgedDevices = mutableListOf<WedgedDeviceInfo>()
 
     /**
@@ -151,11 +154,13 @@ internal class EventDecryptor @Inject constructor(
         }
     }
 
-    private fun markOlmSessionForUnwedging(senderId: String, senderKey: String) {
-        val info = WedgedDeviceInfo(senderId, senderKey)
-        if (!wedgedDevices.contains(info)) {
-            Timber.tag(loggerTag.value).d("Marking device from $senderId key:$senderKey as wedged")
-            wedgedDevices.add(info)
+    private suspend fun markOlmSessionForUnwedging(senderId: String, senderKey: String) {
+        wedgedMutex.withLock {
+            val info = WedgedDeviceInfo(senderId, senderKey)
+            if (!wedgedDevices.contains(info)) {
+                Timber.tag(loggerTag.value).d("Marking device from $senderId key:$senderKey as wedged")
+                wedgedDevices.add(info)
+            }
         }
     }
 
@@ -167,15 +172,17 @@ internal class EventDecryptor @Inject constructor(
         Timber.tag(loggerTag.value).v("Unwedging:  ${wedgedDevices.size} are wedged")
         // get the one that should be retried according to rate limit
         val now = clock.epochMillis()
-        val toUnwedge = wedgedDevices.filter {
-            val lastForcedDate = lastNewSessionForcedDates[it] ?: 0
-            if (now - lastForcedDate < DefaultCryptoService.CRYPTO_MIN_FORCE_SESSION_PERIOD_MILLIS) {
-                Timber.tag(loggerTag.value).d("Unwedging, New session for $it already forced with device at $lastForcedDate")
-                return@filter false
+        val toUnwedge = wedgedMutex.withLock {
+            wedgedDevices.filter {
+                val lastForcedDate = lastNewSessionForcedDates[it] ?: 0
+                if (now - lastForcedDate < DefaultCryptoService.CRYPTO_MIN_FORCE_SESSION_PERIOD_MILLIS) {
+                    Timber.tag(loggerTag.value).d("Unwedging, New session for $it already forced with device at $lastForcedDate")
+                    return@filter false
+                }
+                // let's already mark that we tried now
+                lastNewSessionForcedDates[it] = now
+                true
             }
-            // let's already mark that we tried now
-            lastNewSessionForcedDates[it] = now
-            true
         }
 
         if (toUnwedge.isEmpty()) {
@@ -230,6 +237,15 @@ internal class EventDecryptor @Inject constructor(
                                     withContext(coroutineDispatchers.io) {
                                         sendToDeviceTask.executeRetry(sendToDeviceParams, remainingRetry = SEND_TO_DEVICE_RETRY_COUNT)
                                     }
+
+                                    deviceList.values.flatten().forEach { deviceInfo ->
+                                        wedgedMutex.withLock {
+                                            wedgedDevices.removeAll {
+                                                it.senderKey == deviceInfo.identityKey() &&
+                                                        it.userId == deviceInfo.userId
+                                            }
+                                        }
+                                    }
                                 } catch (failure: Throwable) {
                                     deviceList.flatMap { it.value }.joinToString { it.shortDebugString() }.let {
                                         Timber.tag(loggerTag.value).e(failure, "## Failed to unwedge devices: $it}")