feat(e2ei): migrate mls-conversation after enrolling E2EI certificate (WPB-2928)

cryptography/src/appleMain/kotlin/com/wire/kalium/cryptography/MLSClientImpl.kt

@@ -194,7 +194,11 @@ class MLSClientImpl(
         TODO("Not yet implemented")
     }
 
-    override suspend fun e2eiRotateAll(enrollment: E2EIClient, certificateChain: CertificateChain, newMLSKeyPackageCount: UInt) {
+    override suspend fun e2eiRotateAll(
+        enrollment: E2EIClient,
+        certificateChain: CertificateChain,
+        newMLSKeyPackageCount: UInt
+    ): RotateBundle {
         TODO("Not yet implemented")
     }
 

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

@@ -260,11 +260,13 @@ class MLSClientImpl(
         enrollment: E2EIClient,
         certificateChain: CertificateChain,
         newMLSKeyPackageCount: UInt
-    ) {
-        coreCrypto.e2eiRotateAll(
-            (enrollment as E2EIClientImpl).wireE2eIdentity,
-            certificateChain,
-            newMLSKeyPackageCount
+    ): RotateBundle {
+        return toRotateBundle(
+            coreCrypto.e2eiRotateAll(
+                (enrollment as E2EIClientImpl).wireE2eIdentity,
+                certificateChain,
+                newMLSKeyPackageCount
+            )
         )
     }
 

cryptography/src/commonMain/kotlin/com/wire/kalium/cryptography/MLSClient.kt

@@ -326,7 +326,7 @@ interface MLSClient {
     /**
      * Generate new keypackages after E2EI certificate issued
      */
-    suspend fun e2eiRotateAll(enrollment: E2EIClient, certificateChain: CertificateChain, newMLSKeyPackageCount: UInt)
+    suspend fun e2eiRotateAll(enrollment: E2EIClient, certificateChain: CertificateChain, newMLSKeyPackageCount: UInt): RotateBundle
 
     /**
      * Conversation E2EI Verification Status

cryptography/src/jsMain/kotlin/com/wire/kalium/cryptography/MLSClientImpl.kt

@@ -136,7 +136,7 @@ class MLSClientImpl : MLSClient {
         enrollment: E2EIClient,
         certificateChain: CertificateChain,
         newMLSKeyPackageCount: UInt
-    ) {
+    ): RotateBundle {
         TODO("Not yet implemented")
     }
 

logic/src/commonMain/kotlin/com/wire/kalium/logic/data/conversation/MLSConversationRepository.kt

@@ -21,6 +21,7 @@ package com.wire.kalium.logic.data.conversation
 import com.wire.kalium.cryptography.CommitBundle
 import com.wire.kalium.cryptography.CryptoQualifiedClientId
 import com.wire.kalium.cryptography.CryptoQualifiedID
+import com.wire.kalium.cryptography.E2EIClient
 import com.wire.kalium.logger.obfuscateId
 import com.wire.kalium.logic.CoreFailure
 import com.wire.kalium.logic.MLSFailure
@@ -43,6 +44,8 @@ import com.wire.kalium.logic.functional.Either
 import com.wire.kalium.logic.functional.flatMap
 import com.wire.kalium.logic.functional.flatMapLeft
 import com.wire.kalium.logic.functional.flatten
+import com.wire.kalium.logic.functional.fold
+import com.wire.kalium.logic.functional.foldToEitherWhileRight
 import com.wire.kalium.logic.functional.map
 import com.wire.kalium.logic.functional.onFailure
 import com.wire.kalium.logic.functional.onSuccess
@@ -108,6 +111,11 @@ interface MLSConversationRepository {
     suspend fun observeProposalTimers(): Flow<ProposalTimer>
     suspend fun observeEpochChanges(): Flow<GroupID>
     suspend fun getConversationVerificationStatus(groupID: GroupID): Either<CoreFailure, Conversation.VerificationStatus>
+    suspend fun rotateKeysAndMigrateConversations(
+        clientId: ClientId,
+        e2eiClient: E2EIClient,
+        certificateChain: String
+    ): Either<CoreFailure, Unit>
 }
 
 private enum class CommitStrategy {
@@ -514,6 +522,33 @@ internal class MLSConversationDataSource(
             wrapMLSRequest { mlsClient.isGroupVerified(idMapper.toCryptoModel(groupID)) }
         }.map { it.toModel() }
 
+    override suspend fun rotateKeysAndMigrateConversations(
+        clientId: ClientId,
+        e2eiClient: E2EIClient,
+        certificateChain: String
+    ) = mlsClientProvider.getMLSClient().flatMap { mlsClient ->
+        wrapMLSRequest {
+            mlsClient.e2eiRotateAll(e2eiClient, certificateChain, 10U)
+        }.map { rotateBundle ->
+            // todo: make below API calls atomic when the backend does it in one request
+            // todo: store keypackages to drop, later drop them again
+            kaliumLogger.w("drop old key packages after conversations migration")
+            keyPackageRepository.deleteKeyPackages(clientId, rotateBundle.keyPackageRefsToRemove).flatMapLeft {
+                return Either.Left(it)
+            }
+
+            kaliumLogger.w("upload new key packages including x509 certificate")
+            keyPackageRepository.uploadKeyPackages(clientId, rotateBundle.newKeyPackages).flatMapLeft {
+                return Either.Left(it)
+            }
+
+            kaliumLogger.w("send migration commits after key rotations")
+            rotateBundle.commits.map {
+                sendCommitBundle(GroupID(it.key), it.value)
+            }.foldToEitherWhileRight(Unit) { value, _ -> value }.fold({ return Either.Left(it) }, { })
+        }
+    }
+
     private suspend fun retryOnCommitFailure(
         groupID: GroupID,
         retryOnClientMismatch: Boolean = true,

logic/src/commonMain/kotlin/com/wire/kalium/logic/data/e2ei/E2EIRepository.kt

@@ -26,6 +26,7 @@ import com.wire.kalium.cryptography.NewAcmeOrder
 import com.wire.kalium.logic.CoreFailure
 import com.wire.kalium.logic.data.client.E2EIClientProvider
 import com.wire.kalium.logic.data.client.MLSClientProvider
+import com.wire.kalium.logic.data.conversation.MLSConversationRepository
 import com.wire.kalium.logic.feature.CurrentClientIdProvider
 import com.wire.kalium.logic.functional.Either
 import com.wire.kalium.logic.functional.flatMap
@@ -59,15 +60,16 @@ interface E2EIRepository {
     suspend fun finalize(location: String, prevNonce: String): Either<CoreFailure, Pair<ACMEResponse, String>>
     suspend fun checkOrderRequest(location: String, prevNonce: String): Either<CoreFailure, Pair<ACMEResponse, String>>
     suspend fun certificateRequest(location: String, prevNonce: String): Either<CoreFailure, ACMEResponse>
-    suspend fun initMLSClientWithCertificate(certificateChain: String)
+    suspend fun rotateKeysAndMigrateConversations(certificateChain: String): Either<CoreFailure, Unit>
 }
 
 class E2EIRepositoryImpl(
     private val e2EIApi: E2EIApi,
     private val acmeApi: ACMEApi,
     private val e2EIClientProvider: E2EIClientProvider,
     private val mlsClientProvider: MLSClientProvider,
-    private val currentClientIdProvider: CurrentClientIdProvider
+    private val currentClientIdProvider: CurrentClientIdProvider,
+    private val mlsConversationRepository: MLSConversationRepository
 ) : E2EIRepository {
 
     override suspend fun loadACMEDirectories(): Either<CoreFailure, AcmeDirectory> = wrapApiRequest {
@@ -191,11 +193,11 @@ class E2EIRepositoryImpl(
             }.map { it }
         }
 
-    override suspend fun initMLSClientWithCertificate(certificateChain: String) {
+    override suspend fun rotateKeysAndMigrateConversations(certificateChain: String) =
         e2EIClientProvider.getE2EIClient().flatMap { e2eiClient ->
-            mlsClientProvider.getMLSClient().map {
-                it.e2eiMlsInitOnly(e2eiClient, certificateChain)
+            currentClientIdProvider().flatMap { clientId ->
+                mlsConversationRepository.rotateKeysAndMigrateConversations(clientId, e2eiClient, certificateChain)
             }
         }
-    }
+
 }

logic/src/commonMain/kotlin/com/wire/kalium/logic/data/keypackage/KeyPackageRepository.kt

@@ -43,6 +43,10 @@ interface KeyPackageRepository {
 
     suspend fun uploadNewKeyPackages(clientId: ClientId, amount: Int = 100): Either<CoreFailure, Unit>
 
+    suspend fun uploadKeyPackages(clientId: ClientId, keyPackages: List<ByteArray>): Either<CoreFailure, Unit>
+
+    suspend fun deleteKeyPackages(clientId: ClientId, keyPackages: List<ByteArray>): Either<CoreFailure, Unit>
+
     suspend fun getAvailableKeyPackageCount(clientId: ClientId): Either<NetworkFailure, KeyPackageCountDTO>
 
     suspend fun validKeyPackageCount(clientId: ClientId): Either<CoreFailure, Int>
@@ -87,6 +91,22 @@ class KeyPackageDataSource(
             }
         }
 
+    override suspend fun uploadKeyPackages(
+        clientId: ClientId,
+        keyPackages: List<ByteArray>
+    ): Either<CoreFailure, Unit> =
+        wrapApiRequest {
+            keyPackageApi.uploadKeyPackages(clientId.value, keyPackages.map { it.encodeBase64() })
+        }
+
+    override suspend fun deleteKeyPackages(
+        clientId: ClientId,
+        keyPackages: List<ByteArray>
+    ): Either<CoreFailure, Unit> =
+        wrapApiRequest {
+            keyPackageApi.deleteKeyPackages(clientId.value, keyPackages.map { it.encodeBase64() })
+        }
+
     override suspend fun validKeyPackageCount(clientId: ClientId): Either<CoreFailure, Int> =
         mlsClientProvider.getMLSClient(clientId).flatMap { mlsClient ->
             wrapMLSRequest {

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

@@ -592,7 +592,8 @@ class UserSessionScope internal constructor(
             globalScope.unboundNetworkContainer.acmeApi,
             e2EIClientProvider,
             mlsClientProvider,
-            clientIdProvider
+            clientIdProvider,
+            mlsConversationRepository
         )
 
     private val e2EIClientProvider: E2EIClientProvider by lazy {

logic/src/commonMain/kotlin/com/wire/kalium/logic/feature/e2ei/usecase/EnrollE2EIUseCase.kt

@@ -22,6 +22,7 @@ import com.wire.kalium.logic.E2EIFailure
 import com.wire.kalium.logic.data.e2ei.E2EIRepository
 import com.wire.kalium.logic.functional.Either
 import com.wire.kalium.logic.functional.fold
+import com.wire.kalium.logic.functional.onFailure
 
 /**
  * Issue an E2EI certificate and re-initiate the MLSClient
@@ -111,8 +112,10 @@ class EnrollE2EIUseCaseImpl internal constructor(
             return E2EIEnrollmentResult.Failed(E2EIEnrollmentResult.E2EIStep.Certificate, it).toEitherLeft()
         }, { it })
 
-        // TODO(fix): init after fixing the MLS client initialization mechanism
-        // TODO(revert): e2EIRepository.initMLSClientWithCertificate(certificateRequest.response.decodeToString())
+        e2EIRepository.rotateKeysAndMigrateConversations(certificateRequest.response.decodeToString()).onFailure {
+            return E2EIEnrollmentResult.Failed(E2EIEnrollmentResult.E2EIStep.ConversationMigration, it).toEitherLeft()
+        }
+
         return Either.Right(E2EIEnrollmentResult.Success(certificateRequest.response.decodeToString()))
     }
 
@@ -132,6 +135,7 @@ sealed interface E2EIEnrollmentResult {
         OIDCChallenge,
         CheckOrderRequest,
         FinalizeRequest,
+        ConversationMigration,
         Certificate
     }