From 8b2216290330b174c9e67be32765bec0c74769f9 Mon Sep 17 00:00:00 2001
From: Damien Arrachequesne <damien.arrachequesne@gmail.com>
Date: Tue, 2 May 2023 00:26:44 +0200
Subject: [PATCH] fix(uws): prevent crash when using with middlewares

The class used to accumulate the response headers did not expose the
exact same API as its wrapped type, which could lead to the following
error in some rare cases:

> TypeError: Cannot read properties of undefined (reading 'end')
>    at Polling.onDataRequest (build/transports-uws/polling.js:109:53)
>    at Polling.onRequest (build/transports-uws/polling.js:47:18)
>    at callback (build/userver.js:94:56)
>    at uServer.verify (build/server.js:152:9)

Related: https://github.com/socketio/socket.io/issues/4643
---
 lib/userver.ts      |  1 +
 test/middlewares.js | 27 +++++++++++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/lib/userver.ts b/lib/userver.ts
index 6f4872f2..5616f78f 100644
--- a/lib/userver.ts
+++ b/lib/userver.ts
@@ -294,6 +294,7 @@ class ResponseWrapper {
     this.res.writeStatus(status);
     this.statusWritten = true;
     this.writeBufferedHeaders();
+    return this;
   }
 
   public writeHeader(key: string, value: string) {
diff --git a/test/middlewares.js b/test/middlewares.js
index 4045d5d0..586df621 100644
--- a/test/middlewares.js
+++ b/test/middlewares.js
@@ -4,6 +4,7 @@ const request = require("superagent");
 const { WebSocket } = require("ws");
 const helmet = require("helmet");
 const session = require("express-session");
+const { ClientSocket } = require("./common");
 
 describe("middlewares", () => {
   it("should apply middleware (polling)", (done) => {
@@ -291,4 +292,30 @@ describe("middlewares", () => {
       });
     });
   });
+
+  it("should not be receiving data when getting a message longer than maxHttpBufferSize when polling (with a middleware)", (done) => {
+    const opts = {
+      allowUpgrades: false,
+      transports: ["polling"],
+      maxHttpBufferSize: 5,
+    };
+    const engine = listen(opts, (port) => {
+      engine.use((req, res, next) => {
+        next();
+      });
+
+      const socket = new ClientSocket(`ws://localhost:${port}`);
+      engine.on("connection", (conn) => {
+        conn.on("message", () => {
+          done(new Error("Test invalidation (message is longer than allowed)"));
+        });
+      });
+      socket.on("open", () => {
+        socket.send("aasdasdakjhasdkjhasdkjhasdkjhasdkjhasdkjhasdkjha");
+      });
+      socket.on("close", (reason) => {
+        done();
+      });
+    });
+  });
 });