Merge pull request #7 from Recouse/message-parser-fix

Message parser fix

Sources/EventSource/EventSource.swift

@@ -50,6 +50,8 @@ public final class EventSource {
 
     private var currentRetryCount: Int = 1
 
+    public private(set) var lastMessageId: String = ""
+
     /// Server-sent events channel.
     public let events: AsyncChannel<ChannelSubject> = .init()
 
@@ -58,7 +60,7 @@ public final class EventSource {
         configuration.httpAdditionalHeaders = [
             HTTPHeaderField.accept: Accept.eventStream,
             HTTPHeaderField.cacheControl: CacheControl.noStore,
-            HTTPHeaderField.lastEventID: messageParser.lastMessageId
+            HTTPHeaderField.lastEventID: lastMessageId
         ]
         configuration.timeoutIntervalForRequest = self.timeoutInterval
         configuration.timeoutIntervalForResource = self.timeoutInterval
@@ -77,7 +79,7 @@ public final class EventSource {
 
     public init(
         request: URLRequest,
-        messageParser: MessageParser = .init(),
+        messageParser: MessageParser = .live,
         maxRetryCount: Int = 3,
         retryDelay: Double = 1.0,
         timeoutInterval: TimeInterval = 300
@@ -185,7 +187,7 @@ public final class EventSource {
     public func close() async {
         let previousState = readyState
         readyState = .closed
-        messageParser.reset()
+        lastMessageId = ""
         sessionDelegateTask?.cancel()
         dataTask?.cancel()
         dataTask = nil
@@ -208,7 +210,12 @@ public final class EventSource {
             return
         }
 
-        let messages = messageParser.parsed(from: data)
+        let messages = messageParser.parse(data)
+
+        // Update last message ID
+        if let lastMessageWithId = messages.last(where: { $0.id != nil }) {
+            lastMessageId = lastMessageWithId.id ?? ""
+        }
 
         await messages.asyncForEach {
             await events.send(.message($0))

Sources/EventSource/MeesageParser.swift → Sources/EventSource/MessageParser.swift

@@ -8,17 +8,15 @@
 
 import Foundation
 
-public class MessageParser {
-    public static let lf: UInt8 = 0x0A
-    public static let semicolon: UInt8 = 0x3a
-
-    public private(set) var lastMessageId: String = ""
-
-    public init() {
-
-    }
+public struct MessageParser {
+    public var parse: (_ data: Data) -> [ServerMessage]
+}
+
+public extension MessageParser {
+    static let lf: UInt8 = 0x0A
+    static let colon: UInt8 = 0x3A
 
-    public func parsed(from data: Data) -> [ServerMessage] {
+    static let live = Self(parse: { data in
         // Split message with double newline
         let rawMessages: [Data]
         if #available(macOS 13.0, iOS 16.0, watchOS 9.0, tvOS 16.0, visionOS 1.0, *) {
@@ -30,17 +28,8 @@ public class MessageParser {
         // Parse data to ServerMessage model
         let messages: [ServerMessage] = rawMessages.compactMap(ServerMessage.parse(from:))
 
-        // Update last message ID
-        if let lastMessageWithId = messages.last(where: { $0.id != nil }) {
-            lastMessageId = lastMessageWithId.id ?? ""
-        }
-
         return messages
-    }
-
-    public func reset() {
-        lastMessageId = ""
-    }
+    })
 }
 
 fileprivate extension Data {

Sources/EventSource/ServerMessage.swift

@@ -12,17 +12,20 @@ public struct ServerMessage {
     public var id: String?
     public var event: String?
     public var data: String?
+    public var other: [String: String]?
     public var time: String?
 
     init(
         id: String? = nil,
         event: String? = nil,
         data: String? = nil,
+        other: [String: String]? = nil,
         time: String? = nil
     ) {
         self.id = id
         self.event = event
         self.data = data
+        self.other = other
         self.time = time
     }
 
@@ -39,6 +42,10 @@ public struct ServerMessage {
             return false
         }
 
+        if let other, !other.isEmpty {
+            return false
+        }
+
         if let time, !time.isEmpty {
             return false
         }
@@ -52,25 +59,41 @@ public struct ServerMessage {
         var message = ServerMessage()
 
         for row in rows {
-            let keyValue = row.split(separator: MessageParser.semicolon, maxSplits: 1)
-            let key = keyValue[0].utf8String
-            let value = keyValue[1].utf8String
+            // Skip the line if it is empty or it starts with a colon character
+            if row.isEmpty, row.first == MessageParser.colon {
+                continue
+            }
+
+            let keyValue = row.split(separator: MessageParser.colon, maxSplits: 1)
+            let key = keyValue[0].utf8String.trimmingCharacters(in: .whitespaces)
+            let value = keyValue[safe: 1]?.utf8String.trimmingCharacters(in: .whitespaces)
 
             switch key {
             case "id":
-                message.id = value.trimmingCharacters(in: .whitespacesAndNewlines)
+                message.id = value?.trimmingCharacters(in: .whitespaces)
             case "event":
-                message.event = value.trimmingCharacters(in: .whitespacesAndNewlines)
+                message.event = value?.trimmingCharacters(in: .whitespaces)
             case "data":
                 if let existingData = message.data {
-                    message.data = existingData + "\n" + value.trimmingCharacters(in: .whitespacesAndNewlines)
+                    message.data = existingData + "\n" + (value?.trimmingCharacters(in: .whitespaces) ?? "")
                 } else {
-                    message.data = value.trimmingCharacters(in: .whitespacesAndNewlines)
+                    message.data = value?.trimmingCharacters(in: .whitespaces)
                 }
             case "time":
-                message.time = value.trimmingCharacters(in: .whitespacesAndNewlines)
+                message.time = value?.trimmingCharacters(in: .whitespaces)
             default:
-                continue
+                // If the line is not empty but does not contain a color character
+                // add it to the other fields using the whole line as the field name,
+                // and the empty string as the field value.
+                if row.contains(MessageParser.colon) == false {
+                    let string = row.utf8String
+                    if var other = message.other {
+                        other[string] = ""
+                        message.other = other
+                    } else {
+                        message.other = [string: ""]
+                    }
+                }
             }
         }
 
@@ -87,3 +110,12 @@ fileprivate extension Data {
         String(decoding: self, as: UTF8.self)
     }
 }
+
+fileprivate extension Array {
+    subscript(safe index: Int) -> Element? {
+        guard index >= 0, index < endIndex else {
+            return nil
+        }
+        return self[index]
+    }
+}

Tests/EventSourceTests/MessageParserTests.swift

@@ -8,7 +8,7 @@ import XCTest
 
 final class MessageParserTests: XCTestCase {
     func testMessagesParsing() throws {
-        let parser = MessageParser()
+        let parser = MessageParser.live
 
         let text = """
         data: test 1
@@ -28,7 +28,7 @@ final class MessageParserTests: XCTestCase {
         """
         let data = Data(text.utf8)
 
-        let messages = parser.parsed(from: data)
+        let messages = parser.parse(data)
 
         XCTAssertEqual(messages.count, 5)
 
@@ -57,45 +57,103 @@ final class MessageParserTests: XCTestCase {
     }
 
     func testEmptyData() {
-        let parser = MessageParser()
+        let parser = MessageParser.live
 
         let text = """
         
         
         """
         let data = Data(text.utf8)
 
-        let messages = parser.parsed(from: data)
+        let messages = parser.parse(data)
 
         XCTAssertTrue(messages.isEmpty)
     }
 
-    func testLastEventId() {
-        let parser = MessageParser()
+    func testOtherMessageFormats() {
+        let parser = MessageParser.live
 
         let text = """
-        data: test 1
+        data : test 1
 
-        id: 2
-        data: test 2
+        id : 2
+         data : test 2
 
-        event: add
-        data: test 3
+         event : add
+        data :  test 3
 
-        id: 3
-        event: ping
-        data: test 3
+        id : 4
+        event : ping
+        data : test 4
+        
+        test 5
+        
+        message 6
+        message 6-1
         """
         let data = Data(text.utf8)
+
+        let messages = parser.parse(data)
 
-        XCTAssertEqual(parser.lastMessageId, "")
+        XCTAssertNotNil(messages[0].data)
+        XCTAssertEqual(messages[0].data!, "test 1")
 
-        let _ = parser.parsed(from: data)
+        XCTAssertNotNil(messages[1].id)
+        XCTAssertNotNil(messages[1].data)
+        XCTAssertEqual(messages[1].id!, "2")
+        XCTAssertEqual(messages[1].data!, "test 2")
 
-        XCTAssertEqual(parser.lastMessageId, "3")
+        XCTAssertNotNil(messages[2].event)
+        XCTAssertNotNil(messages[2].data)
+        XCTAssertEqual(messages[2].event!, "add")
+        XCTAssertEqual(messages[2].data!, "test 3")
 
-        parser.reset()
+        XCTAssertNotNil(messages[3].id)
+        XCTAssertNotNil(messages[3].event)
+        XCTAssertNotNil(messages[3].data)
+        XCTAssertEqual(messages[3].id!, "4")
+        XCTAssertEqual(messages[3].event!, "ping")
+        XCTAssertEqual(messages[3].data!, "test 4")
 
-        XCTAssertEqual(parser.lastMessageId, "")
+        XCTAssertNotNil(messages[4].other)
+        XCTAssertEqual(messages[4].other!["test 5"], "")
+
+        XCTAssertNotNil(messages[5].other)
+        XCTAssertEqual(messages[5].other!["message 6"], "")
+        XCTAssertEqual(messages[5].other!["message 6-1"], "")
+    }
+
+    func testJSONData() {
+        let parser = MessageParser.live
+        let jsonDecoder = JSONDecoder()
+
+        let text = """
+        data: {\"id\":\"abcd-1\",\"type\":\"message\",\"content\":\"\\ntest\\n\"}
+
+        id: abcd-2
+        data: {\"id\":\"abcd-2\",\"type\":\"message\",\"content\":\"\\n\\n"}
+        
+        """
+        let data = Data(text.utf8)
+
+        let messages = parser.parse(data)
+
+        XCTAssertNotNil(messages[0].data)
+        XCTAssertNotNil(messages[1].data)
+
+        do {
+            let decoded1 = try jsonDecoder.decode(TestModel.self, from: Data(messages[0].data!.utf8))
+            let decoded2 = try jsonDecoder.decode(TestModel.self, from: Data(messages[1].data!.utf8))
+        } catch {
+            XCTFail("The JSON strings provided in the test data were parsed incorrectly.")
+        }
+    }
+}
+
+fileprivate extension MessageParserTests {
+    struct TestModel: Decodable {
+        let id: String
+        let type: String
+        let content: String
     }
 }