chore: bulletproofing crypto box to cc migration (WPB-14250)

cryptography/src/commonJvmAndroid/kotlin/com.wire.kalium.cryptography/ProteusClientCoreCryptoImpl.kt

@@ -22,6 +22,7 @@ import com.wire.crypto.CoreCrypto
 import com.wire.crypto.CoreCryptoException
 import com.wire.crypto.client.toByteArray
 import com.wire.kalium.cryptography.exceptions.ProteusException
+import com.wire.kalium.cryptography.exceptions.ProteusStorageMigrationException
 import io.ktor.util.decodeBase64Bytes
 import io.ktor.util.encodeBase64
 import kotlinx.coroutines.sync.Mutex
@@ -178,36 +179,44 @@ class ProteusClientCoreCryptoImpl private constructor(
                 acc && File(rootDir).resolve(file).deleteRecursively()
             }
 
-        private suspend fun migrateFromCryptoBoxIfNecessary(coreCrypto: CoreCrypto, rootDir: String) {
-            if (cryptoBoxFilesExists(File(rootDir))) {
-                kaliumLogger.i("migrating from crypto box at: $rootDir")
-                coreCrypto.proteusCryptoboxMigrate(rootDir)
-                kaliumLogger.i("migration successful")
-
-                if (deleteCryptoBoxFiles(rootDir)) {
-                    kaliumLogger.i("successfully deleted old crypto box files")
-                } else {
-                    kaliumLogger.e("Failed to deleted old crypto box files at $rootDir")
-                }
-            }
-        }
-
-        @Suppress("TooGenericExceptionCaught")
+        @Suppress("TooGenericExceptionCaught", "ThrowsCount")
         suspend operator fun invoke(coreCrypto: CoreCrypto, rootDir: String): ProteusClientCoreCryptoImpl {
             try {
                 migrateFromCryptoBoxIfNecessary(coreCrypto, rootDir)
                 coreCrypto.proteusInit()
                 return ProteusClientCoreCryptoImpl(coreCrypto)
+            } catch (exception: ProteusStorageMigrationException) {
+                throw exception
             } catch (e: CoreCryptoException) {
                 throw ProteusException(
-                    e.message,
-                    ProteusException.fromProteusCode(coreCrypto.proteusLastErrorCode().toInt()),
-                    coreCrypto.proteusLastErrorCode().toInt(),
-                    e.cause
+                    message = e.message,
+                    code = ProteusException.fromProteusCode(coreCrypto.proteusLastErrorCode().toInt()),
+                    intCode = coreCrypto.proteusLastErrorCode().toInt(),
+                    cause = e.cause
                 )
             } catch (e: Exception) {
                 throw ProteusException(e.message, ProteusException.Code.UNKNOWN_ERROR, null, e.cause)
             }
         }
+
+        @Suppress("TooGenericExceptionCaught")
+        private suspend fun migrateFromCryptoBoxIfNecessary(coreCrypto: CoreCrypto, rootDir: String) {
+            try {
+                if (cryptoBoxFilesExists(File(rootDir))) {
+                    kaliumLogger.i("migrating from crypto box at: $rootDir")
+                    coreCrypto.proteusCryptoboxMigrate(rootDir)
+                    kaliumLogger.i("migration successful")
+
+                    if (deleteCryptoBoxFiles(rootDir)) {
+                        kaliumLogger.i("successfully deleted old crypto box files")
+                    } else {
+                        kaliumLogger.e("Failed to deleted old crypto box files at $rootDir")
+                    }
+                }
+            } catch (exception: Exception) {
+                kaliumLogger.e("Failed to migrate from crypto box to core crypto, exception: $exception")
+                throw ProteusStorageMigrationException("Failed to migrate from crypto box at $rootDir", exception)
+            }
+        }
     }
 }

cryptography/src/commonMain/kotlin/com/wire/kalium/cryptography/exceptions/ProteusException.kt

@@ -18,7 +18,7 @@
 
 package com.wire.kalium.cryptography.exceptions
 
