From 7c3d82cba23f14ae392169f64b7b5ce9bdf7e256 Mon Sep 17 00:00:00 2001
From: Rick Newton-Rogers <rnro@apple.com>
Date: Thu, 12 Sep 2024 14:15:59 +0100
Subject: [PATCH] review comments

---
 Sources/NIOCore/SingleStepByteToMessageDecoder.swift  |  4 ++--
 .../SingleStepByteToMessageDecoderTest.swift          | 11 +++++++++--
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/Sources/NIOCore/SingleStepByteToMessageDecoder.swift b/Sources/NIOCore/SingleStepByteToMessageDecoder.swift
index fdcaa30b68..47fa34c121 100644
--- a/Sources/NIOCore/SingleStepByteToMessageDecoder.swift
+++ b/Sources/NIOCore/SingleStepByteToMessageDecoder.swift
@@ -233,8 +233,9 @@ public final class NIOSingleStepByteToMessageProcessor<Decoder: NIOSingleStepByt
         }
 
         self._buffer = nil  // To avoid CoW
+        defer { self._buffer = buffer }
+
         let result = try body(&buffer)
-        self._buffer = buffer
         return result
     }
 
@@ -272,7 +273,6 @@ public final class NIOSingleStepByteToMessageProcessor<Decoder: NIOSingleStepByt
             throw ByteToMessageDecoderError.PayloadTooLargeError()
         }
 
-        // force unwrapping is safe because nil case is handled already and asserted above
         if let readerIndex = self._buffer?.readerIndex, readerIndex > 0, self.decoder.shouldReclaimBytes(buffer: self._buffer!) {
             self._buffer!.discardReadBytes()
         }
diff --git a/Tests/NIOCoreTests/SingleStepByteToMessageDecoderTest.swift b/Tests/NIOCoreTests/SingleStepByteToMessageDecoderTest.swift
index c3a190c373..ff4163753a 100644
--- a/Tests/NIOCoreTests/SingleStepByteToMessageDecoderTest.swift
+++ b/Tests/NIOCoreTests/SingleStepByteToMessageDecoderTest.swift
@@ -564,6 +564,7 @@ public final class NIOSingleStepByteToMessageDecoderTest: XCTestCase {
             struct DecodeLastError: Error {}
 
             func decodeLast(buffer: inout NIOCore.ByteBuffer, seenEOF: Bool) throws -> NIOCore.ByteBuffer? {
+                buffer = ByteBuffer() // to allow the decode loop to exit
                 throw DecodeLastError()
             }
 
@@ -574,10 +575,16 @@ public final class NIOSingleStepByteToMessageDecoderTest: XCTestCase {
 
         let decoder = ThrowingOnLastDecoder()
         let b2mp = NIOSingleStepByteToMessageProcessor(decoder)
+        var errorObserved = false
         XCTAssertNoThrow(try b2mp.process(buffer: ByteBuffer(string: "1\n\n2\n3\n")) { line in
-            XCTAssertThrowsError(try b2mp.finishProcessing(seenEOF: true) { _ in }) { error in
-                XCTAssertNotNil(error as? ThrowingOnLastDecoder.DecodeLastError)
+            // We will throw an error to exit the decoding within the nested process call prematurely.
+            // Unless this is carefully handled we can be left in an inconsistent state which the outer call will encounter
+            do {
+                try b2mp.finishProcessing(seenEOF: true) { _ in }
+            } catch let error as ThrowingOnLastDecoder.DecodeLastError {
+                errorObserved = true
             }
         })
+        XCTAssertTrue(errorObserved)
     }
 }