From 088a68a6be1ddf30cb47361561bdbddd1a30bac6 Mon Sep 17 00:00:00 2001
From: Daniel McCartney <dmccartney@gmail.com>
Date: Wed, 23 Aug 2023 19:55:54 +0300
Subject: [PATCH] fix: properly discard malformed messages during batch calls

---
 Sources/XMTP/Conversations.swift        |  8 +++----
 Tests/XMTPTests/ConversationTests.swift | 29 +++++++++++++++++++++++++
 XMTP.podspec                            |  2 +-
 3 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/Sources/XMTP/Conversations.swift b/Sources/XMTP/Conversations.swift
index 95a9d26a..eb5e9c4b 100644
--- a/Sources/XMTP/Conversations.swift
+++ b/Sources/XMTP/Conversations.swift
@@ -46,18 +46,18 @@ public class Conversations {
         for batch in batches {
             messages += try await client.apiClient.batchQuery(request: batch)
                 .responses.flatMap { (res) in
-                    try res.envelopes.compactMap { (envelope) in
+                    res.envelopes.compactMap { (envelope) in
                         let conversation = conversationsByTopic[envelope.contentTopic]
                         if conversation == nil {
                             print("discarding message, unknown conversation \(envelope)")
                             return nil
                         }
-                        let msg = try conversation?.decode(envelope)
-                        if msg == nil {
+                        do {
+                            return try conversation?.decode(envelope)
+                        } catch {
                             print("discarding message, unable to decode \(envelope)")
                             return nil
                         }
-                        return msg
                     }
                 }
         }
diff --git a/Tests/XMTPTests/ConversationTests.swift b/Tests/XMTPTests/ConversationTests.swift
index c43956d8..4fe88330 100644
--- a/Tests/XMTPTests/ConversationTests.swift
+++ b/Tests/XMTPTests/ConversationTests.swift
@@ -434,6 +434,35 @@ class ConversationTests: XCTestCase {
         XCTAssertEqual(bobConversation.topic, messages[2].topic)
     }
 
+    func testProperlyDiscardBadBatchMessages() async throws {
+
+        guard case let .v2(bobConversation) = try await aliceClient.conversations
+            .newConversation(with: bob.address) else {
+            XCTFail("did not get a v2 conversation for bob")
+            return
+        }
+
+        try await bobConversation.send(content: "Hello")
+
+        // Now we send some garbage and expect it to be properly ignored.
+        try await bobClient.apiClient.publish(envelopes: [
+            Envelope(
+                topic: bobConversation.topic,
+                timestamp: Date(),
+                message: Data([1, 2, 3]) // garbage, malformed message
+            )
+        ])
+
+        try await bobConversation.send(content: "Goodbye")
+
+        let messages = try await aliceClient.conversations.listBatchMessages(
+            topics: [bobConversation.topic : nil]
+        )
+        XCTAssertEqual(2, messages.count)
+        XCTAssertEqual("Goodbye", try messages[0].content())
+        XCTAssertEqual("Hello", try messages[1].content())
+    }
+
 	func testImportV1ConversationFromJS() async throws {
 		let jsExportJSONData = Data("""
 		{
diff --git a/XMTP.podspec b/XMTP.podspec
index b1276f68..ae6b671c 100644
--- a/XMTP.podspec
+++ b/XMTP.podspec
@@ -16,7 +16,7 @@ Pod::Spec.new do |spec|
   #
 
   spec.name         = "XMTP"
-  spec.version      = "0.4.9-alpha0"
+  spec.version      = "0.4.10-alpha0"
   spec.summary      = "XMTP SDK Cocoapod"
 
   # This description is used to generate tags and improve search results.