fix(proteus): prevent missing messages by using transactions [WPB-10873]

cryptography/src/androidMain/kotlin/com/wire/kalium/cryptography/ProteusClientCryptoBoxImpl.kt

@@ -110,18 +110,24 @@ class ProteusClientCryptoBoxImpl constructor(
         }
     }
 
-    override suspend fun decrypt(message: ByteArray, sessionId: CryptoSessionId): ByteArray = lock.withLock {
+    override suspend fun <T : Any> decrypt(
+        message: ByteArray,
+        sessionId: CryptoSessionId,
+        handleDecryptedMessage: suspend (decryptedMessage: ByteArray) -> T
+    ): T = lock.withLock {
         withContext(defaultContext) {
             val session = box.tryGetSession(sessionId.value)
             wrapException {
                 if (session != null) {
                     val decryptedMessage = session.decrypt(message)
-                    session.save()
-                    decryptedMessage
+                    handleDecryptedMessage(decryptedMessage).also {
+                        session.save()
+                    }
                 } else {
                     val result = box.initSessionFromMessage(sessionId.value, message)
-                    result.session.save()
-                    result.message
+                    handleDecryptedMessage(result.message).also {
+                        result.session.save()
+                    }
                 }
             }
         }

cryptography/src/appleMain/kotlin/com/wire/kalium/cryptography/ProteusClientCoreCryptoImpl.kt

@@ -31,7 +31,8 @@ import platform.Foundation.URLByAppendingPathComponent
 @Suppress("TooManyFunctions")
 class ProteusClientCoreCryptoImpl private constructor(private val coreCrypto: CoreCrypto) : ProteusClient {
     @Suppress("EmptyFunctionBlock")
-    override suspend fun close() {}
+    override suspend fun close() {
+    }
 
     override fun getIdentity(): ByteArray {
         return ByteArray(0)
@@ -72,18 +73,21 @@ class ProteusClientCoreCryptoImpl private constructor(private val coreCrypto: Co
         wrapException { coreCrypto.proteusSessionFromPrekey(sessionId.value, toUByteList(preKeyCrypto.encodedData.decodeBase64Bytes())) }
     }
 
-    override suspend fun decrypt(message: ByteArray, sessionId: CryptoSessionId): ByteArray {
+    override suspend fun <T : Any> decrypt(
+        message: ByteArray,
+        sessionId: CryptoSessionId,
+        handleDecryptedMessage: suspend (decryptedMessage: ByteArray) -> T
+    ): T {
         val sessionExists = doesSessionExist(sessionId)
 
         return wrapException {
-            if (sessionExists) {
-                val decryptedMessage = toByteArray(coreCrypto.proteusDecrypt(sessionId.value, toUByteList(message)))
-                coreCrypto.proteusSessionSave(sessionId.value)
-                decryptedMessage
+            val decryptedMessage = if (sessionExists) {
+                toByteArray(coreCrypto.proteusDecrypt(sessionId.value, toUByteList(message)))
             } else {
-                val decryptedMessage = toByteArray(coreCrypto.proteusSessionFromMessage(sessionId.value, toUByteList(message)))
+                toByteArray(coreCrypto.proteusSessionFromMessage(sessionId.value, toUByteList(message)))
+            }
+            handleDecryptedMessage(decryptedMessage).also {
                 coreCrypto.proteusSessionSave(sessionId.value)
-                decryptedMessage
             }
         }
     }
@@ -129,7 +133,7 @@ class ProteusClientCoreCryptoImpl private constructor(private val coreCrypto: Co
     }
 
     @Suppress("TooGenericExceptionCaught")
-    private fun <T> wrapException(b: () -> T): T {
+    private inline fun <T> wrapException(b: () -> T): T {
         try {
             return b()
         } catch (e: CryptoException) {

cryptography/src/commonJvmAndroid/kotlin/com.wire.kalium.cryptography/ProteusClientCoreCryptoImpl.kt

@@ -85,15 +85,20 @@ class ProteusClientCoreCryptoImpl private constructor(
         wrapException { coreCrypto.proteusSessionFromPrekey(sessionId.value, preKeyCrypto.encodedData.decodeBase64Bytes()) }
     }
 
-    override suspend fun decrypt(message: ByteArray, sessionId: CryptoSessionId): ByteArray {
+    override suspend fun <T : Any> decrypt(
+        message: ByteArray,
+        sessionId: CryptoSessionId,
+        handleDecryptedMessage: suspend (decryptedMessage: ByteArray) -> T
+    ): T {
         val sessionExists = doesSessionExist(sessionId)
 
         return wrapException {
-            if (sessionExists) {
+            val decryptedMessage = if (sessionExists) {
                 coreCrypto.proteusDecrypt(sessionId.value, message)
             } else {
                 coreCrypto.proteusSessionFromMessage(sessionId.value, message)
             }
+            handleDecryptedMessage(decryptedMessage)
         }
     }
 

cryptography/src/commonMain/kotlin/com/wire/kalium/cryptography/ProteusClient.kt

@@ -81,8 +81,18 @@ interface ProteusClient {
     @Throws(ProteusException::class, CancellationException::class)
     suspend fun createSession(preKeyCrypto: PreKeyCrypto, sessionId: CryptoSessionId)
 
+    /**
+     * Decrypts a message.
+     * In case of success, calls [handleDecryptedMessage] with the decrypted bytes.
+     * @throws ProteusException in case of failure
+     * @throws CancellationException
+     */
     @Throws(ProteusException::class, CancellationException::class)
-    suspend fun decrypt(message: ByteArray, sessionId: CryptoSessionId): ByteArray
+    suspend fun <T : Any> decrypt(
+        message: ByteArray,
+        sessionId: CryptoSessionId,
+        handleDecryptedMessage: suspend (decryptedMessage: ByteArray) -> T
+    ): T
 
     @Throws(ProteusException::class, CancellationException::class)
     suspend fun encrypt(message: ByteArray, sessionId: CryptoSessionId): ByteArray

cryptography/src/commonTest/kotlin/com/wire/kalium/cryptography/ProteusClientTest.kt

@@ -1,20 +1,20 @@
- /*
- * Wire
- * Copyright (C) 2024 Wire Swiss GmbH
- *
- * This program is free software: you can redistribute it and/or modify
- * it under the terms of the GNU General Public License as published by
- * the Free Software Foundation, either version 3 of the License, or
- * (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
- * GNU General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License
- * along with this program. If not, see http://www.gnu.org/licenses/.
- */
+/*
+* Wire
+* Copyright (C) 2024 Wire Swiss GmbH
+*
+* This program is free software: you can redistribute it and/or modify
+* it under the terms of the GNU General Public License as published by
+* the Free Software Foundation, either version 3 of the License, or
+* (at your option) any later version.
+*
+* This program is distributed in the hope that it will be useful,
+* but WITHOUT ANY WARRANTY; without even the implied warranty of
+* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+* GNU General Public License for more details.
+*
+* You should have received a copy of the GNU General Public License
+* along with this program. If not, see http://www.gnu.org/licenses/.
+*/
 
 package com.wire.kalium.cryptography
 
@@ -93,7 +93,7 @@ class ProteusClientTest : BaseProteusClientTest() {
         val message = "Hi Alice!"
         val aliceKey = aliceClient.newPreKeys(0, 10).first()
         val encryptedMessage = bobClient.encryptWithPreKey(message.encodeToByteArray(), aliceKey, aliceSessionId)
-        val decryptedMessage = aliceClient.decrypt(encryptedMessage, bobSessionId)
+        val decryptedMessage = aliceClient.decrypt(encryptedMessage, bobSessionId) { it }
         assertEquals(message, decryptedMessage.decodeToString())
     }
 
@@ -105,11 +105,11 @@ class ProteusClientTest : BaseProteusClientTest() {
         val aliceKey = aliceClient.newPreKeys(0, 10).first()
         val message1 = "Hi Alice!"
         val encryptedMessage1 = bobClient.encryptWithPreKey(message1.encodeToByteArray(), aliceKey, aliceSessionId)
-        aliceClient.decrypt(encryptedMessage1, bobSessionId)
+        aliceClient.decrypt(encryptedMessage1, bobSessionId) {}
 
         val message2 = "Hi again Alice!"
         val encryptedMessage2 = bobClient.encrypt(message2.encodeToByteArray(), aliceSessionId)
-        val decryptedMessage2 = aliceClient.decrypt(encryptedMessage2, bobSessionId)
+        val decryptedMessage2 = aliceClient.decrypt(encryptedMessage2, bobSessionId) { it }
 
         assertEquals(message2, decryptedMessage2.decodeToString())
     }
@@ -124,10 +124,10 @@ class ProteusClientTest : BaseProteusClientTest() {
         val aliceKey = aliceClient.newPreKeys(0, 10).first()
         val message1 = "Hi Alice!"
         val encryptedMessage1 = bobClient.encryptWithPreKey(message1.encodeToByteArray(), aliceKey, aliceSessionId)
-        aliceClient.decrypt(encryptedMessage1, bobSessionId)
+        aliceClient.decrypt(encryptedMessage1, bobSessionId) {}
 
         val exception: ProteusException = assertFailsWith {
-            aliceClient.decrypt(encryptedMessage1, bobSessionId)
+            aliceClient.decrypt(encryptedMessage1, bobSessionId) {}
         }
         assertEquals(ProteusException.Code.DUPLICATE_MESSAGE, exception.code)
     }
@@ -188,8 +188,44 @@ class ProteusClientTest : BaseProteusClientTest() {
         }
     }
 
+    // TODO: Implement on CoreCrypto as well once it supports transactions
+    @IgnoreJS
+    @IgnoreJvm
+    @IgnoreIOS
+    @Test
+    fun givenNonEncryptedClient_whenThrowingDuringTransaction_thenShouldNotSaveSessionAndBeAbleToDecryptAgain() = runTest {
+        val aliceRef = createProteusStoreRef(alice.id)
+        val failedAliceClient = createProteusClient(aliceRef)
+        val bobClient = createProteusClient(createProteusStoreRef(bob.id))
+
+        val aliceKey = failedAliceClient.newPreKeys(0, 10).first()
+        val message1 = "Hi Alice!"
+
+        var decryptedCount = 0
+
+        val encryptedMessage1 = bobClient.encryptWithPreKey(message1.encodeToByteArray(), aliceKey, aliceSessionId)
+        try {
+            failedAliceClient.decrypt(encryptedMessage1, bobSessionId) {
+                decryptedCount++
+                throw NullPointerException("")
+            }
+        } catch (ignore: Throwable) {
+            /** No-op **/
+        }
+        // Assume that the app crashed after decrypting but before saving session.
+        // Trying to decrypt again should succeed.
+
+        val secondAliceClient = createProteusClient(aliceRef)
+
+        val result = secondAliceClient.decrypt(encryptedMessage1, bobSessionId) { result ->
+            decryptedCount++
+            result
+        }
+        assertEquals(message1, result.decodeToString())
+        assertEquals(2, decryptedCount)
+    }
+
     companion object {
         val PROTEUS_DB_SECRET = ProteusDBSecret("secret")
     }
-
 }

cryptography/src/jsMain/kotlin/com/wire/kalium/cryptography/ProteusClientCryptoBoxImpl.kt

@@ -102,12 +102,13 @@ class ProteusClientCryptoBoxImpl : ProteusClient {
         box.session_from_prekey(sessionId.value, preKeyBundle.toArrayBuffer()).await()
     }
 
-    override suspend fun decrypt(
+    override suspend fun <T : Any> decrypt(
         message: ByteArray,
-        sessionId: CryptoSessionId
-    ): ByteArray {
+        sessionId: CryptoSessionId,
+        handleDecryptedMessage: suspend (decryptedMessage: ByteArray) -> T
+    ): T {
         val decryptedMessage = box.decrypt(sessionId.value, message.toArrayBuffer()).await()
-        return Int8Array(decryptedMessage.buffer).unsafeCast<ByteArray>()
+        return handleDecryptedMessage(Int8Array(decryptedMessage.buffer).unsafeCast<ByteArray>())
     }
 
     override suspend fun encrypt(

cryptography/src/jvmMain/kotlin/com/wire/kalium/cryptography/ProteusClientCryptoBoxImpl.kt

@@ -27,17 +27,13 @@ import java.util.Base64
 import kotlin.coroutines.CoroutineContext
 
 @Suppress("TooManyFunctions")
-class ProteusClientCryptoBoxImpl constructor(
+class ProteusClientCryptoBoxImpl(
     rootDir: String
 ) : ProteusClient {
 
-    private val path: String
+    private val path: String = rootDir
     private lateinit var box: CryptoBox
 
-    init {
-        path = rootDir
-    }
-
     fun openOrCreate() {
         val directory = File(path)
         box = wrapException {
@@ -84,14 +80,18 @@ class ProteusClientCryptoBoxImpl constructor(
         wrapException { box.encryptFromPreKeys(sessionId.value, toPreKey(preKeyCrypto), ByteArray(0)) }
     }
 
-    override suspend fun decrypt(message: ByteArray, sessionId: CryptoSessionId): ByteArray {
-        return wrapException { box.decrypt(sessionId.value, message) }
+    override suspend fun <T : Any> decrypt(
+        message: ByteArray,
+        sessionId: CryptoSessionId,
+        handleDecryptedMessage: suspend (decryptedMessage: ByteArray) -> T
+    ): T = wrapException {
+        handleDecryptedMessage(box.decrypt(sessionId.value, message))
     }
 
     override suspend fun encrypt(message: ByteArray, sessionId: CryptoSessionId): ByteArray {
         return wrapException {
             box.encryptFromSession(sessionId.value, message)
-        }?.let { it } ?: throw ProteusException(null, ProteusException.Code.SESSION_NOT_FOUND)
+        } ?: throw ProteusException(null, ProteusException.Code.SESSION_NOT_FOUND)
     }
 
     override suspend fun encryptBatched(message: ByteArray, sessionIds: List<CryptoSessionId>): Map<CryptoSessionId, ByteArray> {
@@ -121,7 +121,7 @@ class ProteusClientCryptoBoxImpl constructor(
     }
 
     @Suppress("TooGenericExceptionCaught")
-    private fun <T> wrapException(b: () -> T): T {
+    private inline fun <T> wrapException(b: () -> T): T {
         try {
             return b()
         } catch (e: CryptoException) {

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

@@ -59,44 +59,43 @@ internal class NewMessageEventHandlerImpl(
 
     override suspend fun handleNewProteusMessage(event: Event.Conversation.NewMessage, deliveryInfo: EventDeliveryInfo) {
         val eventLogger = logger.createEventProcessingLogger(event)
-        proteusMessageUnpacker.unpackProteusMessage(event)
-            .onFailure {
-                val logMap = mapOf(
-                    "event" to event.toLogMap(),
-                    "errorInfo" to "$it",
-                    "protocol" to "Proteus"
-                )
-
-                if (it is ProteusFailure && it.proteusException.code == ProteusException.Code.DUPLICATE_MESSAGE) {
-                    logger.i("Ignoring duplicate event: ${logMap.toJsonElement()}")
-                    return
-                }
+        proteusMessageUnpacker.unpackProteusMessage(event) {
+            processApplicationMessage(it, deliveryInfo)
+            it
+        }.onSuccess {
+            eventLogger.logSuccess(
+                "protocol" to "Proteus",
+                "messageType" to it.messageTypeDescription,
+            )
+        }.onFailure {
+            val logMap = mapOf(
+                "event" to event.toLogMap(),
+                "errorInfo" to "$it",
+                "protocol" to "Proteus"
+            )
 
-                logger.e("Failed to decrypt event: ${logMap.toJsonElement()}")
+            if (it is ProteusFailure && it.proteusException.code == ProteusException.Code.DUPLICATE_MESSAGE) {
+                logger.i("Ignoring duplicate event: ${logMap.toJsonElement()}")
+                return
+            }
 
-                applicationMessageHandler.handleDecryptionError(
-                    eventId = event.id,
-                    conversationId = event.conversationId,
-                    messageInstant = event.messageInstant,
+            logger.e("Failed to decrypt event: ${logMap.toJsonElement()}")
+
+            applicationMessageHandler.handleDecryptionError(
+                eventId = event.id,
+                conversationId = event.conversationId,
+                messageInstant = event.messageInstant,
+                senderUserId = event.senderUserId,
+                senderClientId = event.senderClientId,
+                content = MessageContent.FailedDecryption(
+                    encodedData = event.encryptedExternalContent?.data,
+                    isDecryptionResolved = false,
                     senderUserId = event.senderUserId,
-                    senderClientId = event.senderClientId,
-                    content = MessageContent.FailedDecryption(
-                        encodedData = event.encryptedExternalContent?.data,
-                        isDecryptionResolved = false,
-                        senderUserId = event.senderUserId,
-                        clientId = ClientId(event.senderClientId.value)
-                    )
-                )
-                eventLogger.logFailure(it, "protocol" to "Proteus")
-            }.onSuccess {
-                if (it is MessageUnpackResult.ApplicationMessage) {
-                    processApplicationMessage(it, deliveryInfo)
-                }
-                eventLogger.logSuccess(
-                    "protocol" to "Proteus",
-                    "messageType" to it.messageTypeDescription,
+                    clientId = ClientId(event.senderClientId.value)
                 )
-            }
+            )
+            eventLogger.logFailure(it, "protocol" to "Proteus")
+        }
     }
 
     override suspend fun handleNewMLSMessage(event: Event.Conversation.NewMLSMessage, deliveryInfo: EventDeliveryInfo) {