Fix WebSocket continuation frame data loss in processFrames

processFrames matched only .binary and .text opcodes; .continuation frames
fell through to default: break. For any fragmented WebSocket message the
intermediate and final continuation frames were silently discarded, leaving
a partial buffer that then corrupted the next unrelated delivery.

Introduce FrameAccumulator (internal, testable) that accepts all data-frame
opcodes and flushes once fin==true. Use it in both ServerConnectionHandler
and ClientConnectionHandler processFrames. Add FrameAccumulatorTests covering
two- and three-fragment assembly, empty payload, non-fin continuation, and
buffer-clear-after-flush to prevent cross-message contamination.

Co-Authored-By: Mean-E <mean-e@users.noreply.github.com>

Author: JustinMeans

Sources/Transmission/ClientManager.swift

@@ -191,18 +191,17 @@ private enum ClientConnectionHandler {
         system: TransmissionSystem,
         node: RemoteNode
     ) async throws {
-        var buffer = Data()
+        var accumulator = FrameAccumulator()
 
         for try await frame in inbound {
             switch frame.opcode {
             case .text, .binary:
-                var data = frame.data
-                if let bytes = data.readBytes(length: data.readableBytes) {
-                    buffer.append(contentsOf: bytes)
+                if let message = accumulator.feed(frame) {
+                    await system.decodeAndDeliver(data: message, from: node)
                 }
-                if frame.fin {
-                    await system.decodeAndDeliver(data: buffer, from: node)
-                    buffer = Data()
+            case .continuation:
+                if let message = accumulator.feed(frame) {
+                    await system.decodeAndDeliver(data: message, from: node)
                 }
             case .ping:
                 let pong = WebSocketFrame(fin: true, opcode: .pong, data: frame.data)

Sources/Transmission/ServerManager.swift

@@ -140,19 +140,19 @@ private enum ServerConnectionHandler {
         system: TransmissionSystem,
         node: RemoteNode
     ) async throws {
-        var buffer = Data()
+        var accumulator = FrameAccumulator()
 
         for try await frame in inbound {
             switch frame.opcode {
             case .text, .binary:
-                var data = frame.data
-                if let bytes = data.readBytes(length: data.readableBytes) {
-                    buffer.append(contentsOf: bytes)
-                    system.metrics.recordMessageReceived(bytes: bytes.count)
+                if let message = accumulator.feed(frame) {
+                    system.metrics.recordMessageReceived(bytes: message.count)
+                    await system.decodeAndDeliver(data: message, from: node)
                 }
-                if frame.fin {
-                    await system.decodeAndDeliver(data: buffer, from: node)
-                    buffer = Data()
+            case .continuation:
+                if let message = accumulator.feed(frame) {
+                    system.metrics.recordMessageReceived(bytes: message.count)
+                    await system.decodeAndDeliver(data: message, from: node)
                 }
             case .ping:
                 let pong = WebSocketFrame(fin: true, opcode: .pong, data: frame.data)
@@ -168,6 +168,38 @@ private enum ServerConnectionHandler {
     }
 }
 
+/// Reassembles a WebSocket message from one or more frames.
+///
+/// WebSocket messages may be split across multiple frames: the first frame carries
+/// opcode `.binary` or `.text` with `fin == false`, intermediate frames carry
+/// opcode `.continuation` with `fin == false`, and the last fragment carries
+/// opcode `.continuation` with `fin == true`. Callers must feed every data frame
+/// (including continuation frames) into `feed(_:)` and act on the returned `Data`
+/// when a complete message has been assembled.
+///
+/// Previously `processFrames` only matched `.binary` and `.text` opcodes, silently
+/// discarding any `.continuation` frames via `default: break`. This caused:
+/// - Fragmented message data to be silently dropped.
+/// - The partial buffer from the first fragment to be flushed with the NEXT unrelated
+///   message, corrupting every subsequent delivery.
+struct FrameAccumulator {
+    private var buffer = Data()
+
+    /// Feed one inbound frame. Returns the assembled message data when the final
+    /// fragment (or a standalone non-fragmented frame) has been received;
+    /// returns nil when more frames are still needed to complete the message.
+    mutating func feed(_ frame: WebSocketFrame) -> Data? {
+        var frameData = frame.data
+        if let bytes = frameData.readBytes(length: frameData.readableBytes) {
+            buffer.append(contentsOf: bytes)
+        }
+        guard frame.fin else { return nil }
+        let message = buffer
+        buffer = Data()
+        return message
+    }
+}
+
 private enum UpgradeResult: Sendable {
     case websocket(NIOAsyncChannel<WebSocketFrame, WebSocketFrame>, clientNodeID: String?)
     case notUpgraded

Tests/TransmissionTests/FrameAccumulatorTests.swift

@@ -0,0 +1,127 @@
+import Testing
+import Foundation
+import NIO
+import NIOWebSocket
+@testable import Transmission
+
+/// Tests for FrameAccumulator, which reassembles WebSocket message fragments.
+///
+/// Before the fix, processFrames matched only .binary and .text opcodes. Continuation
+/// frames fell through to `default: break`, silently discarding their data and leaving
+/// the partial buffer unflushed. The next unrelated message would then have the stale
+/// bytes prepended, corrupting delivery.
+///
+/// The fix introduces FrameAccumulator which handles .continuation frames identically
+/// to the initial .binary/.text fragment: accumulate bytes, flush when fin == true.
+@Suite("FrameAccumulator Tests")
+struct FrameAccumulatorTests {
+
+    // MARK: - Helpers
+
+    private func makeFrame(bytes: [UInt8], opcode: WebSocketOpcode, fin: Bool) -> WebSocketFrame {
+        var buffer = ByteBuffer()
+        buffer.writeBytes(bytes)
+        return WebSocketFrame(fin: fin, opcode: opcode, data: buffer)
+    }
+
+    // MARK: - Non-fragmented (single frame, fin == true)
+
+    @Test("Single binary frame with fin=true returns assembled data immediately")
+    func singleBinaryFrame() {
+        var acc = FrameAccumulator()
+        let frame = makeFrame(bytes: [0x01, 0x02, 0x03], opcode: .binary, fin: true)
+        let result = acc.feed(frame)
+        #expect(result == Data([0x01, 0x02, 0x03]))
+    }
+
+    @Test("Single text frame with fin=true returns assembled data immediately")
+    func singleTextFrame() {
+        var acc = FrameAccumulator()
+        let frame = makeFrame(bytes: Array("hello".utf8), opcode: .text, fin: true)
+        let result = acc.feed(frame)
+        #expect(result == Data("hello".utf8))
+    }
+
+    @Test("Empty single frame returns empty Data (not nil)")
+    func emptyPayloadFrame() {
+        var acc = FrameAccumulator()
+        let frame = makeFrame(bytes: [], opcode: .binary, fin: true)
+        let result = acc.feed(frame)
+        #expect(result == Data())
+    }
+
+    // MARK: - Fragmented messages (the bug site)
+
+    @Test("Fragmented message: binary(fin=false) + continuation(fin=true) assembles correctly")
+    func twoFragmentMessage() {
+        var acc = FrameAccumulator()
+
+        // First fragment: opcode .binary, fin == false — should return nil (not complete)
+        let first = makeFrame(bytes: [0xAA, 0xBB], opcode: .binary, fin: false)
+        let midResult = acc.feed(first)
+        #expect(midResult == nil, "Intermediate frame must not flush the buffer")
+
+        // Final fragment: opcode .continuation, fin == true — must return assembled data
+        let last = makeFrame(bytes: [0xCC, 0xDD], opcode: .continuation, fin: true)
+        let finalResult = acc.feed(last)
+        #expect(finalResult == Data([0xAA, 0xBB, 0xCC, 0xDD]),
+                "Continuation frame data must be appended and flushed when fin == true")
+    }
+
+    @Test("Three-fragment message assembles in order")
+    func threeFragmentMessage() {
+        var acc = FrameAccumulator()
+
+        let f1 = makeFrame(bytes: [0x01], opcode: .binary, fin: false)
+        let f2 = makeFrame(bytes: [0x02], opcode: .continuation, fin: false)
+        let f3 = makeFrame(bytes: [0x03], opcode: .continuation, fin: true)
+
+        #expect(acc.feed(f1) == nil)
+        #expect(acc.feed(f2) == nil)
+        let result = acc.feed(f3)
+        #expect(result == Data([0x01, 0x02, 0x03]))
+    }
+
+    @Test("Continuation with fin=false does not flush")
+    func continuationWithoutFin() {
+        var acc = FrameAccumulator()
+        let f1 = makeFrame(bytes: [0xDE], opcode: .binary, fin: false)
+        let f2 = makeFrame(bytes: [0xAD], opcode: .continuation, fin: false)
+
+        #expect(acc.feed(f1) == nil)
+        #expect(acc.feed(f2) == nil, "Continuation without fin must not flush")
+    }
+
+    // MARK: - Buffer is cleared after flush (no stale data leaks into next message)
+
+    @Test("Buffer resets after a complete message; next message starts clean")
+    func bufferClearsAfterFlush() {
+        var acc = FrameAccumulator()
+
+        // Message 1: single frame
+        let msg1 = makeFrame(bytes: [0x11], opcode: .binary, fin: true)
+        let result1 = acc.feed(msg1)
+        #expect(result1 == Data([0x11]))
+
+        // Message 2: must NOT contain bytes from message 1
+        let msg2 = makeFrame(bytes: [0x22], opcode: .binary, fin: true)
+        let result2 = acc.feed(msg2)
+        #expect(result2 == Data([0x22]),
+                "Stale bytes from the previous message must not contaminate the next")
+    }
+
+    @Test("Buffer resets after fragmented message; next message starts clean")
+    func bufferClearsAfterFragmentedFlush() {
+        var acc = FrameAccumulator()
+
+        // Fragmented message 1
+        #expect(acc.feed(makeFrame(bytes: [0xAA], opcode: .binary, fin: false)) == nil)
+        let result1 = acc.feed(makeFrame(bytes: [0xBB], opcode: .continuation, fin: true))
+        #expect(result1 == Data([0xAA, 0xBB]))
+
+        // Independent message 2: must not include 0xAA or 0xBB
+        let result2 = acc.feed(makeFrame(bytes: [0xFF], opcode: .binary, fin: true))
+        #expect(result2 == Data([0xFF]),
+                "Previously accumulated bytes must be cleared after fin frame flushes")
+    }
+}