@ibgib/space-gib 0.0.3 → 0.0.5

package/src/server/serve-gib/handlers/ws/ws-helper.mts

@@ -116,6 +116,7 @@ export function performConnect(req: IncomingMessage, socket: Socket): boolean {
  */
 export class WebSocketFrameDecoder {
     private buffer: Buffer = Buffer.alloc(0);
+    private fragments: Buffer[] = [];
 
     /** Adds new data chunk from TCP socket. */
     addChunk(chunk: Buffer): void {
@@ -123,66 +124,93 @@ export class WebSocketFrameDecoder {
     }
 
     /**
-     * Parses the next complete text frame from the accumulated buffer.
+     * Parses the next complete text message, assembling fragments if necessary.
      * Returns:
-     * - `string` if a complete text frame was decoded successfully (and consumes those bytes).
+     * - `string` if a complete text message was assembled/decoded (and consumes those bytes).
      * - `null` if the buffer is incomplete or we need more data.
      * - `throws Error` on protocol errors or close frames.
      */
     nextFrame(): string | null {
-        if (this.buffer.length < 2) {
-            return null;
-        }
+        while (true) {
+            if (this.buffer.length < 2) {
+                return null;
+            }
 
-        const opcode = this.buffer[0] & 0x0f;
-        if (opcode === 0x8) {
-            throw new Error('Close frame received');
-        }
-        if (opcode !== 0x1) {
-            throw new Error(`Unsupported WebSocket opcode: 0x${opcode.toString(16)}`);
-        }
+            const firstByte = this.buffer[0];
+            const fin = (firstByte & 0x80) !== 0;
+            const opcode = firstByte & 0x0f;
 
-        const secondByte = this.buffer[1];
-        const masked = (secondByte & 0x80) !== 0;
-        let payloadLen = secondByte & 0x7f;
-        let headerLen = 2;
+            if (opcode === 0x8) {
+                throw new Error('Close frame received');
+            }
 
-        if (payloadLen === 126) {
-            if (this.buffer.length < 4) {
-                return null;
+            const isContinuation = opcode === 0x0;
+            const isText = opcode === 0x1;
+
+            if (!isContinuation && !isText) {
+                throw new Error(`Unsupported WebSocket opcode: 0x${opcode.toString(16)}`);
             }
-            payloadLen = this.buffer.readUInt16BE(2);
-            headerLen = 4;
-        } else if (payloadLen === 127) {
-            if (this.buffer.length < 10) {
-                return null;
+
+            // Protocol checks
+            if (isContinuation && this.fragments.length === 0) {
+                throw new Error('Continuation frame received without an active message.');
+            }
+            if (isText && this.fragments.length > 0) {
+                throw new Error('New message frame received before previous one completed.');
             }
-            payloadLen = Number(this.buffer.readBigUInt64BE(2));
-            headerLen = 10;
-        }
 
-        const maskLen = masked ? 4 : 0;
-        const totalFrameLen = headerLen + maskLen + payloadLen;
+            const secondByte = this.buffer[1];
+            const masked = (secondByte & 0x80) !== 0;
+            let payloadLen = secondByte & 0x7f;
+            let headerLen = 2;
+
+            if (payloadLen === 126) {
+                if (this.buffer.length < 4) {
+                    return null;
+                }
+                payloadLen = this.buffer.readUInt16BE(2);
+                headerLen = 4;
+            } else if (payloadLen === 127) {
+                if (this.buffer.length < 10) {
+                    return null;
+                }
+                payloadLen = Number(this.buffer.readBigUInt64BE(2));
+                headerLen = 10;
+            }
 
-        if (this.buffer.length < totalFrameLen) {
-            return null;
-        }
+            const maskLen = masked ? 4 : 0;
+            const totalFrameLen = headerLen + maskLen + payloadLen;
 
-        const frameData = this.buffer.subarray(0, totalFrameLen);
-        this.buffer = this.buffer.subarray(totalFrameLen);
+            if (this.buffer.length < totalFrameLen) {
+                return null;
+            }
 
-        let payloadStart = headerLen;
-        if (masked) {
-            const mask = frameData.subarray(payloadStart, payloadStart + 4);
-            payloadStart += 4;
-            const payload = Buffer.from(frameData.subarray(payloadStart));
-            for (let i = 0; i < payload.length; i++) {
-                payload[i] ^= mask[i % 4];
+            // Consume frame from buffer
+            const frameData = this.buffer.subarray(0, totalFrameLen);
+            this.buffer = this.buffer.subarray(totalFrameLen);
+
+            let payloadStart = headerLen;
+            let payload: Buffer;
+            if (masked) {
+                const mask = frameData.subarray(payloadStart, payloadStart + 4);
+                payloadStart += 4;
+                payload = Buffer.from(frameData.subarray(payloadStart));
+                for (let i = 0; i < payload.length; i++) {
+                    payload[i] ^= mask[i % 4];
+                }
+            } else {
+                payload = Buffer.from(frameData.subarray(payloadStart));
             }
-            return payload.toString('utf-8');
-        } else {
-            const payload = frameData.subarray(payloadStart);
-            return payload.toString('utf-8');
+
+            this.fragments.push(payload);
+
+            if (fin) {
+                const completePayload = Buffer.concat(this.fragments);
+                this.fragments = [];
+                return completePayload.toString('utf-8');
+            }
+            
+            // Loop again to process any subsequent frames already inside the TCP buffer
         }
     }
 }