fix(mls): set removal-keys for 1on1 calls from conversation-response …

…(WPB-10743) 🍒 (#3019) * Commit with unresolved merge conflicts * fix merge conflicts * fix merge conflicts * fix merge conflicts * fix merge conflicts * fix stalling test * fix tests * fix tests --------- Co-authored-by: Mojtaba Chenani <[email protected]> Co-authored-by: Mohamad Jaara <[email protected]>
wireapp · Sep 26, 2024 · 1dd7cb2 · 1dd7cb2
1 parent 8268f21
commit 1dd7cb2
Show file tree

Hide file tree

Showing 11 changed files with 226 additions and 42 deletions.
diff --git a/.../commonMain/kotlin/com/wire/kalium/logic/data/conversation/ConversationGroupRepository.kt b/.../commonMain/kotlin/com/wire/kalium/logic/data/conversation/ConversationGroupRepository.kt
@@ -154,6 +154,7 @@ internal class ConversationGroupRepositoryImpl(
         val conversationEntity = conversationMapper.fromApiModelToDaoModel(
             conversationResponse, mlsGroupState = ConversationEntity.GroupState.PENDING_CREATION, selfTeamId
         )
+        val mlsPublicKeys = conversationMapper.fromApiModel(conversationResponse.publicKeys)
         val protocol = protocolInfoMapper.fromEntity(conversationEntity.protocolInfo)
 
         return wrapStorageRequest {
@@ -166,7 +167,8 @@ internal class ConversationGroupRepositoryImpl(
                 is Conversation.ProtocolInfo.MLSCapable -> mlsConversationRepository.establishMLSGroup(
                     groupID = protocol.groupId,
                     members = usersList + selfUserId,
-                    allowSkippingUsersWithoutKeyPackages = true
+                    publicKeys = mlsPublicKeys,
+                    allowSkippingUsersWithoutKeyPackages = true,
                 ).map { it.notAddedUsers }
             }
         }.flatMap { protocolSpecificAdditionFailures ->

diff --git a/logic/src/commonMain/kotlin/com/wire/kalium/logic/data/conversation/ConversationMapper.kt b/logic/src/commonMain/kotlin/com/wire/kalium/logic/data/conversation/ConversationMapper.kt
@@ -28,6 +28,7 @@ import com.wire.kalium.logic.data.id.toApi
 import com.wire.kalium.logic.data.id.toDao
 import com.wire.kalium.logic.data.id.toModel
 import com.wire.kalium.logic.data.message.MessagePreview
+import com.wire.kalium.logic.data.mlspublickeys.MLSPublicKeys
 import com.wire.kalium.logic.data.user.AvailabilityStatusMapper
 import com.wire.kalium.logic.data.user.BotService
 import com.wire.kalium.logic.data.user.Connection
@@ -41,6 +42,7 @@ import com.wire.kalium.network.api.authenticated.conversation.ConvTeamInfo
 import com.wire.kalium.network.api.authenticated.conversation.ConversationResponse
 import com.wire.kalium.network.api.authenticated.conversation.CreateConversationRequest
 import com.wire.kalium.network.api.authenticated.conversation.ReceiptMode
+import com.wire.kalium.network.api.authenticated.serverpublickey.MLSPublicKeysDTO
 import com.wire.kalium.network.api.model.ConversationAccessDTO
 import com.wire.kalium.network.api.model.ConversationAccessRoleDTO
 import com.wire.kalium.persistence.dao.conversation.ConversationEntity
@@ -59,6 +61,7 @@ import kotlin.time.toDuration
 
 interface ConversationMapper {
     fun fromApiModelToDaoModel(apiModel: ConversationResponse, mlsGroupState: GroupState?, selfUserTeamId: TeamId?): ConversationEntity
+    fun fromApiModel(mlsPublicKeysDTO: MLSPublicKeysDTO?): MLSPublicKeys?
     fun fromDaoModel(daoModel: ConversationViewEntity): Conversation
     fun fromDaoModel(daoModel: ConversationEntity): Conversation
     fun fromDaoModelToDetails(
@@ -136,6 +139,12 @@ internal class ConversationMapperImpl(
         legalHoldStatus = ConversationEntity.LegalHoldStatus.DISABLED
     )
 
+    override fun fromApiModel(mlsPublicKeysDTO: MLSPublicKeysDTO?) = mlsPublicKeysDTO?.let {
+        MLSPublicKeys(
+            removal = mlsPublicKeysDTO.removal
+        )
+    }
+
     override fun fromDaoModel(daoModel: ConversationViewEntity): Conversation = with(daoModel) {
         val lastReadDateEntity = if (type == ConversationEntity.Type.CONNECTION_PENDING) Instant.UNIX_FIRST_DATE
         else lastReadDate

diff --git a/...rc/commonMain/kotlin/com/wire/kalium/logic/data/conversation/MLSConversationRepository.kt b/...rc/commonMain/kotlin/com/wire/kalium/logic/data/conversation/MLSConversationRepository.kt
@@ -45,7 +45,9 @@ import com.wire.kalium.logic.data.id.toModel
 import com.wire.kalium.logic.data.keypackage.KeyPackageLimitsProvider
 import com.wire.kalium.logic.data.keypackage.KeyPackageRepository
 import com.wire.kalium.logic.data.mls.CipherSuite
+import com.wire.kalium.logic.data.mlspublickeys.MLSPublicKeys
 import com.wire.kalium.logic.data.mlspublickeys.MLSPublicKeysRepository
+import com.wire.kalium.logic.data.mlspublickeys.getRemovalKey
 import com.wire.kalium.logic.data.user.UserId
 import com.wire.kalium.logic.di.MapperProvider
 import com.wire.kalium.logic.data.e2ei.RevocationListChecker
@@ -123,6 +125,7 @@ interface MLSConversationRepository {
     suspend fun establishMLSGroup(
         groupID: GroupID,
         members: List<UserId>,
+        publicKeys: MLSPublicKeys? = null,
         allowSkippingUsersWithoutKeyPackages: Boolean = false
     ): Either<CoreFailure, MLSAdditionResult>
 
@@ -554,16 +557,18 @@ internal class MLSConversationDataSource(
     override suspend fun establishMLSGroup(
         groupID: GroupID,
         members: List<UserId>,
-        allowSkippingUsersWithoutKeyPackages: Boolean,
+        publicKeys: MLSPublicKeys?,
+        allowSkippingUsersWithoutKeyPackages: Boolean
     ): Either<CoreFailure, MLSAdditionResult> = withContext(serialDispatcher) {
-        mlsClientProvider.getMLSClient().flatMap<MLSAdditionResult, CoreFailure, MLSClient> {
-            mlsPublicKeysRepository.getKeyForCipherSuite(
-                CipherSuite.fromTag(it.getDefaultCipherSuite())
-            ).flatMap { key ->
+        mlsClientProvider.getMLSClient().flatMap<MLSAdditionResult, CoreFailure, MLSClient> { mlsClient ->
+            val cipherSuite = CipherSuite.fromTag(mlsClient.getDefaultCipherSuite())
+            val keys = publicKeys?.getRemovalKey(cipherSuite) ?: mlsPublicKeysRepository.getKeyForCipherSuite(cipherSuite)
+
+            keys.flatMap { externalSenders ->
                 establishMLSGroup(
                     groupID = groupID,
                     members = members,
-                    externalSenders = key,
+                    externalSenders = externalSenders,
                     allowPartialMemberList = allowSkippingUsersWithoutKeyPackages
                 )
             }

diff --git a/...src/commonMain/kotlin/com/wire/kalium/logic/data/mlspublickeys/MLSPublicKeysRepository.kt b/...src/commonMain/kotlin/com/wire/kalium/logic/data/mlspublickeys/MLSPublicKeysRepository.kt
@@ -34,6 +34,15 @@ data class MLSPublicKeys(
     val removal: Map<String, String>?
 )
 
+fun MLSPublicKeys.getRemovalKey(cipherSuite: CipherSuite): Either<CoreFailure, ByteArray> {
+    val mlsPublicKeysMapper: MLSPublicKeysMapper = MapperProvider.mlsPublicKeyMapper()
+    val keySignature = mlsPublicKeysMapper.fromCipherSuite(cipherSuite)
+    val key = this.removal?.let { removalKeys ->
+        removalKeys[keySignature.value]
+    } ?: return Either.Left(MLSFailure.Generic(IllegalStateException("No key found for cipher suite $cipherSuite")))
+    return key.decodeBase64Bytes().right()
+}
+
 interface MLSPublicKeysRepository {
     suspend fun fetchKeys(): Either<CoreFailure, MLSPublicKeys>
     suspend fun getKeys(): Either<CoreFailure, MLSPublicKeys>
@@ -42,7 +51,6 @@ interface MLSPublicKeysRepository {
 
 class MLSPublicKeysRepositoryImpl(
     private val mlsPublicKeyApi: MLSPublicKeyApi,
-    private val mlsPublicKeysMapper: MLSPublicKeysMapper = MapperProvider.mlsPublicKeyMapper()
 ) : MLSPublicKeysRepository {
 
     // TODO: make it thread safe
@@ -60,14 +68,8 @@ class MLSPublicKeysRepositoryImpl(
     }
 
     override suspend fun getKeyForCipherSuite(cipherSuite: CipherSuite): Either<CoreFailure, ByteArray> {
-
         return getKeys().flatMap { serverPublicKeys ->
-            val keySignature = mlsPublicKeysMapper.fromCipherSuite(cipherSuite)
-            val key = serverPublicKeys.removal?.let { removalKeys ->
-                removalKeys[keySignature.value]
-            } ?: return Either.Left(MLSFailure.Generic(IllegalStateException("No key found for cipher suite $cipherSuite")))
-            key.decodeBase64Bytes().right()
+            serverPublicKeys.getRemovalKey(cipherSuite)
         }
     }
-
 }
diff --git a/...monTest/kotlin/com/wire/kalium/logic/data/conversation/ConversationGroupRepositoryTest.kt b/...monTest/kotlin/com/wire/kalium/logic/data/conversation/ConversationGroupRepositoryTest.kt
@@ -422,7 +422,7 @@ class ConversationGroupRepositoryTest {
             }.wasInvoked(once)
 
             coVerify {
-                mlsConversationRepository.establishMLSGroup(any(), any(), eq(true))
+                mlsConversationRepository.establishMLSGroup(any(), any(), any(), eq(true))
             }.wasInvoked(once)
 
             coVerify {
@@ -465,7 +465,7 @@ class ConversationGroupRepositoryTest {
                 }.wasInvoked(once)
 
                 coVerify {
-                    mlsConversationRepository.establishMLSGroup(any(), any(), eq(true))
+                    mlsConversationRepository.establishMLSGroup(any(), any(), any(), eq(true))
                 }.wasInvoked(once)
 
                 coVerify {
@@ -1723,11 +1723,10 @@ class ConversationGroupRepositoryTest {
                 legalHoldHandler
             )
 
-        suspend fun withMlsConversationEstablished(additionResult: MLSAdditionResult): Arrangement {
+        suspend fun withMlsConversationEstablished(additionResult: MLSAdditionResult): Arrangement = apply {
             coEvery {
-                mlsConversationRepository.establishMLSGroup(any(), any(), any())
+                mlsConversationRepository.establishMLSGroup(any(), any(), any(), any())
             }.returns(Either.Right(additionResult))
-            return this
         }
 
         /**

diff --git a/...ommonTest/kotlin/com/wire/kalium/logic/data/conversation/MLSConversationRepositoryTest.kt b/...ommonTest/kotlin/com/wire/kalium/logic/data/conversation/MLSConversationRepositoryTest.kt
@@ -40,6 +40,7 @@ import com.wire.kalium.logic.data.conversation.MLSConversationRepositoryTest.Arr
 import com.wire.kalium.logic.data.conversation.MLSConversationRepositoryTest.Arrangement.Companion.CRYPTO_CLIENT_ID
 import com.wire.kalium.logic.data.conversation.MLSConversationRepositoryTest.Arrangement.Companion.E2EI_CONVERSATION_CLIENT_INFO_ENTITY
 import com.wire.kalium.logic.data.conversation.MLSConversationRepositoryTest.Arrangement.Companion.KEY_PACKAGE
+import com.wire.kalium.logic.data.conversation.MLSConversationRepositoryTest.Arrangement.Companion.MLS_PUBLIC_KEY
 import com.wire.kalium.logic.data.conversation.MLSConversationRepositoryTest.Arrangement.Companion.ROTATE_BUNDLE
 import com.wire.kalium.logic.data.conversation.MLSConversationRepositoryTest.Arrangement.Companion.TEST_FAILURE
 import com.wire.kalium.logic.data.conversation.MLSConversationRepositoryTest.Arrangement.Companion.WIRE_IDENTITY
@@ -98,26 +99,30 @@ import io.mockative.matches
 import io.mockative.mock
 import io.mockative.once
 import io.mockative.twice
+import io.mockative.verify
 import kotlinx.coroutines.async
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.yield
 import kotlinx.datetime.Instant
+import kotlin.test.BeforeTest
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertIs
 
 class MLSConversationRepositoryTest {
 
+    @BeforeTest
+
     @Test
-    fun givenCommitMessage_whenDecryptingMessage_thenEmitEpochChange() = runTest(TestKaliumDispatcher.default) {
+    fun givenCommitMessage_whenDecryptingMessage_thenEmitEpochChange() = runTest {
         val (arrangement, mlsConversationRepository) = Arrangement(testKaliumDispatcher)
             .withGetMLSClientSuccessful()
             .withDecryptMLSMessageSuccessful(Arrangement.DECRYPTED_MESSAGE_BUNDLE)
             .arrange()
 
-        val epochChange = async(TestKaliumDispatcher.default) {
+        val epochChange = async() {
             arrangement.epochsFlow.first()
         }
         yield()
@@ -168,7 +173,7 @@ class MLSConversationRepositoryTest {
             .withSendCommitBundleSuccessful()
             .arrange()
 
-        val result = mlsConversationRepository.establishMLSGroup(Arrangement.GROUP_ID, listOf(TestConversation.USER_1))
+        val result = mlsConversationRepository.establishMLSGroup(Arrangement.GROUP_ID, listOf(TestConversation.USER_1), publicKeys = null)
         result.shouldSucceed()
 
         coVerify {
@@ -280,6 +285,84 @@ class MLSConversationRepositoryTest {
         }.wasNotInvoked()
     }
 
+    @Test
+    fun givenPublicKeysIsNotNull_whenCallingEstablishMLSGroup_ThenGetPublicKeysRepositoryNotCalled() = runTest {
+        val (arrangement, mlsConversationRepository) = Arrangement(kaliumDispatcher = testKaliumDispatcher)
+            .withGetDefaultCipherSuite(CipherSuite.MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519)
+            .withCommitPendingProposalsReturningNothing()
+            .withClaimKeyPackagesSuccessful()
+            .withGetMLSClientSuccessful()
+            .withKeyForCipherSuite()
+            .withAddMLSMemberSuccessful()
+            .withSendCommitBundleSuccessful()
+            .arrange()
+
+        val result =
+            mlsConversationRepository.establishMLSGroup(Arrangement.GROUP_ID, listOf(TestConversation.USER_1), publicKeys = MLS_PUBLIC_KEY)
+        result.shouldSucceed()
+
+        coVerify {
+            arrangement.mlsClient.createConversation(
+                groupId = eq(Arrangement.RAW_GROUP_ID),
+                externalSenders = any())
+        }.wasInvoked(once)
+
+        coVerify {
+            arrangement.mlsClient.addMember(
+                groupId = eq(Arrangement.RAW_GROUP_ID),
+                membersKeyPackages = any())
+        }.wasInvoked(once)
+
+        coVerify {
+            arrangement.mlsMessageApi.sendCommitBundle(any<MLSMessageApi.CommitBundle>())
+        }.wasInvoked(once)
+
+        coVerify {
+            arrangement.mlsClient.commitAccepted(eq(Arrangement.RAW_GROUP_ID))
+        }.wasInvoked(once)
+
+        coVerify {
+            arrangement.mlsPublicKeysRepository.getKeyForCipherSuite(any())
+        }.wasNotInvoked()
+    }
+
+    @Test
+    fun givenPublicKeysIsNull_whenCallingEstablishMLSGroup_ThenGetPublicKeysRepositoryIsCalled() = runTest {
+        val (arrangement, mlsConversationRepository) = Arrangement(testKaliumDispatcher)
+            .withGetDefaultCipherSuite(CipherSuite.MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519)
+            .withCommitPendingProposalsReturningNothing()
+            .withClaimKeyPackagesSuccessful()
+            .withGetMLSClientSuccessful()
+            .withKeyForCipherSuite()
+            .withAddMLSMemberSuccessful()
+            .withSendCommitBundleSuccessful()
+            .arrange()
+
+        val result =
+            mlsConversationRepository.establishMLSGroup(Arrangement.GROUP_ID, listOf(TestConversation.USER_1), publicKeys = null)
+        result.shouldSucceed()
+
+        coVerify {
+            arrangement.mlsClient.createConversation(eq(Arrangement.RAW_GROUP_ID), any())
+        }.wasInvoked(once)
+
+        coVerify {
+            arrangement.mlsClient.addMember(eq(Arrangement.RAW_GROUP_ID), any())
+        }.wasInvoked(once)
+
+        coVerify {
+            arrangement.mlsMessageApi.sendCommitBundle(any<MLSMessageApi.CommitBundle>())
+        }.wasInvoked(once)
+
+        coVerify {
+            arrangement.mlsClient.commitAccepted(eq(Arrangement.RAW_GROUP_ID))
+        }.wasInvoked(once)
+
+        coVerify {
+            arrangement.mlsPublicKeysRepository.getKeyForCipherSuite(any())
+        }.wasInvoked(once)
+    }
+
     @Test
     fun givenNewCrlDistributionPoints_whenEstablishingMLSGroup_thenCheckRevocationList() = runTest {
         val (arrangement, mlsConversationRepository) = Arrangement(testKaliumDispatcher)
@@ -329,7 +412,7 @@ class MLSConversationRepositoryTest {
             .withWaitUntilLiveSuccessful()
             .arrange()
 
-        val result = mlsConversationRepository.establishMLSGroup(Arrangement.GROUP_ID, listOf(TestConversation.USER_1))
+        val result = mlsConversationRepository.establishMLSGroup(Arrangement.GROUP_ID, listOf(TestConversation.USER_1), publicKeys = null)
         result.shouldSucceed()
 
         coVerify {
@@ -357,7 +440,7 @@ class MLSConversationRepositoryTest {
             .withSendCommitBundleFailing(Arrangement.MLS_STALE_MESSAGE_ERROR)
             .arrange()
 
-        val result = mlsConversationRepository.establishMLSGroup(Arrangement.GROUP_ID, listOf(TestConversation.USER_1))
+        val result = mlsConversationRepository.establishMLSGroup(Arrangement.GROUP_ID, listOf(TestConversation.USER_1), publicKeys = null)
         result.shouldFail()
 
         coVerify {
@@ -385,7 +468,7 @@ class MLSConversationRepositoryTest {
             .withSendCommitBundleSuccessful()
             .arrange()
 
-        val result = mlsConversationRepository.establishMLSGroup(Arrangement.GROUP_ID, listOf(TestConversation.USER_1))
+        val result = mlsConversationRepository.establishMLSGroup(Arrangement.GROUP_ID, listOf(TestConversation.USER_1), publicKeys = null)
         result.shouldSucceed()
 
         coVerify {
@@ -410,7 +493,7 @@ class MLSConversationRepositoryTest {
             .withSendCommitBundleSuccessful()
             .arrange()
 
-        val result = mlsConversationRepository.establishMLSGroup(Arrangement.GROUP_ID, emptyList())
+        val result = mlsConversationRepository.establishMLSGroup(Arrangement.GROUP_ID, emptyList(), publicKeys = null)
         result.shouldSucceed()
 
         coVerify {

diff --git a/...tlin/com/wire/kalium/logic/feature/conversation/JoinExistingMLSConversationUseCaseTest.kt b/...tlin/com/wire/kalium/logic/feature/conversation/JoinExistingMLSConversationUseCaseTest.kt
@@ -131,7 +131,12 @@ class JoinExistingMLSConversationUseCaseTest {
             joinExistingMLSConversationsUseCase(Arrangement.MLS_UNESTABLISHED_GROUP_CONVERSATION.id).shouldSucceed()
 
             coVerify {
-                arrangement.mlsConversationRepository.establishMLSGroup(eq(Arrangement.GROUP_ID3), eq(emptyList()), any())
+                arrangement.mlsConversationRepository.establishMLSGroup(
+                    groupID = Arrangement.GROUP_ID3,
+                    members = emptyList(),
+                    publicKeys = null,
+                    allowSkippingUsersWithoutKeyPackages = false
+                )
             }.wasNotInvoked()
         }
 
@@ -148,7 +153,12 @@ class JoinExistingMLSConversationUseCaseTest {
             joinExistingMLSConversationsUseCase(Arrangement.MLS_UNESTABLISHED_SELF_CONVERSATION.id).shouldSucceed()
 
             coVerify {
-                arrangement.mlsConversationRepository.establishMLSGroup(eq(Arrangement.GROUP_ID_SELF), eq(emptyList()), any())
+                arrangement.mlsConversationRepository.establishMLSGroup(
+                    groupID = Arrangement.GROUP_ID_SELF,
+                    members = emptyList(),
+                    publicKeys = null,
+                    allowSkippingUsersWithoutKeyPackages = false
+                )
             }.wasInvoked(once)
         }
 
@@ -167,7 +177,12 @@ class JoinExistingMLSConversationUseCaseTest {
             joinExistingMLSConversationsUseCase(Arrangement.MLS_UNESTABLISHED_ONE_ONE_ONE_CONVERSATION.id).shouldSucceed()
 
             coVerify {
-                arrangement.mlsConversationRepository.establishMLSGroup(eq(Arrangement.GROUP_ID_ONE_ON_ONE), eq(members), any())
+                arrangement.mlsConversationRepository.establishMLSGroup(
+                    groupID = Arrangement.GROUP_ID_ONE_ON_ONE,
+                    members = members,
+                    publicKeys = null,
+                    allowSkippingUsersWithoutKeyPackages = false
+                )
             }.wasInvoked(once)
         }
 
@@ -256,7 +271,7 @@ class JoinExistingMLSConversationUseCaseTest {
 
         suspend fun withEstablishMLSGroupSuccessful(additionResult: MLSAdditionResult) = apply {
             coEvery {
-                mlsConversationRepository.establishMLSGroup(any(), any(), any())
+                mlsConversationRepository.establishMLSGroup(any(), any(), any(), any())
             }.returns(Either.Right(additionResult))
         }