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
@@ -108,6 +109,7 @@ 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)
 }
 
 private enum class CommitStrategy {
@@ -514,6 +516,26 @@ internal class MLSConversationDataSource(
             wrapMLSRequest { mlsClient.isGroupVerified(idMapper.toCryptoModel(groupID)) }
         }.map { it.toModel() }
 
+    override suspend fun rotateKeysAndMigrateConversations(clientId: ClientId, e2eiClient: E2EIClient, certificateChain: String) {
+        mlsClientProvider.getMLSClient().map { mlsClient ->
+            wrapMLSRequest {
+                mlsClient.e2eiRotateAll(e2eiClient, certificateChain, 10U)
+            }.map { rotateBundle ->
+                // todo: make below API calls atomic when the backend does it in one request
+                kaliumLogger.w("drop old key packages after conversations migration")
+                keyPackageRepository.deleteKeyPackages(clientId, rotateBundle.keyPackageRefsToRemove)
+
+                kaliumLogger.w("upload new key packages including x509 certificate")
+                keyPackageRepository.uploadKeyPackages(clientId, rotateBundle.newKeyPackages)
+
+                kaliumLogger.w("send migration commits after key rotations")
+                rotateBundle.commits.forEach {
+                    sendCommitBundle(GroupID(it.key), it.value)
+                }
+            }
+        }
+    }
+
     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
@@ -60,14 +61,16 @@ interface E2EIRepository {
     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)
 }
 
 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 {
@@ -198,4 +201,12 @@ class E2EIRepositoryImpl(
             }
         }
     }
+
+    override suspend fun rotateKeysAndMigrateConversations(certificateChain: String) {
+        e2EIClientProvider.getE2EIClient().flatMap { e2eiClient ->
+            currentClientIdProvider().map { 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/e2ei/EnrollE2EIUseCase.kt

@@ -111,8 +111,8 @@ 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())
+
         return Either.Right(E2EIEnrollmentResult.Success(certificateRequest.response.decodeToString()))
     }