fix: prepend a header to each WebTransport chunk

socketio · Jul 31, 2023 · 8a24ceb · 8a24ceb
1 parent 7dd1350
commit 8a24ceb
Show file tree

Hide file tree

Showing 5 changed files with 118 additions and 114 deletions.
diff --git a/lib/server.ts b/lib/server.ts
@@ -16,12 +16,11 @@ import type { CookieSerializeOptions } from "cookie";
 import type { CorsOptions, CorsOptionsDelegate } from "cors";
 import type { Duplex } from "stream";
 import { WebTransport } from "./transports/webtransport";
-import { TextDecoder } from "util";
+import { createPacketDecoderStream } from "engine.io-parser";
 
 const debug = debugModule("engine");
 
 const kResponseHeaders = Symbol("responseHeaders");
-const TEXT_DECODER = new TextDecoder();
 
 type Transport = "polling" | "websocket";
 
@@ -149,15 +148,13 @@ type Middleware = (
   next: (err?: any) => void
 ) => void;
 
-function parseSessionId(handshake: string) {
-  if (handshake.startsWith("0{")) {
-    try {
-      const parsed = JSON.parse(handshake.substring(1));
-      if (typeof parsed.sid === "string") {
-        return parsed.sid;
-      }
-    } catch (e) {}
-  }
+function parseSessionId(data: string) {
+  try {
+    const parsed = JSON.parse(data);
+    if (typeof parsed.sid === "string") {
+      return parsed.sid;
+    }
+  } catch (e) {}
 }
 
 export abstract class BaseServer extends EventEmitter {
@@ -536,7 +533,11 @@ export abstract class BaseServer extends EventEmitter {
     }
 
     const stream = result.value;
-    const reader = stream.readable.getReader();
+    const transformStream = createPacketDecoderStream(
+      this.opts.maxHttpBufferSize,
+      "nodebuffer"
+    );
+    const reader = stream.readable.pipeThrough(transformStream).getReader();
 
     // reading the first packet of the stream
     const { value, done } = await reader.read();
@@ -546,12 +547,13 @@ export abstract class BaseServer extends EventEmitter {
     }
 
     clearTimeout(timeout);
-    const handshake = TEXT_DECODER.decode(value);
 
-    // handshake is either
-    // "0" => new session
-    // '0{"sid":"xxxx"}' => upgrade
-    if (handshake === "0") {
+    if (value.type !== "open") {
+      debug("invalid WebTransport handshake");
+      return session.close();
+    }
+
+    if (value.data === undefined) {
       const transport = new WebTransport(session, stream, reader);
 
       // note: we cannot use "this.generateId()", because there is no "req" argument
@@ -572,7 +574,7 @@ export abstract class BaseServer extends EventEmitter {
       return;
     }
 
-    const sid = parseSessionId(handshake);
+    const sid = parseSessionId(value.data);
 
     if (!sid) {
       debug("invalid WebTransport handshake");

diff --git a/lib/transports/webtransport.ts b/lib/transports/webtransport.ts
@@ -1,21 +1,9 @@
 import { Transport } from "../transport";
 import debugModule from "debug";
+import { createPacketEncoderStream } from "engine.io-parser";
 
 const debug = debugModule("engine:webtransport");
 
-const BINARY_HEADER = Buffer.of(54);
-
-function shouldIncludeBinaryHeader(packet, encoded) {
-  // 48 === "0".charCodeAt(0) (OPEN packet type)
-  // 54 === "6".charCodeAt(0) (NOOP packet type)
-  return (
-    packet.type === "message" &&
-    typeof packet.data !== "string" &&
-    encoded[0] >= 48 &&
-    encoded[0] <= 54
-  );
-}
-
 /**
  * Reference: https://developer.mozilla.org/en-US/docs/Web/API/WebTransport_API
  */
@@ -24,24 +12,24 @@ export class WebTransport extends Transport {
 
   constructor(private readonly session, stream, reader) {
     super({ _query: { EIO: "4" } });
-    this.writer = stream.writable.getWriter();
+
+    const transformStream = createPacketEncoderStream();
+    transformStream.readable.pipeTo(stream.writable);
+    this.writer = transformStream.writable.getWriter();
+
     (async () => {
-      let binaryFlag = false;
-      while (true) {
-        const { value, done } = await reader.read();
-        if (done) {
-          debug("session is closed");
-          break;
-        }
-        debug("received chunk: %o", value);
-        if (!binaryFlag && value.byteLength === 1 && value[0] === 54) {
-          binaryFlag = true;
-          continue;
+      try {
+        while (true) {
+          const { value, done } = await reader.read();
+          if (done) {
+            debug("session is closed");
+            break;
+          }
+          debug("received chunk: %o", value);
+          this.onPacket(value);
         }
-        this.onPacket(
-          this.parser.decodePacketFromBinary(value, binaryFlag, "nodebuffer")
-        );
-        binaryFlag = false;
+      } catch (e) {
+        debug("error while reading: %s", e.message);
       }
     })();
 
@@ -58,26 +46,20 @@ export class WebTransport extends Transport {
     return true;
   }
 
-  send(packets) {
+  async send(packets) {
     this.writable = false;
 
     for (let i = 0; i < packets.length; i++) {
       const packet = packets[i];
-      const isLast = i + 1 === packets.length;
-
-      this.parser.encodePacketToBinary(packet, (data) => {
-        if (shouldIncludeBinaryHeader(packet, data)) {
-          debug("writing binary header");
-          this.writer.write(BINARY_HEADER);
-        }
-        debug("writing chunk: %o", data);
-        this.writer.write(data);
-        if (isLast) {
-          this.writable = true;
-          this.emit("drain");
-        }
-      });
+      try {
+        await this.writer.write(packet);
+      } catch (e) {
+        debug("error while writing: %s", e.message);
+      }
     }
+
+    this.writable = true;
+    this.emit("drain");
   }
 
   doClose(fn) {

diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -39,7 +39,7 @@
     "cookie": "~0.4.1",
     "cors": "~2.8.5",
     "debug": "~4.3.1",
-    "engine.io-parser": "~5.1.0",
+    "engine.io-parser": "~5.2.0",
     "ws": "~8.11.0"
   },
   "devDependencies": {