feat: discover legal hold when receiving message [WPB-5998]

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

@@ -1250,8 +1250,8 @@ class UserSessionScope internal constructor(
         get() = NewMessageEventHandlerImpl(
             proteusUnpacker,
             mlsUnpacker,
-            conversationRepository,
             applicationMessageHandler,
+            legalHoldHandler,
             { conversationId, messageId ->
                 messages.ephemeralMessageDeletionHandler.startSelfDeletion(conversationId, messageId)
             },
@@ -1389,18 +1389,20 @@ class UserSessionScope internal constructor(
 
     private val legalHoldSystemMessagesHandler = LegalHoldSystemMessagesHandlerImpl(
         selfUserId = userId,
-        membersHavingLegalHoldClient = membersHavingLegalHoldClient,
         persistMessage = persistMessage,
         conversationRepository = conversationRepository,
-        messageRepository = messageRepository,
+        messageRepository = messageRepository
     )
 
     private val legalHoldHandler = LegalHoldHandlerImpl(
         selfUserId = userId,
         persistOtherUserClients = persistOtherUserClients,
         fetchSelfClientsFromRemote = fetchSelfClientsFromRemote,
         observeLegalHoldStateForUser = observeLegalHoldStateForUser,
+        membersHavingLegalHoldClient = membersHavingLegalHoldClient,
         userConfigRepository = userConfigRepository,
+        conversationRepository = conversationRepository,
+        observeSyncState = observeSyncState,
         legalHoldSystemMessagesHandler = legalHoldSystemMessagesHandler,
     )
 

logic/src/commonMain/kotlin/com/wire/kalium/logic/sync/receiver/conversation/message/NewMessageEventHandler.kt

@@ -23,7 +23,6 @@ import com.wire.kalium.logger.KaliumLogger
 import com.wire.kalium.logic.ProteusFailure
 import com.wire.kalium.logic.data.conversation.ClientId
 import com.wire.kalium.logic.data.conversation.Conversation
-import com.wire.kalium.logic.data.conversation.ConversationRepository
 import com.wire.kalium.logic.data.event.Event
 import com.wire.kalium.logic.data.event.EventLoggingStatus
 import com.wire.kalium.logic.data.event.logEventProcessing
@@ -34,6 +33,7 @@ import com.wire.kalium.logic.feature.message.StaleEpochVerifier
 import com.wire.kalium.logic.functional.onFailure
 import com.wire.kalium.logic.functional.onSuccess
 import com.wire.kalium.logic.kaliumLogger
+import com.wire.kalium.logic.sync.receiver.handler.legalhold.LegalHoldHandler
 import com.wire.kalium.util.serialization.toJsonElement
 import kotlinx.datetime.toInstant
 
@@ -46,8 +46,8 @@ internal interface NewMessageEventHandler {
 internal class NewMessageEventHandlerImpl(
     private val proteusMessageUnpacker: ProteusMessageUnpacker,
     private val mlsMessageUnpacker: MLSMessageUnpacker,
-    private val conversationRepository: ConversationRepository,
     private val applicationMessageHandler: ApplicationMessageHandler,
+    private val legalHoldHandler: LegalHoldHandler,
     private val enqueueSelfDeletion: (conversationId: ConversationId, messageId: String) -> Unit,
     private val selfUserId: UserId,
     private val staleEpochVerifier: StaleEpochVerifier
@@ -86,13 +86,10 @@ internal class NewMessageEventHandlerImpl(
                 )
             }.onSuccess {
                 if (it is MessageUnpackResult.ApplicationMessage) {
-                    handleSuccessfulResult(it)
                     if (it.content.legalHoldStatus != Conversation.LegalHoldStatus.UNKNOWN) {
-                        conversationRepository.updateLegalHoldStatus(
-                            conversationId = it.conversationId,
-                            legalHoldStatus = it.content.legalHoldStatus
-                        )
+                        legalHoldHandler.handleNewMessage(it, event.live)
                     }
+                    handleSuccessfulResult(it)
                     onMessageInserted(it)
                 }
                 kaliumLogger
@@ -143,13 +140,10 @@ internal class NewMessageEventHandlerImpl(
             }.onSuccess {
                 it.forEach {
                     if (it is MessageUnpackResult.ApplicationMessage) {
-                        handleSuccessfulResult(it)
                         if (it.content.legalHoldStatus != Conversation.LegalHoldStatus.UNKNOWN) {
-                            conversationRepository.updateLegalHoldStatus(
-                                conversationId = it.conversationId,
-                                legalHoldStatus = it.content.legalHoldStatus
-                            )
+                            legalHoldHandler.handleNewMessage(it, event.live)
                         }
+                        handleSuccessfulResult(it)
                         onMessageInserted(it)
                     }
                 }

logic/src/commonMain/kotlin/com/wire/kalium/logic/sync/receiver/handler/legalhold/LegalHoldHandler.kt

@@ -19,19 +19,39 @@ package com.wire.kalium.logic.sync.receiver.handler.legalhold
 
 import com.wire.kalium.logic.CoreFailure
 import com.wire.kalium.logic.configuration.UserConfigRepository
+import com.wire.kalium.logic.data.conversation.Conversation
+import com.wire.kalium.logic.data.conversation.ConversationRepository
 import com.wire.kalium.logic.data.event.Event
+import com.wire.kalium.logic.data.id.ConversationId
+import com.wire.kalium.logic.data.sync.SyncState
 import com.wire.kalium.logic.data.user.UserId
 import com.wire.kalium.logic.feature.client.FetchSelfClientsFromRemoteUseCase
 import com.wire.kalium.logic.feature.client.PersistOtherUserClientsUseCase
 import com.wire.kalium.logic.feature.legalhold.LegalHoldState
+import com.wire.kalium.logic.feature.legalhold.MembersHavingLegalHoldClientUseCase
 import com.wire.kalium.logic.feature.legalhold.ObserveLegalHoldStateForUserUseCase
 import com.wire.kalium.logic.functional.Either
+import com.wire.kalium.logic.functional.flatMap
+import com.wire.kalium.logic.functional.foldToEitherWhileRight
+import com.wire.kalium.logic.functional.getOrElse
+import com.wire.kalium.logic.functional.map
 import com.wire.kalium.logic.kaliumLogger
+import com.wire.kalium.logic.sync.ObserveSyncStateUseCase
+import com.wire.kalium.logic.sync.receiver.conversation.message.MessageUnpackResult
+import com.wire.kalium.logic.util.TriggerBuffer
+import com.wire.kalium.util.DateTimeUtil
+import com.wire.kalium.util.KaliumDispatcher
+import com.wire.kalium.util.KaliumDispatcherImpl
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.firstOrNull
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.launch
 
 internal interface LegalHoldHandler {
     suspend fun handleEnable(legalHoldEnabled: Event.User.LegalHoldEnabled): Either<CoreFailure, Unit>
     suspend fun handleDisable(legalHoldDisabled: Event.User.LegalHoldDisabled): Either<CoreFailure, Unit>
+    suspend fun handleNewMessage(message: MessageUnpackResult.ApplicationMessage, live: Boolean): Either<CoreFailure, Unit>
 }
 
 @Suppress("LongParameterList")
@@ -40,9 +60,24 @@ internal class LegalHoldHandlerImpl internal constructor(
     private val persistOtherUserClients: PersistOtherUserClientsUseCase,
     private val fetchSelfClientsFromRemote: FetchSelfClientsFromRemoteUseCase,
     private val observeLegalHoldStateForUser: ObserveLegalHoldStateForUserUseCase,
+    private val membersHavingLegalHoldClient: MembersHavingLegalHoldClientUseCase,
     private val userConfigRepository: UserConfigRepository,
+    private val conversationRepository: ConversationRepository,
     private val legalHoldSystemMessagesHandler: LegalHoldSystemMessagesHandler,
+    observeSyncState: ObserveSyncStateUseCase,
+    kaliumDispatcher: KaliumDispatcher = KaliumDispatcherImpl,
 ) : LegalHoldHandler {
+    private val scope = CoroutineScope(kaliumDispatcher.default)
+    private val bufferedUpdatedConversationIds =
+        TriggerBuffer<ConversationId>(observeSyncState().distinctUntilChanged().map { it == SyncState.Live }, scope)
+
+    init {
+        scope.launch {
+            bufferedUpdatedConversationIds.observe()
+                .collect { handleUpdatedConversations(it) }
+        }
+    }
+
     override suspend fun handleEnable(legalHoldEnabled: Event.User.LegalHoldEnabled): Either<CoreFailure, Unit> {
         kaliumLogger.i("legal hold enabled for user ${legalHoldEnabled.userId.toLogString()}")
         // check if the user has already been under legal hold prior to this event
@@ -54,7 +89,8 @@ internal class LegalHoldHandlerImpl internal constructor(
             if (selfUserId == legalHoldEnabled.userId) { // notify only for self user
                 userConfigRepository.setLegalHoldChangeNotified(false)
             }
-            legalHoldSystemMessagesHandler.handleEnable(legalHoldEnabled.userId)
+            handleConversationsForUser(legalHoldEnabled.userId)
+            legalHoldSystemMessagesHandler.handleEnabledForUser(legalHoldEnabled.userId, DateTimeUtil.currentIsoDateTimeString())
         }
 
         return Either.Right(Unit)
@@ -71,12 +107,26 @@ internal class LegalHoldHandlerImpl internal constructor(
             if (selfUserId == legalHoldDisabled.userId) { // notify only for self user
                 userConfigRepository.setLegalHoldChangeNotified(false)
             }
-            legalHoldSystemMessagesHandler.handleDisable(legalHoldDisabled.userId)
+            handleConversationsForUser(legalHoldDisabled.userId)
+            legalHoldSystemMessagesHandler.handleDisabledForUser(legalHoldDisabled.userId, DateTimeUtil.currentIsoDateTimeString())
         }
 
         return Either.Right(Unit)
     }
 
+    override suspend fun handleNewMessage(message: MessageUnpackResult.ApplicationMessage, live: Boolean): Either<CoreFailure, Unit> {
+        val isStatusChangedForConversation = when (val legalHoldStatus = message.content.legalHoldStatus) {
+            Conversation.LegalHoldStatus.ENABLED, Conversation.LegalHoldStatus.DISABLED ->
+                handleForConversation(message.conversationId, legalHoldStatus, message.timestampIso)
+            else -> false
+        }
+        if (isStatusChangedForConversation) {
+            if (live) handleUpdatedConversations(listOf(message.conversationId), message.timestampIso) // handle it right away
+            else bufferedUpdatedConversationIds.add(message.conversationId) // buffer and handle after sync
+        }
+        return Either.Right(Unit)
+    }
+
     private suspend fun processEvent(selfUserId: UserId, userId: UserId) {
         if (selfUserId == userId) {
             userConfigRepository.deleteLegalHoldRequest()
@@ -88,4 +138,66 @@ internal class LegalHoldHandlerImpl internal constructor(
 
     private suspend fun isUserUnderLegalHold(userId: UserId): Boolean =
         observeLegalHoldStateForUser(userId).firstOrNull() == LegalHoldState.Enabled
+
+    private suspend fun handleForConversation(
+        conversationId: ConversationId,
+        newStatus: Conversation.LegalHoldStatus,
+        systemMessageTimestampIso: String = DateTimeUtil.currentIsoDateTimeString(),
+    ): Boolean =
+        if (newStatus != Conversation.LegalHoldStatus.UNKNOWN) {
+            conversationRepository.updateLegalHoldStatus(conversationId, newStatus)
+                .getOrElse(false)
+                .also { isChanged -> // if conversation legal hold status has changed, create system message for it
+                    if (isChanged) {
+                        when (newStatus) {
+                            Conversation.LegalHoldStatus.DISABLED ->
+                                legalHoldSystemMessagesHandler.handleDisabledForConversation(conversationId, systemMessageTimestampIso)
+                            Conversation.LegalHoldStatus.ENABLED ->
+                                legalHoldSystemMessagesHandler.handleEnabledForConversation(conversationId, systemMessageTimestampIso)
+                            else -> { /* do nothing */ }
+                        }
+                    }
+                }
+        } else false
+
+    private suspend fun handleConversationsForUser(userId: UserId) {
+        conversationRepository.getConversationsByUserId(userId).map { conversations ->
+            conversations.forEach { conversation ->
+                // create system message for conversation if needed
+                membersHavingLegalHoldClient(conversation.id)
+                    .map { if (it.isEmpty()) Conversation.LegalHoldStatus.DISABLED else Conversation.LegalHoldStatus.ENABLED }
+                    .map { newLegalHoldStatus -> handleForConversation(conversation.id, newLegalHoldStatus) }
+            }
+        }
+    }
+
+    private suspend fun handleUpdatedConversations(
+        conversationIds: List<ConversationId>,
+        systemMessageTimestampIso: String = DateTimeUtil.currentIsoDateTimeString(),
+    ) {
+        conversationIds
+            .foldToEitherWhileRight(mapOf<UserId, Boolean>()) { conversationId, acc ->
+                conversationRepository.getConversationMembers(conversationId)
+                    .flatMap { members ->
+                        membersHavingLegalHoldClient(conversationId)
+                            .map { membersHavingLegalHoldClient ->
+                                members.associateWith { membersHavingLegalHoldClient.contains(it) } + acc
+                            }
+                    }
+            }
+            .map {
+                it.forEach { (userId, userHasBeenUnderLegalHold) ->
+                    // TODO: to be optimized - send empty message and handle legal hold discovery after sending a message
+                    processEvent(selfUserId, userId)
+                    val userIsNowUnderLegalHold = isUserUnderLegalHold(userId)
+                    if (userHasBeenUnderLegalHold != userIsNowUnderLegalHold) {
+                        if (selfUserId == userId) { // notify only for self user
+                            userConfigRepository.setLegalHoldChangeNotified(false)
+                        }
+                        if (userIsNowUnderLegalHold) legalHoldSystemMessagesHandler.handleEnabledForUser(userId, systemMessageTimestampIso)
+                        else legalHoldSystemMessagesHandler.handleDisabledForUser(userId, systemMessageTimestampIso)
+                    }
+                }
+            }
+    }
 }