Merge pull request #501 from streamich/ws-improvements

WebSocket improvements
streamich · Dec 30, 2023 · 38392b3 · 38392b3
2 parents cc314a5 + c209dc2
commit 38392b3
Show file tree

Hide file tree

Showing 2 changed files with 356 additions and 38 deletions.
diff --git a/src/reactive-rpc/server/ws/server/WsServerConnection.ts b/src/reactive-rpc/server/ws/server/WsServerConnection.ts
@@ -1,73 +1,104 @@
-import * as net from 'net';
 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 {listToUint8} from '../../../../util/buffers/concat';
 import type {WsFrameEncoder} from '../codec/WsFrameEncoder';
 
+export type WsServerConnectionSocket = stream.Duplex;
+
 export class WsServerConnection {
   public closed: boolean = false;
   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);
   };
 
+  private _fragments: Uint8Array[] = [];
+  private _fragmentsSize: number = 0;
+  public readonly defaultOnFragment = (isLast: boolean, data: Uint8Array, isUtf8: boolean): void => {
+    const fragments = this._fragments;
+    this._fragmentsSize += data.length;
+    if (this._fragmentsSize > this.maxIncomingMessage) {
+      this.onClose(1009, 'TOO_LARGE');
+      return;
+    }
+    fragments.push(data);
+    if (!isLast) return;
+    this._fragments = [];
+    this._fragmentsSize = 0;
+    const message = listToUint8(fragments);
+    this.onmessage(message, isUtf8);
+  };
+
   public onmessage: (data: Uint8Array, isUtf8: boolean) => void = () => {};
+  public onfragment: (isLast: boolean, data: Uint8Array, isUtf8: boolean) => void = this.defaultOnFragment;
   public onping: (data: Uint8Array | null) => void = this.defaultOnPing;
   public onpong: (data: Uint8Array | null) => void = () => {};
   public onclose: (code: number, reason: string) => void = () => {};
 
-  constructor(protected readonly encoder: WsFrameEncoder, public readonly socket: net.Socket) {
+  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) this.onClose(frame.code, frame.reason);
-          else if (frame instanceof WsFrameHeader) {
-            if (this.stream) {
+          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);
+            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) {