pref: database and e2ee

- fix ff preview for e2e messages

core/src/main/kotlin/me/rhunk/snapenhance/core/ModContext.kt

@@ -27,6 +27,7 @@ import me.rhunk.snapenhance.core.features.Feature
 import me.rhunk.snapenhance.core.logger.CoreLogger
 import me.rhunk.snapenhance.core.manager.impl.ActionManager
 import me.rhunk.snapenhance.core.manager.impl.FeatureManager
+import me.rhunk.snapenhance.core.messaging.CoreMessagingBridge
 import me.rhunk.snapenhance.core.messaging.MessageSender
 import me.rhunk.snapenhance.core.scripting.CoreScriptRuntime
 import me.rhunk.snapenhance.core.util.media.HttpServer
@@ -61,6 +62,7 @@ class ModContext {
     val eventDispatcher = EventDispatcher(this)
     val native = NativeLib()
     val scriptRuntime by lazy { CoreScriptRuntime(androidContext, log) }
+    val messagingBridge = CoreMessagingBridge(this)
 
     val isDeveloper by lazy { config.scripting.developerMode.get() }
 

core/src/main/kotlin/me/rhunk/snapenhance/core/SnapEnhance.kt

@@ -19,7 +19,6 @@ import me.rhunk.snapenhance.core.bridge.loadFromBridge
 import me.rhunk.snapenhance.core.data.SnapClassCache
 import me.rhunk.snapenhance.core.event.events.impl.SnapWidgetBroadcastReceiveEvent
 import me.rhunk.snapenhance.core.event.events.impl.UnaryCallEvent
-import me.rhunk.snapenhance.core.messaging.CoreMessagingBridge
 import me.rhunk.snapenhance.core.util.LSPatchUpdater
 import me.rhunk.snapenhance.core.util.hook.HookStage
 import me.rhunk.snapenhance.core.util.hook.hook
@@ -117,7 +116,7 @@ class SnapEnhance {
                     logCritical(null, throwable)
                 }
             }
-            bridgeClient.registerMessagingBridge(CoreMessagingBridge(this))
+            bridgeClient.registerMessagingBridge(messagingBridge)
 
             reloadConfig()
             actionManager.init()

core/src/main/kotlin/me/rhunk/snapenhance/core/database/DatabaseAccess.kt

@@ -26,6 +26,10 @@ class DatabaseAccess(private val context: ModContext) : Manager {
         context.androidContext.getDatabasePath("main.db")
     }
 
