Restore manual Client cleanup on session logout

libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/RustClientSessionDelegate.kt

@@ -19,6 +19,7 @@ import org.matrix.rustcomponents.sdk.ClientDelegate
 import org.matrix.rustcomponents.sdk.ClientSessionDelegate
 import org.matrix.rustcomponents.sdk.Session
 import timber.log.Timber
+import java.lang.ref.WeakReference
 import java.util.concurrent.atomic.AtomicBoolean
 
 /**
@@ -43,13 +44,20 @@ class RustClientSessionDelegate(
     private val updateTokensDispatcher = coroutineDispatchers.io.limitedParallelism(1)
 
     // This Client needs to be set up as soon as possible so `didReceiveAuthError` can work properly.
-    private var client: RustMatrixClient? = null
+    private var client: WeakReference<RustMatrixClient> = WeakReference(null)
 
     /**
      * Sets the [ClientDelegate] for the [RustMatrixClient], and keeps a reference to the client so it can be used later.
      */
     fun bindClient(client: RustMatrixClient) {
-        this.client = client
+        this.client = WeakReference(client)
+    }
+
+    /**
+     * Clears the current client reference.
+     */
+    fun clearCurrentClient() {
+        this.client.clear()
     }
 
     override fun saveSessionInKeychain(session: Session) {
@@ -81,7 +89,7 @@ class RustClientSessionDelegate(
             clientLog.v("didReceiveAuthError -> do the cleanup")
             // TODO handle isSoftLogout parameter.
             appCoroutineScope.launch(updateTokensDispatcher) {
-                val currentClient = client
+                val currentClient = client.get()
                 if (currentClient == null) {
                     clientLog.w("didReceiveAuthError -> no client, exiting")
                     isLoggingOut.set(false)
@@ -101,7 +109,7 @@ class RustClientSessionDelegate(
                 } else {
                     clientLog.d("No session data found.")
                 }
-                client?.logout(userInitiated = false, ignoreSdkError = true)
+                currentClient.logout(userInitiated = false, ignoreSdkError = true)
             }.invokeOnCompletion {
                 if (it != null) {
                     clientLog.e(it, "Failed to remove session data.")

libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/RustMatrixClient.kt

@@ -482,6 +482,13 @@ class RustMatrixClient(
         clientDelegateTaskHandle?.cancelAndDestroy()
         notificationSettingsService.destroy()
         verificationService.destroy()
+
+        sessionDelegate.clearCurrentClient()
+        innerRoomListService.destroy()
+        notificationService.destroy()
+        notificationProcessSetup.destroy()
+        encryptionService.destroy()
+        innerClient.destroy()
     }
 
     override suspend fun getCacheSize(): Long {

libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/auth/RustMatrixAuthenticationService.kt

@@ -156,6 +156,10 @@ class RustMatrixAuthenticationService @Inject constructor(
                     )
                 newMatrixClientObserver?.invoke(rustMatrixClientFactory.create(client))
                 sessionStore.storeData(sessionData)
+
+                // Clean up the strong reference held here since it's no longer necessary
+                currentClient = null
+
                 SessionId(sessionData.userId)
             }.mapFailure { failure ->
                 failure.mapAuthenticationException()
@@ -229,6 +233,10 @@ class RustMatrixAuthenticationService @Inject constructor(
                 pendingOidcAuthorizationData = null
                 newMatrixClientObserver?.invoke(rustMatrixClientFactory.create(client))
                 sessionStore.storeData(sessionData)
+
+                // Clean up the strong reference held here since it's no longer necessary
+                currentClient = null
+
                 SessionId(sessionData.userId)
             }.mapFailure { failure ->
                 failure.mapAuthenticationException()
@@ -264,6 +272,10 @@ class RustMatrixAuthenticationService @Inject constructor(
                     )
                 newMatrixClientObserver?.invoke(rustMatrixClientFactory.create(client))
                 sessionStore.storeData(sessionData)
+
+                // Clean up the strong reference held here since it's no longer necessary
+                currentClient = null
+
                 SessionId(sessionData.userId)
             }.mapFailure {
                 when (it) {

libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/encryption/RustEncryptionService.kt

@@ -186,11 +186,11 @@ internal class RustEncryptionService(
     }
 
     override suspend fun deviceCurve25519(): String? {
-        return service.curve25519Key()
+        return runCatching { service.curve25519Key() }.getOrNull()
     }
 
     override suspend fun deviceEd25519(): String? {
-        return service.ed25519Key()
+        return runCatching { service.ed25519Key() }.getOrNull()
     }
 
     override suspend fun startIdentityReset(): Result<IdentityResetHandle?> {
@@ -219,4 +219,8 @@ internal class RustEncryptionService(
             // requestFromHomeserverIfNeeded = true,
         ) ?: error("User identity not found")
     }
+
+    fun destroy() {
+        service.destroy()
+    }
 }

libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/notification/RustNotificationService.kt

@@ -35,4 +35,8 @@ class RustNotificationService(
             }
         }
     }
+
+    fun destroy() {
+        notificationClient.destroy()
+    }
 }