Skip to content

Commit

Permalink
cover rotateKeysAndMigrationConversation with test in MLSConversation…
Browse files Browse the repository at this point in the history 
…Repository
  • Loading branch information
@mchenani
mchenani committed Oct 25, 2023
1 parent ebacf96 commit ac1543f
Show file tree
Hide file tree
Showing 3 changed files with 267 additions and 5 deletions.
16 changes: 12 additions & 4 deletions ...rc/commonMain/kotlin/com/wire/kalium/logic/data/conversation/MLSConversationRepository.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,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
Expand Down Expand Up @@ -529,19 +531,25 @@ internal class MLSConversationDataSource(
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)
keyPackageRepository.deleteKeyPackages(clientId, rotateBundle.keyPackageRefsToRemove).flatMapLeft {
return Either.Left(it)
}

kaliumLogger.w("upload new key packages including x509 certificate")
keyPackageRepository.uploadKeyPackages(clientId, rotateBundle.newKeyPackages)
keyPackageRepository.uploadKeyPackages(clientId, rotateBundle.newKeyPackages).flatMapLeft {
return Either.Left(it)
}

kaliumLogger.w("send migration commits after key rotations")
rotateBundle.commits.forEach {
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,
Expand Down
175 changes: 174 additions & 1 deletion ...ommonTest/kotlin/com/wire/kalium/logic/data/conversation/MLSConversationRepositoryTest.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,16 @@
package com.wire.kalium.logic.data.conversation

import com.wire.kalium.cryptography.CommitBundle
import com.wire.kalium.cryptography.E2EIClient
import com.wire.kalium.cryptography.E2EIConversationState
import com.wire.kalium.cryptography.GroupInfoBundle
import com.wire.kalium.cryptography.GroupInfoEncryptionType
import com.wire.kalium.cryptography.MLSClient
import com.wire.kalium.cryptography.RatchetTreeType
import com.wire.kalium.cryptography.RotateBundle
import com.wire.kalium.logic.CoreFailure
import com.wire.kalium.logic.data.client.MLSClientProvider
import com.wire.kalium.logic.data.conversation.MLSConversationRepositoryTest.Arrangement.Companion.TEST_FAILURE
import com.wire.kalium.logic.data.event.Event
import com.wire.kalium.logic.data.id.GroupID
import com.wire.kalium.logic.data.id.QualifiedClientID
Expand All @@ -35,6 +38,7 @@ import com.wire.kalium.logic.data.mlspublickeys.KeyType
import com.wire.kalium.logic.data.mlspublickeys.MLSPublicKey
import com.wire.kalium.logic.data.mlspublickeys.MLSPublicKeysRepository
import com.wire.kalium.logic.di.MapperProvider
import com.wire.kalium.logic.framework.TestClient
import com.wire.kalium.logic.framework.TestConversation
import com.wire.kalium.logic.framework.TestUser
import com.wire.kalium.logic.functional.Either
Expand Down Expand Up @@ -1100,6 +1104,149 @@ class MLSConversationRepositoryTest {
.wasNotInvoked()
}

@Test
fun givenSuccessResponse_whenRotatingKeysAndMigratingConversation_thenReturnsSuccess() = runTest {
val (arrangement, mlsConversationRepository) = Arrangement()
.withGetMLSClientSuccessful()
.withRotateAllSuccessful()
.withSendCommitBundleSuccessful()
.withUploadKeyPackagesReturning(Either.Right(Unit))
.withDeleteKeyPackagesReturning(Either.Right(Unit))
.arrange()

assertEquals(
Either.Right(Unit),
mlsConversationRepository.rotateKeysAndMigrateConversations(TestClient.CLIENT_ID, arrangement.e2eiClient, "")
)

verify(arrangement.mlsClient)
.suspendFunction(arrangement.mlsClient::e2eiRotateAll)
.with(any(), any(), any())
.wasInvoked(once)

verify(arrangement.keyPackageRepository)
.suspendFunction(arrangement.keyPackageRepository::deleteKeyPackages)
.with(any(), any())
.wasInvoked(once)

verify(arrangement.keyPackageRepository)
.suspendFunction(arrangement.keyPackageRepository::uploadKeyPackages)
.with(any(), any())
.wasInvoked(once)

verify(arrangement.mlsMessageApi)
.suspendFunction(arrangement.mlsMessageApi::sendCommitBundle)
.with(anyInstanceOf(MLSMessageApi.CommitBundle::class))
.wasInvoked(once)
}

@Test
fun givenDropKeypackagesFailed_whenRotatingKeysAndMigratingConversation_thenReturnsFailure() = runTest {
val (arrangement, mlsConversationRepository) = Arrangement()
.withGetMLSClientSuccessful()
.withRotateAllSuccessful()
.withUploadKeyPackagesReturning(Either.Right(Unit))
.withDeleteKeyPackagesReturning(TEST_FAILURE)
.withSendCommitBundleSuccessful()
.arrange()

assertEquals(
TEST_FAILURE,
mlsConversationRepository.rotateKeysAndMigrateConversations(TestClient.CLIENT_ID, arrangement.e2eiClient, "")
)

verify(arrangement.mlsClient)
.suspendFunction(arrangement.mlsClient::e2eiRotateAll)
.with(any(), any(), any())
.wasInvoked(once)

verify(arrangement.keyPackageRepository)
.suspendFunction(arrangement.keyPackageRepository::deleteKeyPackages)
.with(any(), any())
.wasInvoked(once)

verify(arrangement.keyPackageRepository)
.suspendFunction(arrangement.keyPackageRepository::uploadKeyPackages)
.with(any(), any())
.wasNotInvoked()

verify(arrangement.mlsMessageApi)
.suspendFunction(arrangement.mlsMessageApi::sendCommitBundle)
.with(anyInstanceOf(MLSMessageApi.CommitBundle::class))
.wasNotInvoked()
}

@Test
fun givenUploadKeypackagesFailed_whenRotatingKeysAndMigratingConversation_thenReturnsFailure() = runTest {
val (arrangement, mlsConversationRepository) = Arrangement()
.withGetMLSClientSuccessful()
.withRotateAllSuccessful()
.withUploadKeyPackagesReturning(TEST_FAILURE)
.withDeleteKeyPackagesReturning(Either.Right(Unit))
.withSendCommitBundleSuccessful()
.arrange()

assertEquals(
TEST_FAILURE,
mlsConversationRepository.rotateKeysAndMigrateConversations(TestClient.CLIENT_ID, arrangement.e2eiClient, "")
)

verify(arrangement.mlsClient)
.suspendFunction(arrangement.mlsClient::e2eiRotateAll)
.with(any(), any(), any())
.wasInvoked(once)

verify(arrangement.keyPackageRepository)
.suspendFunction(arrangement.keyPackageRepository::deleteKeyPackages)
.with(any(), any())
.wasInvoked(once)

verify(arrangement.keyPackageRepository)
.suspendFunction(arrangement.keyPackageRepository::uploadKeyPackages)
.with(any(), any())
.wasInvoked(once)

verify(arrangement.mlsMessageApi)
.suspendFunction(arrangement.mlsMessageApi::sendCommitBundle)
.with(anyInstanceOf(MLSMessageApi.CommitBundle::class))
.wasNotInvoked()
}

@Test
fun givenSendingCommitBundlesFails_whenRotatingKeysAndMigratingConversation_thenReturnsFailure() = runTest {
val (arrangement, mlsConversationRepository) = Arrangement()
.withGetMLSClientSuccessful()
.withRotateAllSuccessful()
.withUploadKeyPackagesReturning(Either.Right(Unit))
.withDeleteKeyPackagesReturning(Either.Right(Unit))
.withSendCommitBundleFailing(Arrangement.MLS_CLIENT_MISMATCH_ERROR, times = 1)
.arrange()


val result = mlsConversationRepository.rotateKeysAndMigrateConversations(TestClient.CLIENT_ID, arrangement.e2eiClient, "")
result.shouldFail()

verify(arrangement.mlsClient)
.suspendFunction(arrangement.mlsClient::e2eiRotateAll)
.with(any(), any(), any())
.wasInvoked(once)

verify(arrangement.keyPackageRepository)
.suspendFunction(arrangement.keyPackageRepository::deleteKeyPackages)
.with(any(), any())
.wasInvoked(once)

verify(arrangement.keyPackageRepository)
.suspendFunction(arrangement.keyPackageRepository::uploadKeyPackages)
.with(any(), any())
.wasInvoked(once)

verify(arrangement.mlsMessageApi)
.suspendFunction(arrangement.mlsMessageApi::sendCommitBundle)
.with(anyInstanceOf(MLSMessageApi.CommitBundle::class))
.wasInvoked(once)
}

private class Arrangement {

@Mock
Expand All @@ -1126,6 +1273,9 @@ class MLSConversationRepositoryTest {
@Mock
val mlsClient = mock(classOf<MLSClient>())

@Mock
val e2eiClient = mock(classOf<E2EIClient>())

@Mock
val syncManager = mock(SyncManager::class)

Expand Down Expand Up @@ -1172,6 +1322,20 @@ class MLSConversationRepositoryTest {
.then { Either.Right(keyPackages) }
}

fun withUploadKeyPackagesReturning(result: Either<CoreFailure, Unit>) = apply {
given(keyPackageRepository)
.suspendFunction(keyPackageRepository::uploadKeyPackages)
.whenInvokedWith(anything(), anything())
.thenReturn(result)
}

fun withDeleteKeyPackagesReturning(result: Either<CoreFailure, Unit>) = apply {
given(keyPackageRepository)
.suspendFunction(keyPackageRepository::deleteKeyPackages)
.whenInvokedWith(anything(), anything())
.thenReturn(result)
}

fun withGetPublicKeysSuccessful() = apply {
given(mlsPublicKeysRepository)
.suspendFunction(mlsPublicKeysRepository::getKeys)
Expand All @@ -1193,6 +1357,13 @@ class MLSConversationRepositoryTest {
.then { Either.Left(failure) }
}

fun withRotateAllSuccessful() = apply {
given(mlsClient)
.suspendFunction(mlsClient::e2eiRotateAll)
.whenInvokedWith(anything(), anything(), anything())
.thenReturn(ROTATE_BUNDLE)
}

fun withAddMLSMemberSuccessful() = apply {
given(mlsClient)
.suspendFunction(mlsClient::addMember)
Expand Down Expand Up @@ -1330,7 +1501,8 @@ class MLSConversationRepositoryTest {
proposalTimersFlow
)

internal companion object {
companion object {
val TEST_FAILURE = Either.Left(CoreFailure.Unknown(Throwable("an error")))
const val EPOCH = 5UL
const val RAW_GROUP_ID = "groupId"
val TIME = DateTimeUtil.currentIsoDateTimeString()
Expand Down Expand Up @@ -1359,6 +1531,7 @@ class MLSConversationRepositoryTest {
PUBLIC_GROUP_STATE
)
val COMMIT_BUNDLE = CommitBundle(COMMIT, WELCOME, PUBLIC_GROUP_STATE_BUNDLE)
val ROTATE_BUNDLE = RotateBundle(mapOf(RAW_GROUP_ID to COMMIT_BUNDLE), emptyList(), emptyList())
val DECRYPTED_MESSAGE_BUNDLE = com.wire.kalium.cryptography.DecryptedMessageBundle(
message = null,
commitDelay = null,
Expand Down
Loading

0 comments on commit ac1543f

Please sign in to comment.