fix: secure message draft saving [WPB-11325]

persistence/src/commonMain/db_user/com/wire/kalium/persistence/Messages.sq

@@ -488,6 +488,9 @@ UPDATE Message
 SET id = :newId
 WHERE id = :oldId AND conversation_id = :conversationId;
 
+getMessage:
+SELECT * FROM Message WHERE id = ? AND conversation_id = ?;
+
 selectById:
 SELECT * FROM MessageDetailsView WHERE id = ? AND conversationId = ?;
 

persistence/src/commonMain/kotlin/com/wire/kalium/persistence/dao/message/draft/MessageDraftDAOImpl.kt

@@ -18,7 +18,9 @@
 package com.wire.kalium.persistence.dao.message.draft
 
 import app.cash.sqldelight.coroutines.asFlow
+import com.wire.kalium.persistence.ConversationsQueries
 import com.wire.kalium.persistence.MessageDraftsQueries
+import com.wire.kalium.persistence.MessagesQueries
 import com.wire.kalium.persistence.dao.ConversationIDEntity
 import com.wire.kalium.persistence.dao.message.draft.MessageDraftMapper.toDao
 import com.wire.kalium.persistence.util.mapToList
@@ -29,11 +31,36 @@ import kotlin.coroutines.CoroutineContext
 
 class MessageDraftDAOImpl internal constructor(
     private val queries: MessageDraftsQueries,
+    private val messagesQueries: MessagesQueries,
+    private val conversationsQueries: ConversationsQueries,
     private val coroutineContext: CoroutineContext,
 ) : MessageDraftDAO {
 
     override suspend fun upsertMessageDraft(messageDraft: MessageDraftEntity) =
         withContext(coroutineContext) {
+            val conversationExists = conversationsQueries.selectConversationByQualifiedId(messageDraft.conversationId)
+                .executeAsOneOrNull() != null
+
+            if (!conversationExists) {
+                return@withContext
+            }
+
+            if (messageDraft.editMessageId != null) {
+                val messageExists = messagesQueries.getMessage(messageDraft.editMessageId, messageDraft.conversationId)
+                    .executeAsOneOrNull() != null
+                if (!messageExists) {
+                    return@withContext
+                }
+            }
+
+            if (messageDraft.quotedMessageId != null) {
+                val quotedMessageExists = messagesQueries.getMessage(messageDraft.quotedMessageId, messageDraft.conversationId)
+                    .executeAsOneOrNull() != null
+                if (!quotedMessageExists) {
+                    return@withContext
+                }
+            }
+
             queries.upsertDraft(
                 conversation_id = messageDraft.conversationId,
                 text = messageDraft.text,

persistence/src/commonMain/kotlin/com/wire/kalium/persistence/db/UserDatabaseBuilder.kt

@@ -256,6 +256,8 @@ class UserDatabaseBuilder internal constructor(
 
     val messageDraftDAO = MessageDraftDAOImpl(
         database.messageDraftsQueries,
+        database.messagesQueries,
+        database.conversationsQueries,
         queriesContext
     )
 

persistence/src/commonTest/kotlin/com/wire/kalium/persistence/dao/message/draft/MessageDraftDAOTest.kt

@@ -113,6 +113,48 @@ class MessageDraftDAOTest : BaseDatabaseTest() {
         assertEquals(null, removedResult)
     }
 
+    @Test
+    fun givenMessageIsRemoved_whenUpsertingDraft_thenItShouldIgnore() = runTest {
+        // Given
+        insertInitialData()
+        messageDAO.deleteMessage("editMessageId", conversationEntity1.id)
+
+        // When
+        messageDraftDAO.upsertMessageDraft(MESSAGE_DRAFT)
+
+        // Then
+        val result = messageDraftDAO.getMessageDraft(MESSAGE_DRAFT.conversationId)
+        assertEquals(null, result)
+    }
+
+    @Test
+    fun givenConversationIsRemoved_whenUpsertingDraft_thenItShouldIgnore() = runTest {
+        // Given
+        insertInitialData()
+        conversationDAO.deleteConversationByQualifiedID(conversationEntity1.id)
+
+        // When
+        messageDraftDAO.upsertMessageDraft(MESSAGE_DRAFT)
+
+        // Then
+        val result = messageDraftDAO.getMessageDraft(MESSAGE_DRAFT.conversationId)
+        assertEquals(null, result)
+    }
+
+    @Test
+    fun givenQuotedMessageIsRemoved_whenUpsertingDraft_thenItShouldIgnore() = runTest {
+        // Given
+        insertInitialData()
+        messageDAO.deleteMessage("quotedMessageId", conversationEntity1.id)
+
+        // When
+        messageDraftDAO.upsertMessageDraft(MESSAGE_DRAFT)
+
+        // Then
+        val result = messageDraftDAO.getMessageDraft(MESSAGE_DRAFT.conversationId)
+        assertEquals(null, result)
+    }
+
     private suspend fun insertInitialData() {
         userDAO.upsertUsers(listOf(userEntity1))
         conversationDAO.insertConversation(conversationEntity1)