feat(reactive-rpc): 🎸 add support for fragmented ws messages

Author: streamich

src/reactive-rpc/server/ws/server/WsServerConnection.ts

@@ -2,7 +2,6 @@ import * as crypto from 'crypto';
 import * as stream from 'stream';
 import {WsCloseFrame, WsFrameDecoder, WsFrameHeader, WsFrameOpcode, WsPingFrame, WsPongFrame} from '../codec';
 import {utf8Size} from '../../../../util/strings/utf8';
-import {FanOut} from 'thingies/es2020/fanout';
 import type {WsFrameEncoder} from '../codec/WsFrameEncoder';
 
 export type WsServerConnectionSocket = stream.Duplex;
@@ -12,12 +11,6 @@ export class WsServerConnection {
   public maxIncomingMessage: number = 2 * 1024 * 1024;
   public maxBackpressure: number = 2 * 1024 * 1024;
 
-  /**
-   * If this is not null, then the connection is receiving a stream: a sequence
-   * of fragment frames.
-   */
-  protected stream: FanOut<Uint8Array> | null = null;
-
   public readonly defaultOnPing = (data: Uint8Array | null): void => {
     this.sendPong(data);
   };
@@ -30,49 +23,64 @@ export class WsServerConnection {
 
   constructor(protected readonly encoder: WsFrameEncoder, public readonly socket: WsServerConnectionSocket) {
     const decoder = new WsFrameDecoder();
-    let currentFrame: WsFrameHeader | null = null;
+    let currentFrameHeader: WsFrameHeader | null = null;
+    let fragmentStartFrameHeader: WsFrameHeader | null = null;
     const handleData = (data: Uint8Array): void => {
       try {
         decoder.push(data);
-        if (currentFrame) {
-          const length = currentFrame.length;
-          if (length <= decoder.reader.size()) {
-            const buf = new Uint8Array(length);
-            decoder.copyFrameData(currentFrame, buf, 0);
-            const isText = currentFrame.opcode === WsFrameOpcode.TEXT;
-            currentFrame = null;
-            this.onmessage(buf, isText);
+        main: while (true) {
+          if (currentFrameHeader instanceof WsFrameHeader) {
+            const length = currentFrameHeader.length;
+            if (length > this.maxIncomingMessage) {
+              this.onClose(1009, 'TOO_LARGE');
+              return;
+            }
+            if (length <= decoder.reader.size()) {
+              const buf = new Uint8Array(length);
+              decoder.copyFrameData(currentFrameHeader, buf, 0);
+              if (fragmentStartFrameHeader instanceof WsFrameHeader) {
+                const isText = fragmentStartFrameHeader.opcode === WsFrameOpcode.TEXT;
+                const isLast = currentFrameHeader.fin === 1;
+                currentFrameHeader = null;
+                if (isLast) fragmentStartFrameHeader = null;
+                this.onfragment(isLast, buf, isText)
+              } else {
+                const isText = currentFrameHeader.opcode === WsFrameOpcode.TEXT;
+                currentFrameHeader = null;
+                this.onmessage(buf, isText);
+              }
+            } else break;
           }
-        }
-        while (true) {
           const frame = decoder.readFrameHeader();
           if (!frame) break;
-          else if (frame instanceof WsPingFrame) this.onping(frame.data);
-          else if (frame instanceof WsPongFrame) this.onpong(frame.data);
-          else if (frame instanceof WsCloseFrame) {
+          if (frame instanceof WsPingFrame) {
+            this.onping(frame.data);
+            continue main;
+          }
+          if (frame instanceof WsPongFrame) {
+            this.onpong(frame.data);
+            continue main;
+          }
+          if (frame instanceof WsCloseFrame) {
             decoder.readCloseFrameData(frame);
             this.onClose(frame.code, frame.reason);
-          } else if (frame instanceof WsFrameHeader) {
-            if (this.stream) {
+            continue main;
+          }
+          if (frame instanceof WsFrameHeader) {
+            if (fragmentStartFrameHeader) {
               if (frame.opcode !== WsFrameOpcode.CONTINUE) {
                 this.onClose(1002, 'DATA');
                 return;
               }
-              throw new Error('streaming not implemented');
+              currentFrameHeader = frame;
             }
-            const length = frame.length;
-            if (length > this.maxIncomingMessage) {
-              this.onClose(1009, 'TOO_LARGE');
-              return;
-            }
-            if (length <= decoder.reader.size()) {
-              const buf = new Uint8Array(length);
-              decoder.copyFrameData(frame, buf, 0);
-              const isText = frame.opcode === WsFrameOpcode.TEXT;
-              this.onmessage(buf, isText);
-            } else {
-              currentFrame = frame;
+            if (frame.fin === 0) {
+              fragmentStartFrameHeader = frame;
+              currentFrameHeader = frame;
+              continue main;
             }
+            currentFrameHeader = frame;
+            continue main;
           }
         }
       } catch (error) {

src/reactive-rpc/server/ws/server/__tests__/WsServerConnection.spec.ts

@@ -4,6 +4,7 @@ import {WsFrameEncoder} from '../../codec/WsFrameEncoder';
 import {until} from 'thingies';
 import {WsFrameOpcode} from '../../codec';
 import {bufferToUint8Array} from '../../../../../util/buffers/bufferToUint8Array';
+import {listToUint8} from '../../../../../util/buffers/concat';
 
 const setup = () => {
   const socket = new stream.PassThrough();
@@ -80,15 +81,40 @@ describe('.onmessage', () => {
         messages.push([data, isUtf8]);
       };
       const pingFrame = encoder.encodePong(Buffer.from([0x01, 0x02, 0x03]));
-      const closeFrame = encoder.encodeHdr(1, WsFrameOpcode.BINARY, 3, 0);
+      const frame = encoder.encodeHdr(1, WsFrameOpcode.BINARY, 3, 0);
       encoder.writer.buf(Buffer.from([0x01, 0x02, 0x03]), 3);
       const payload = encoder.writer.flush();
       socket.write(pingFrame);
-      socket.write(closeFrame);
+      socket.write(frame);
       socket.write(payload);
       await until(() => messages.length === 1);
       expect(messages[0]).toEqual([new Uint8Array([0x01, 0x02, 0x03]), false]);
     });
+
+    test('two binary data frames', async () => {
+      const {socket, encoder, connection} = setup();
+      const messages: [data: Uint8Array, isUtf8: boolean][] = [];
+      connection.onmessage = (data: Uint8Array, isUtf8: boolean): void => {
+        messages.push([data, isUtf8]);
+      };
+      const pingFrame = encoder.encodePong(Buffer.from([0x01, 0x02, 0x03]));
+      const frame1 = encoder.encodeHdr(1, WsFrameOpcode.BINARY, 3, 0);
+      encoder.writer.buf(Buffer.from([0x01, 0x02, 0x03]), 3);
+      const payload1 = encoder.writer.flush();
+      const frame2 = encoder.encodeHdr(1, WsFrameOpcode.BINARY, 3, 0);
+      encoder.writer.buf(Buffer.from([0x04, 0x05, 0x06]), 3);
+      const payload2 = encoder.writer.flush();
+      socket.write(pingFrame);
+      socket.write(listToUint8([
+        frame1,
+        payload1,
+        frame2,
+        payload2,
+      ]));
+      await until(() => messages.length === 2);
+      expect(messages[0]).toEqual([new Uint8Array([0x01, 0x02, 0x03]), false]);
+      expect(messages[1]).toEqual([new Uint8Array([0x04, 0x05, 0x06]), false]);
+    });
 
     test('text frame', async () => {
       const {socket, encoder, connection} = setup();
@@ -148,3 +174,36 @@ describe('.onmessage', () => {
     });
   });
 });
+
+describe('.fragment', () => {
+  test('parses out message fragments', async () => {
+    const {socket, encoder, connection} = setup();
+    const fragments: [isLast: boolean, data: Uint8Array, isUtf8: boolean][] = [];
+    connection.onfragment = (isLast: boolean, data: Uint8Array, isUtf8: boolean): void => {
+      fragments.push([isLast, data, isUtf8]);
+    };
+    const pingFrame = encoder.encodePong(Buffer.from([0x01, 0x02, 0x03]));
+    const buf1 = encoder.encodeHdr(0, WsFrameOpcode.BINARY, 3, 0);
+    encoder.writer.buf(Buffer.from([0x01, 0x02, 0x03]), 3);
+    const buf2 = encoder.writer.flush();
+    const buf3 = encoder.encodeHdr(0, WsFrameOpcode.CONTINUE, 3, 0);
+    encoder.writer.buf(Buffer.from([0x04, 0x05, 0x06]), 3);
+    const buf4 = encoder.writer.flush();
+    const buf5 = encoder.encodeHdr(1, WsFrameOpcode.CONTINUE, 3, 0);
+    encoder.writer.buf(Buffer.from([0x07, 0x08, 0x09]), 3);
+    const buf6 = encoder.writer.flush();
+    socket.write(pingFrame);
+    socket.write(buf1);
+    socket.write(buf2);
+    socket.write(buf3);
+    socket.write(buf4);
+    socket.write(buf5);
+    socket.write(buf6);
+    await until(() => fragments.length === 3);
+    expect(fragments).toEqual([
+      [false, new Uint8Array([0x01, 0x02, 0x03]), false],
+      [false, new Uint8Array([0x04, 0x05, 0x06]), false],
+      [true, new Uint8Array([0x07, 0x08, 0x09]), false],
+    ]);
+  });
+});