+    private val dmOtherParticipantCache by lazy {
+        getAllDMOtherParticipants().toMutableMap()
+    }
+
     private fun openMain(): SQLiteDatabase {
         return SQLiteDatabase.openDatabase(
             mainDatabase.absolutePath,
@@ -95,7 +99,7 @@ class DatabaseAccess(private val context: ModContext) : Manager {
     val myUserId by lazy {
         safeDatabaseOperation(openArroyo()) { arroyoDatabase: SQLiteDatabase ->
             arroyoDatabase.rawQuery(buildString {
-                append("SELECT * FROM required_values WHERE key = 'USERID'")
+                append("SELECT value FROM required_values WHERE key = 'USERID'")
             }, null).use { query ->
                 if (!query.moveToFirst()) {
                     return@safeDatabaseOperation null
@@ -163,7 +167,7 @@ class DatabaseAccess(private val context: ModContext) : Manager {
     fun getConversationType(conversationId: String): Int? {
         return safeDatabaseOperation(openArroyo()) { database ->
             database.rawQuery(
-                "SELECT * FROM user_conversation WHERE client_conversation_id = ?",
+                "SELECT conversation_type FROM user_conversation WHERE client_conversation_id = ?",
                 arrayOf(conversationId)
             ).use { query ->
                 if (!query.moveToFirst()) {
@@ -186,10 +190,30 @@ class DatabaseAccess(private val context: ModContext) : Manager {
         }
     }
 
+    private fun getAllDMOtherParticipants(): Map<String, String?> {
+        return safeDatabaseOperation(openArroyo()) { arroyoDatabase: SQLiteDatabase ->
+            arroyoDatabase.rawQuery(
+                "SELECT client_conversation_id, user_id FROM user_conversation WHERE conversation_type = 0",
+                null
+            ).use { query ->
+                val participants = mutableMapOf<String, String>()
+                if (!query.moveToFirst()) {
+                    return@safeDatabaseOperation null
+                }
+                do {
+                    participants[query.getString(query.getColumnIndex("client_conversation_id"))] =
+                        query.getString(query.getColumnIndex("user_id"))
+                } while (query.moveToNext())
+                participants
+            }
+        } ?: emptyMap()
+    }
+
     fun getDMOtherParticipant(conversationId: String): String? {
+        if (dmOtherParticipantCache.containsKey(conversationId)) return dmOtherParticipantCache[conversationId]
         return safeDatabaseOperation(openArroyo()) { cursor ->
             cursor.rawQuery(
-                "SELECT * FROM user_conversation WHERE client_conversation_id = ? AND conversation_type = 0",
+                "SELECT user_id FROM user_conversation WHERE client_conversation_id = ? AND conversation_type = 0",
                 arrayOf(conversationId)
             ).use { query ->
                 val participants = mutableListOf<String>()
@@ -201,7 +225,7 @@ class DatabaseAccess(private val context: ModContext) : Manager {
                 } while (query.moveToNext())
                 participants.firstOrNull { it != myUserId }
             }
-        }
+        }.also { dmOtherParticipantCache[conversationId] = it }
     }
 
 
@@ -214,7 +238,7 @@ class DatabaseAccess(private val context: ModContext) : Manager {
     fun getConversationParticipants(conversationId: String): List<String>? {
         return safeDatabaseOperation(openArroyo()) { arroyoDatabase: SQLiteDatabase ->
             arroyoDatabase.rawQuery(
-                "SELECT * FROM user_conversation WHERE client_conversation_id = ?",
+                "SELECT user_id FROM user_conversation WHERE client_conversation_id = ?",
                 arrayOf(conversationId)
             ).use {
                 if (!it.moveToFirst()) {

core/src/main/kotlin/me/rhunk/snapenhance/core/features/impl/experiments/EndToEndEncryption.kt

@@ -12,6 +12,7 @@ import android.view.ViewGroup.MarginLayoutParams
 import android.widget.Button
 import android.widget.TextView
 import me.rhunk.snapenhance.common.data.ContentType
+import me.rhunk.snapenhance.common.data.MessageState
 import me.rhunk.snapenhance.common.data.MessagingRuleType
 import me.rhunk.snapenhance.common.data.RuleState
 import me.rhunk.snapenhance.common.util.protobuf.ProtoEditor
@@ -268,25 +269,26 @@ class EndToEndEncryption : MessagingRuleFeature(
         }.digest()
     }
 
-    private fun messageHook(conversationId: String, messageId: Long, senderId: String, messageContent: MessageContent) {
+    fun tryDecryptMessage(senderId: String, clientMessageId: Long, conversationId: String, contentType: ContentType, messageBuffer: ByteArray): Pair<ContentType, ByteArray> {
-        if (messageContent.contentType != ContentType.STATUS && decryptedMessageCache.containsKey(messageId)) {
+        if (contentType != ContentType.STATUS && decryptedMessageCache.containsKey(clientMessageId)) {
-            val (contentType, buffer) = decryptedMessageCache[messageId]!!
+            return decryptedMessageCache[clientMessageId]!!
-            messageContent.contentType = contentType
-            messageContent.content = buffer
-            return
         }
 
-        val reader = ProtoReader(messageContent.content)
+        val reader = ProtoReader(messageBuffer)
-        messageContent.contentType = fixContentType(messageContent.contentType!!, reader)
+        var outputBuffer = messageBuffer
+        var outputContentType = fixContentType(contentType, reader) ?: contentType
+        val conversationParticipants by lazy {
+            getE2EParticipants(conversationId)
+        }
 
         fun setMessageContent(buffer: ByteArray) {
-            messageContent.content = buffer
+            outputBuffer = buffer
-            messageContent.contentType = fixContentType(messageContent.contentType, ProtoReader(buffer))
+            outputContentType = fixContentType(outputContentType, ProtoReader(buffer)) ?: outputContentType
-            decryptedMessageCache[messageId] = messageContent.contentType!! to buffer
+            decryptedMessageCache[clientMessageId] = outputContentType to buffer
         }
 
         fun replaceMessageText(text: String) {
-            messageContent.content = ProtoWriter().apply {
+            outputBuffer = ProtoWriter().apply {
                 from(2) {
                     addString(1, text)
                 }
@@ -297,9 +299,6 @@ class EndToEndEncryption : MessagingRuleFeature(
         reader.followPath(2, 1) {
             val messageTypeId = getVarInt(1)?.toInt() ?: return@followPath
             val isMe = context.database.myUserId == senderId
-            val conversationParticipants by lazy {
-                getE2EParticipants(conversationId)
-            }
 
             if (messageTypeId == ENCRYPTED_MESSAGE_ID) {
                 runCatching {
@@ -314,7 +313,7 @@ class EndToEndEncryption : MessagingRuleFeature(
                             setMessageContent(
                                 e2eeInterface.decryptMessage(participantId, ciphertext, iv)
                             )
-                            encryptedMessages.add(messageId)
+                            encryptedMessages.add(clientMessageId)
                             return@eachBuffer
                         }
 
@@ -323,14 +322,14 @@ class EndToEndEncryption : MessagingRuleFeature(
                         setMessageContent(
                             e2eeInterface.decryptMessage(senderId, ciphertext, iv)
                         )
-                        encryptedMessages.add(messageId)
+                        encryptedMessages.add(clientMessageId)
                     }
                 }.onFailure {
-                    context.log.error("Failed to decrypt message id: $messageId", it)
+                    context.log.error("Failed to decrypt message id: $clientMessageId", it)
-                    messageContent.contentType = ContentType.CHAT
+                    outputContentType = ContentType.CHAT
-                    messageContent.content = ProtoWriter().apply {
+                    outputBuffer = ProtoWriter().apply {
                         from(2) {
-                            addString(1, "Failed to decrypt message, id=$messageId. Check logcat for more details.")
+                            addString(1, "Failed to decrypt message, id=$clientMessageId. Check logcat for more details.")
                         }
                     }.toByteArray()
                 }
@@ -354,15 +353,23 @@ class EndToEndEncryption : MessagingRuleFeature(
 
             when (messageTypeId) {
                 REQUEST_PK_MESSAGE_ID -> {
-                    pkRequests[messageId] = payload
+                    pkRequests[clientMessageId] = payload
                     replaceMessageText("You just received a public key request. Click below to accept it.")
                 }
                 RESPONSE_SK_MESSAGE_ID -> {
-                    secretResponses[messageId] = payload
+                    secretResponses[clientMessageId] = payload
                     replaceMessageText("Your friend just accepted your public key. Click below to accept the secret.")
                 }
             }
         }
+
+        return outputContentType to outputBuffer
+    }
+
+    private fun messageHook(conversationId: String, messageId: Long, senderId: String, messageContent: MessageContent) {
+        val (contentType, buffer) = tryDecryptMessage(senderId, messageId, conversationId, messageContent.contentType ?: ContentType.CHAT, messageContent.content)
+        messageContent.contentType = contentType
+        messageContent.content = buffer
     }
 
     override fun asyncInit() {
@@ -474,6 +481,7 @@ class EndToEndEncryption : MessagingRuleFeature(
 
         context.event.subscribe(BuildMessageEvent::class, priority = 0) { event ->
             val message = event.message
+            if (message.messageState != MessageState.COMMITTED) return@subscribe
             val conversationId = message.messageDescriptor.conversationId.toString()
             messageHook(
                 conversationId = conversationId,

core/src/main/kotlin/me/rhunk/snapenhance/core/features/impl/ui/FriendFeedMessagePreview.kt

@@ -15,8 +15,10 @@ import me.rhunk.snapenhance.common.util.protobuf.ProtoReader
 import me.rhunk.snapenhance.core.event.events.impl.BindViewEvent
 import me.rhunk.snapenhance.core.features.Feature
 import me.rhunk.snapenhance.core.features.FeatureLoadParams
+import me.rhunk.snapenhance.core.features.impl.experiments.EndToEndEncryption
 import me.rhunk.snapenhance.core.ui.addForegroundDrawable
 import me.rhunk.snapenhance.core.ui.removeForegroundDrawable
+import me.rhunk.snapenhance.core.util.EvictingMap
 import kotlin.math.absoluteValue
 
 @SuppressLint("DiscouragedApi")
@@ -27,12 +29,17 @@ class FriendFeedMessagePreview : Feature("FriendFeedMessagePreview", loadParams
         ).getColor(0, 0)
     }
 
+    private val friendNameCache = EvictingMap<String, String>(100)
+
     private fun getDimens(name: String) = context.resources.getDimensionPixelSize(context.resources.getIdentifier(name, "dimen", Constants.SNAPCHAT_PACKAGE_NAME))
 
     override fun onActivityCreate() {
         val setting = context.config.userInterface.friendFeedMessagePreview
         if (setting.globalState != true) return
 
+        val hasE2EE = context.config.experimental.e2eEncryption.globalState == true
+        val endToEndEncryption by lazy { context.feature(EndToEndEncryption::class) }
+
         val ffItemId = context.resources.getIdentifier("ff_item", "id", Constants.SNAPCHAT_PACKAGE_NAME)
 
         val secondaryTextSize = getDimens("ff_feed_cell_secondary_text_size").toFloat()
@@ -58,17 +65,29 @@ class FriendFeedMessagePreview : Feature("FriendFeedMessagePreview", loadParams
                 frameLayout.removeForegroundDrawable("ffItem")
 
                 val stringMessages = context.database.getMessagesFromConversationId(conversationId, setting.amount.get().absoluteValue)?.mapNotNull { message ->
-                    val messageContainer = message.messageContent
+                    val messageContainer =
-                        ?.let { ProtoReader(it) }
+                        message.messageContent
-                        ?.followPath(4, 4)
+                            ?.let { ProtoReader(it) }
+                            ?.followPath(4, 4)?.let { messageReader ->
+                                takeIf { hasE2EE }?.let takeIf@{
+                                    endToEndEncryption.tryDecryptMessage(
+                                        senderId = message.senderId ?: return@takeIf null,
+                                        clientMessageId = message.clientMessageId.toLong(),
+                                        conversationId =  message.clientConversationId ?: return@takeIf null,
+                                        contentType = ContentType.fromId(message.contentType),
+                                        messageBuffer = messageReader.getBuffer()
+                                    ).second
+                                }?.let { ProtoReader(it) } ?: messageReader
+                            }
                         ?: return@mapNotNull null
 
                     val messageString = messageContainer.getString(2, 1)
                         ?: ContentType.fromMessageContainer(messageContainer)?.name
                         ?: return@mapNotNull null
 
-                    val friendName = context.database.getFriendInfo(message.senderId ?: return@mapNotNull null)?.let { it.displayName?: it.mutableUsername } ?: "Unknown"
+                    val friendName = friendNameCache.getOrPut(message.senderId ?: return@mapNotNull null) {
-
+                        context.database.getFriendInfo(message.senderId ?: return@mapNotNull null)?.let { it.displayName?: it.mutableUsername } ?: "Unknown"
+                    }
                     "$friendName: $messageString"
                 }?.reversed() ?: return@friendFeedItem
 

core/src/main/kotlin/me/rhunk/snapenhance/core/messaging/CoreMessagingBridge.kt

@@ -36,7 +36,7 @@ class CoreMessagingBridge(
                 val callback = CallbackBuilder(
                     context.mappings.getMappedClass("callbacks", "FetchMessageCallback")
                 ).override("onFetchMessageComplete") { param ->
-                    val message = me.rhunk.snapenhance.core.wrapper.impl.Message(param.arg(1)).toBridge()
+                    val message = me.rhunk.snapenhance.core.wrapper.impl.Message(param.arg(0)).toBridge()
                     continuation.resumeWith(Result.success(message))
                 }
                 .override("onServerRequest", shouldUnhook = false) {}
@@ -47,7 +47,7 @@ class CoreMessagingBridge(
                 context.classCache.conversationManager.methods.first { it.name == "fetchMessage" }.invoke(
                     conversationManager,
                     SnapUUID.fromString(conversationId).instanceNonNull(),
-                    clientMessageId,
+                    clientMessageId.toLong(),
                     callback
                 )
             }