fix: properly discard malformed messages during batch calls

xmtp · Aug 23, 2023 · fb8f791 · fb8f791
1 parent 38687a8
commit fb8f791
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 4 deletions.
diff --git 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
@@ -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("""
 		{