LineBasedFrameDecoder: can be a NIOSingleStepByteToMessageDecoder (#217)

Co-authored-by: Johannes Weiss <[email protected]>
apple · Feb 6, 2024 · cdd1580 · cdd1580
1 parent 363da63
commit cdd1580
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 1 deletion.
diff --git a/Sources/NIOExtras/LineBasedFrameDecoder.swift b/Sources/NIOExtras/LineBasedFrameDecoder.swift
@@ -30,7 +30,7 @@ import NIOCore
 ///     | ABC | DEF | GHI |
 ///     +-----+-----+-----+
 ///
-public class LineBasedFrameDecoder: ByteToMessageDecoder {
+public class LineBasedFrameDecoder: ByteToMessageDecoder & NIOSingleStepByteToMessageDecoder {
     /// `ByteBuffer` is the expected type passed in.
     public typealias InboundIn = ByteBuffer
     /// `ByteBuffer`s will be passed to the next stage.
@@ -57,6 +57,14 @@ public class LineBasedFrameDecoder: ByteToMessageDecoder {
         }
     }
 
+    /// Decode data in the supplied buffer.
+    /// - Parameters:
+    ///   - buffer: Buffer containing data to decode.
+    /// - Returns: The decoded object or `nil` if we require more bytes.
+    public func decode(buffer: inout NIOCore.ByteBuffer) throws -> NIOCore.ByteBuffer? {
+        return try self.findNextFrame(buffer: &buffer)
+    }
+
     /// Decode all remaining data.
     /// If it is not possible to consume all the data then ``NIOExtrasErrors/LeftOverBytesError`` is reported via `context.fireErrorCaught`
     /// - Parameters:
@@ -72,6 +80,20 @@ public class LineBasedFrameDecoder: ByteToMessageDecoder {
         return .needMoreData
     }
 
+    /// Decode all remaining data.
+    /// If it is not possible to consume all the data then ``NIOExtrasErrors/LeftOverBytesError`` is reported via `context.fireErrorCaught`
+    /// - Parameters:
+    ///   - buffer: Buffer containing the data to decode.
+    ///   - seenEOF: Has end of file been seen.
+    /// - Returns: The decoded object or `nil` if we require more bytes.
+    public func decodeLast(buffer: inout ByteBuffer, seenEOF: Bool) throws -> InboundOut? {
+        let decoded = try self.decode(buffer: &buffer)
+        if buffer.readableBytes > 0 {
+            throw NIOExtrasErrors.LeftOverBytesError(leftOverBytes: buffer)
+        }
+        return decoded
+    }
+
     private func findNextFrame(buffer: inout ByteBuffer) throws -> ByteBuffer? {
         let view = buffer.readableBytesView.dropFirst(self.lastScanOffset)
         // look for the delimiter

diff --git a/Tests/NIOExtrasTests/LineBasedFrameDecoderTest.swift b/Tests/NIOExtrasTests/LineBasedFrameDecoderTest.swift
@@ -214,4 +214,59 @@ class LineBasedFrameDecoderTest: XCTestCase {
             XCTFail("Unexpected error: \(error)")
         }
     }
+
+    func testBasicSingleStep() {
+        let decoder = LineBasedFrameDecoder()
+        let b2mp = NIOSingleStepByteToMessageProcessor(decoder)
+        var callCount = 0
+        XCTAssertNoThrow(try b2mp.process(buffer: ByteBuffer(string: "1\n\n2\n3\n")) { line in
+            callCount += 1
+            switch callCount {
+            case 1:
+                XCTAssertEqual(ByteBuffer(string: "1"), line)
+            case 2:
+                XCTAssertEqual(ByteBuffer(string: ""), line)
+            case 3:
+                XCTAssertEqual(ByteBuffer(string: "2"), line)
+            case 4:
+                XCTAssertEqual(ByteBuffer(string: "3"), line)
+            default:
+                XCTFail("not expecting call no \(callCount)")
+            }
+        })
+    }
+
+    func testBasicSingleStepNoNewlineComingButEOF() {
+        let decoder = LineBasedFrameDecoder()
+        let b2mp = NIOSingleStepByteToMessageProcessor(decoder)
+        XCTAssertNoThrow(try b2mp.process(buffer: ByteBuffer(string: "new newline eva\r")) { line in
+            XCTFail("not taking calls")
+        })
+        XCTAssertThrowsError(try b2mp.finishProcessing(seenEOF: true, { line in
+            XCTFail("not taking calls")
+        })) { error in
+            if let error = error as? NIOExtrasErrors.LeftOverBytesError {
+                XCTAssertEqual(ByteBuffer(string: "new newline eva\r"), error.leftOverBytes)
+            } else {
+                XCTFail("unexpected error: \(error)")
+            }
+        }
+    }
+
+    func testBasicSingleStepNoNewlineOrEOFComing() {
+        let decoder = LineBasedFrameDecoder()
+        let b2mp = NIOSingleStepByteToMessageProcessor(decoder)
+        XCTAssertNoThrow(try b2mp.process(buffer: ByteBuffer(string: "new newline eva\r")) { line in
+            XCTFail("not taking calls")
+        })
+        XCTAssertThrowsError(try b2mp.finishProcessing(seenEOF: false, { line in
+            XCTFail("not taking calls")
+        })) { error in
+            if let error = error as? NIOExtrasErrors.LeftOverBytesError {
+                XCTAssertEqual(ByteBuffer(string: "new newline eva\r"), error.leftOverBytes)
+            } else {
+                XCTFail("unexpected error: \(error)")
+            }
+        }
+    }
 }