-class ProteusException(message: String?, val code: Code, val intCode: Int?, cause: Throwable? = null) : Exception(message, cause) {
+open class ProteusException(message: String?, val code: Code, val intCode: Int?, cause: Throwable? = null) : Exception(message, cause) {
 
     constructor(message: String?, code: Int, cause: Throwable? = null) : this(
         message,
@@ -199,3 +199,6 @@ class ProteusException(message: String?, val code: Code, val intCode: Int?, caus
         }
     }
 }
+
+class ProteusStorageMigrationException(override val message: String, val rootCause: Throwable? = null) :
+    ProteusException(message, Int.MIN_VALUE, null)

data/src/commonMain/kotlin/com/wire/kalium/logic/data/logout/LogoutReason.kt

@@ -46,5 +46,11 @@ enum class LogoutReason {
     /**
      * Session Expired.
      */
-    SESSION_EXPIRED;
+    SESSION_EXPIRED,
+
+    /**
+     * The migration to CC failed.
+     * This will trigger a cleanup of the local client data and prepare for a fresh start without losing data.
+     */
+    MIGRATION_TO_CC_FAILED
 }

logic/src/commonMain/kotlin/com/wire/kalium/logic/data/client/ProteusClientProvider.kt

@@ -18,9 +18,11 @@
 
 package com.wire.kalium.logic.data.client
 
+import com.wire.kalium.cryptography.CoreCryptoCentral
 import com.wire.kalium.cryptography.ProteusClient
 import com.wire.kalium.cryptography.coreCryptoCentral
 import com.wire.kalium.cryptography.cryptoboxProteusClient
+import com.wire.kalium.cryptography.exceptions.ProteusStorageMigrationException
 import com.wire.kalium.logger.KaliumLogLevel
 import com.wire.kalium.logger.obfuscateId
 import com.wire.kalium.logic.CoreFailure
@@ -58,20 +60,15 @@ class ProteusClientProviderImpl(
     private val userId: UserId,
     private val passphraseStorage: PassphraseStorage,
     private val kaliumConfigs: KaliumConfigs,
-    private val dispatcher: KaliumDispatcher = KaliumDispatcherImpl
+    private val dispatcher: KaliumDispatcher = KaliumDispatcherImpl,
+    private val proteusMigrationRecoveryHandler: ProteusMigrationRecoveryHandler
 ) : ProteusClientProvider {
 
     private var _proteusClient: ProteusClient? = null
     private val mutex = Mutex()
 
     override suspend fun clearLocalFiles() {
-        mutex.withLock {
-            withContext(dispatcher.io) {
-                _proteusClient?.close()
-                _proteusClient = null
-                FileUtil.deleteDirectory(rootProteusPath)
-            }
-        }
+        mutex.withLock { removeLocalFiles() }
     }
 
     override suspend fun getOrCreate(): ProteusClient {
@@ -109,7 +106,6 @@ class ProteusClientProviderImpl(
                     databaseKey = SecurityHelperImpl(passphraseStorage).proteusDBSecret(userId).value
                 )
             } catch (e: Exception) {
-
                 val logMap = mapOf(
                     "userId" to userId.value.obfuscateId(),
                     "exception" to e,
@@ -119,7 +115,7 @@ class ProteusClientProviderImpl(
                 kaliumLogger.logStructuredJson(KaliumLogLevel.ERROR, TAG, logMap)
                 throw e
             }
-            central.proteusClient()
+            getCentralProteusClientOrError(central)
         } else {
             cryptoboxProteusClient(
                 rootDir = rootProteusPath,
@@ -129,6 +125,34 @@ class ProteusClientProviderImpl(
         }
     }
 
+    private suspend fun getCentralProteusClientOrError(central: CoreCryptoCentral): ProteusClient {
+        return try {
+            central.proteusClient()
+        } catch (exception: ProteusStorageMigrationException) {
+            proteusMigrationRecoveryHandler.clearClientData { removeLocalFiles() }
+            val logMap = mapOf(
+                "userId" to userId.value.obfuscateId(),
+                "exception" to exception,
+                "message" to exception.message,
+                "stackTrace" to exception.stackTraceToString()
+            )
+            kaliumLogger.withTextTag(TAG).logStructuredJson(KaliumLogLevel.ERROR, "Proteus storage migration failed", logMap)
+            throw exception
+        }
+    }
+
+    /**
+     * Actually deletes the proteus local files.
+     * Important! It is the caller responsibility to use the mutex, DON'T add a mutex here or it will be dead lock it.
+     */
+    private suspend fun removeLocalFiles() {
+        withContext(dispatcher.io) {
+            _proteusClient?.close()
+            _proteusClient = null
+            FileUtil.deleteDirectory(rootProteusPath)
+        }
+    }
+
     private companion object {
         const val TAG = "ProteusClientProvider"
     }

logic/src/commonMain/kotlin/com/wire/kalium/logic/data/client/ProteusMigrationRecoveryHandler.kt

@@ -0,0 +1,30 @@
+/*
+ * Wire
+ * Copyright (C) 2024 Wire Swiss GmbH
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program. If not, see http://www.gnu.org/licenses/.
+ */
+package com.wire.kalium.logic.data.client
+
+import com.wire.kalium.logic.data.logout.LogoutReason
+
+/**
+ * Handles the migration error of a proteus client storage from CryptoBox to CoreCrypto.
+ * It will perform a logout, using [LogoutReason.MIGRATION_TO_CC_FAILED] as the reason.
+ *
+ * This achieves that the client data is cleared and the user is logged out without losing content.
+ */
+interface ProteusMigrationRecoveryHandler {
+    suspend fun clearClientData(clearLocalFiles: suspend () -> Unit)
+}

logic/src/commonMain/kotlin/com/wire/kalium/logic/data/session/SessionMapper.kt

@@ -107,6 +107,7 @@ internal class SessionMapperImpl : SessionMapper {
             LogoutReason.REMOVED_CLIENT -> LogoutReasonEntity.REMOVED_CLIENT
             LogoutReason.DELETED_ACCOUNT -> LogoutReasonEntity.DELETED_ACCOUNT
             LogoutReason.SESSION_EXPIRED -> LogoutReasonEntity.SESSION_EXPIRED
+            LogoutReason.MIGRATION_TO_CC_FAILED -> LogoutReasonEntity.MIGRATION_TO_CC_FAILED
         }
 
     override fun toSsoIdEntity(ssoId: SsoId?): SsoIdEntity? =
@@ -140,6 +141,7 @@ internal class SessionMapperImpl : SessionMapper {
             LogoutReasonEntity.REMOVED_CLIENT -> LogoutReason.REMOVED_CLIENT
             LogoutReasonEntity.DELETED_ACCOUNT -> LogoutReason.DELETED_ACCOUNT
             LogoutReasonEntity.SESSION_EXPIRED -> LogoutReason.SESSION_EXPIRED
+            LogoutReasonEntity.MIGRATION_TO_CC_FAILED -> LogoutReason.MIGRATION_TO_CC_FAILED
         }
 
     override fun fromEntityToProxyCredentialsDTO(proxyCredentialsEntity: ProxyCredentialsEntity): ProxyCredentialsDTO =

logic/src/commonMain/kotlin/com/wire/kalium/logic/feature/UserSessionScope.kt

@@ -51,6 +51,7 @@ import com.wire.kalium.logic.data.client.MLSClientProvider
 import com.wire.kalium.logic.data.client.MLSClientProviderImpl
 import com.wire.kalium.logic.data.client.ProteusClientProvider
 import com.wire.kalium.logic.data.client.ProteusClientProviderImpl
+import com.wire.kalium.logic.data.client.ProteusMigrationRecoveryHandler
 import com.wire.kalium.logic.data.client.remote.ClientRemoteDataSource
 import com.wire.kalium.logic.data.client.remote.ClientRemoteRepository
 import com.wire.kalium.logic.data.connection.ConnectionDataSource
@@ -187,6 +188,7 @@ import com.wire.kalium.logic.feature.client.IsAllowedToRegisterMLSClientUseCase
 import com.wire.kalium.logic.feature.client.IsAllowedToRegisterMLSClientUseCaseImpl
 import com.wire.kalium.logic.feature.client.MLSClientManager
 import com.wire.kalium.logic.feature.client.MLSClientManagerImpl
+import com.wire.kalium.logic.feature.client.ProteusMigrationRecoveryHandlerImpl
 import com.wire.kalium.logic.feature.client.RegisterMLSClientUseCase
 import com.wire.kalium.logic.feature.client.RegisterMLSClientUseCaseImpl
 import com.wire.kalium.logic.feature.connection.ConnectionScope
@@ -623,12 +625,17 @@ class UserSessionScope internal constructor(
     private val updateKeyingMaterialThresholdProvider: UpdateKeyingMaterialThresholdProvider
         get() = UpdateKeyingMaterialThresholdProviderImpl(kaliumConfigs)
 
+    private val proteusMigrationRecoveryHandler: ProteusMigrationRecoveryHandler by lazy {
+        ProteusMigrationRecoveryHandlerImpl(lazy { logout })
+    }
+
     val proteusClientProvider: ProteusClientProvider by lazy {
         ProteusClientProviderImpl(
             rootProteusPath = rootPathsProvider.rootProteusPath(userId),
             userId = userId,
             passphraseStorage = globalPreferences.passphraseStorage,
-            kaliumConfigs = kaliumConfigs
+            kaliumConfigs = kaliumConfigs,
+            proteusMigrationRecoveryHandler = proteusMigrationRecoveryHandler
         )
     }
 
@@ -918,11 +925,12 @@ class UserSessionScope internal constructor(
             kaliumFileSystem = kaliumFileSystem
         )
 
-    private val eventGatherer: EventGatherer get() = EventGathererImpl(
-        eventRepository = eventRepository,
-        incrementalSyncRepository = incrementalSyncRepository,
-        logger = userScopedLogger
-    )
+    private val eventGatherer: EventGatherer
+        get() = EventGathererImpl(
+            eventRepository = eventRepository,
+            incrementalSyncRepository = incrementalSyncRepository,
+            logger = userScopedLogger
+        )
 
     private val eventProcessor: EventProcessor by lazy {
         EventProcessorImpl(

logic/src/commonMain/kotlin/com/wire/kalium/logic/feature/auth/LogoutUseCase.kt

@@ -31,6 +31,7 @@ import com.wire.kalium.logic.feature.call.usecase.ObserveEstablishedCallsUseCase
 import com.wire.kalium.logic.feature.client.ClearClientDataUseCase
 import com.wire.kalium.logic.feature.session.DeregisterTokenUseCase
 import com.wire.kalium.logic.featureFlags.KaliumConfigs
+import com.wire.kalium.logic.kaliumLogger
 import com.wire.kalium.logic.sync.UserSessionWorkScheduler
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.cancel
@@ -106,6 +107,9 @@ internal class LogoutUseCaseImpl @Suppress("LongParameterList") constructor(
                 }
 
                 LogoutReason.SELF_SOFT_LOGOUT -> clearCurrentClientIdAndFirebaseTokenFlag()
+                LogoutReason.MIGRATION_TO_CC_FAILED -> prepareForCoreCryptoMigrationRecovery()
+            }.also {
+                kaliumLogger.withTextTag(TAG).d("Logout reason: $reason")
             }
 
             userConfigRepository.clearE2EISettings()
@@ -115,6 +119,13 @@ internal class LogoutUseCaseImpl @Suppress("LongParameterList") constructor(
         }.let { if (waitUntilCompletes) it.join() else it }
     }
 
+    private suspend fun prepareForCoreCryptoMigrationRecovery() {
+        clearClientDataUseCase()
+        logoutRepository.clearClientRelatedLocalMetadata()
+        clientRepository.clearRetainedClientId()
+        pushTokenRepository.setUpdateFirebaseTokenFlag(true)
+    }
+
     private suspend fun clearCurrentClientIdAndFirebaseTokenFlag() {
         clientRepository.clearCurrentClientId()
         clientRepository.clearNewClients()
@@ -146,5 +157,6 @@ internal class LogoutUseCaseImpl @Suppress("LongParameterList") constructor(
 
     companion object {
         const val CLEAR_DATA_DELAY = 1000L
+        const val TAG = "LogoutUseCase"
     }
 }

logic/src/commonMain/kotlin/com/wire/kalium/logic/feature/client/ProteusMigrationRecoveryHandlerImpl.kt

@@ -0,0 +1,52 @@
+/*
+ * Wire
+ * Copyright (C) 2024 Wire Swiss GmbH
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program. If not, see http://www.gnu.org/licenses/.
+ */
+package com.wire.kalium.logic.feature.client
+
+import com.wire.kalium.logic.data.client.ProteusMigrationRecoveryHandler
+import com.wire.kalium.logic.data.logout.LogoutReason
+import com.wire.kalium.logic.feature.auth.LogoutUseCase
+import com.wire.kalium.logic.kaliumLogger
+
+internal class ProteusMigrationRecoveryHandlerImpl(
+    private val logoutUseCase: Lazy<LogoutUseCase>
+) : ProteusMigrationRecoveryHandler {
+
+    /**
+     * Handles the migration error of a proteus client storage from CryptoBox to CoreCrypto.
+     * It will perform a logout, using [LogoutReason.MIGRATION_TO_CC_FAILED] as the reason.
+     *
+     * This achieves that the client data is cleared and the user is logged out without losing content.
+     */
+    @Suppress("TooGenericExceptionCaught")
+    override suspend fun clearClientData(clearLocalFiles: suspend () -> Unit) {
+        try {
+            kaliumLogger.withTextTag(TAG).i("Starting the recovery from failed Proteus storage migration")
+            clearLocalFiles()
+            logoutUseCase.value(LogoutReason.MIGRATION_TO_CC_FAILED, true)
+        } catch (e: Exception) {
+            kaliumLogger.withTextTag(TAG).e("Fatal, error while clearing client data: $e")
+            throw e
+        } finally {
+            kaliumLogger.withTextTag(TAG).i("Finished the recovery from failed Proteus storage migration")
+        }
+    }
+
+    private companion object {
+        const val TAG = "ProteusMigrationRecoveryHandler"
+    }
+}