feat: add the "initial_headers" and "headers" events

Those events will be emitted before the response headers are written to the socket: - "initial_headers": on the first request of the connection - "headers": on all requests (HTTP long-polling and WebSocket upgrade) Syntax: ```js server.on("initial_headers", (headers, req) => { headers["test"] = "123"; headers["set-cookie"] = "mycookie=456"; }); server.on("headers", (headers, req) => { headers["test"] = "789"; }); ``` Related: - #557 - socketio/socket.io#3630
socketio · Apr 30, 2021 · 2527543 · 2527543
1 parent 7096e98
commit 2527543
Show file tree

Hide file tree

Showing 6 changed files with 224 additions and 55 deletions.
diff --git a/README.md b/README.md
@@ -204,6 +204,19 @@ The main server/manager. _Inherits from EventEmitter_.
     - Fired when a new connection is established.
     - **Arguments**
       - `Socket`: a Socket object
+
+- `initial_headers`
+    - Fired on the first request of the connection, before writing the response headers
+    - **Arguments**
+      - `headers` (`Object`): a hash of headers
+      - `req` (`http.IncomingMessage`): the request
+
+- `headers`
+    - Fired on the all requests of the connection, before writing the response headers
+    - **Arguments**
+      - `headers` (`Object`): a hash of headers
+      - `req` (`http.IncomingMessage`): the request
+
 - `connection_error`
     - Fired when an error occurs when establishing the connection.
     - **Arguments**

diff --git a/lib/server.js b/lib/server.js
@@ -84,6 +84,25 @@ class Server extends EventEmitter {
       perMessageDeflate: this.opts.perMessageDeflate,
       maxPayload: this.opts.maxHttpBufferSize
     });
+
+    if (typeof this.ws.on === "function") {
+      this.ws.on("headers", (headersArray, req) => {
+        // note: 'ws' uses an array of headers, while Engine.IO uses an object (response.writeHead() accepts both formats)
+        // we could also try to parse the array and then sync the values, but that will be error-prone
+        const additionalHeaders = {};
+
+        const isInitialRequest = !req._query.sid;
+        if (isInitialRequest) {
+          this.emit("initial_headers", additionalHeaders, req);
+        }
+
+        this.emit("headers", additionalHeaders, req);
+
+        Object.keys(additionalHeaders).forEach(key => {
+          headersArray.push(`${key}: ${additionalHeaders[key]}`);
+        });
+      });
+    }
   }
 
   /**
@@ -328,15 +347,19 @@ class Server extends EventEmitter {
     const socket = new Socket(id, this, transport, req, protocol);
     const self = this;
 
-    if (this.opts.cookie) {
-      transport.on("headers", headers => {
-        headers["Set-Cookie"] = cookieMod.serialize(
-          this.opts.cookie.name,
-          id,
-          this.opts.cookie
-        );
-      });
-    }
+    transport.on("headers", (headers, req) => {
+      const isInitialRequest = !req._query.sid;
+
+      if (isInitialRequest) {
+        if (this.opts.cookie) {
+          headers["Set-Cookie"] = [
+            cookieMod.serialize(this.opts.cookie.name, id, this.opts.cookie)
+          ];
+        }
+        this.emit("initial_headers", headers, req);
+      }
+      this.emit("headers", headers, req);
+    });
 
     transport.onRequest(req);
 

diff --git a/lib/socket.js b/lib/socket.js
@@ -81,49 +81,48 @@ class Socket extends EventEmitter {
    * @api private
    */
   onPacket(packet) {
-    if ("open" === this.readyState) {
-      // export packet event
-      debug("packet");
-      this.emit("packet", packet);
-
-      // Reset ping timeout on any packet, incoming data is a good sign of
-      // other side's liveness
-      this.resetPingTimeout(
-        this.server.opts.pingInterval + this.server.opts.pingTimeout
-      );
+    if ("open" !== this.readyState) {
+      return debug("packet received with closed socket");
+    }
+    // export packet event
+    debug(`received packet ${packet.type}`);
+    this.emit("packet", packet);
+
+    // Reset ping timeout on any packet, incoming data is a good sign of
+    // other side's liveness
+    this.resetPingTimeout(
+      this.server.opts.pingInterval + this.server.opts.pingTimeout
+    );
 
-      switch (packet.type) {
-        case "ping":
-          if (this.transport.protocol !== 3) {
-            this.onError("invalid heartbeat direction");
-            return;
-          }
-          debug("got ping");
-          this.sendPacket("pong");
-          this.emit("heartbeat");
-          break;
-
-        case "pong":
-          if (this.transport.protocol === 3) {
-            this.onError("invalid heartbeat direction");
-            return;
-          }
-          debug("got pong");
-          this.schedulePing();
-          this.emit("heartbeat");
-          break;
-
-        case "error":
-          this.onClose("parse error");
-          break;
-
-        case "message":
-          this.emit("data", packet.data);
-          this.emit("message", packet.data);
-          break;
-      }
-    } else {
-      debug("packet received with closed socket");
+    switch (packet.type) {
+      case "ping":
+        if (this.transport.protocol !== 3) {
+          this.onError("invalid heartbeat direction");
+          return;
+        }
+        debug("got ping");
+        this.sendPacket("pong");
+        this.emit("heartbeat");
+        break;
+
+      case "pong":
+        if (this.transport.protocol === 3) {
+          this.onError("invalid heartbeat direction");
+          return;
+        }
+        debug("got pong");
+        this.schedulePing();
+        this.emit("heartbeat");
+        break;
+
+      case "error":
+        this.onClose("parse error");
+        break;
+
+      case "message":
+        this.emit("data", packet.data);
+        this.emit("message", packet.data);
+        break;
     }
   }
 

diff --git a/lib/transports/polling.js b/lib/transports/polling.js
@@ -385,7 +385,7 @@ class Polling extends Transport {
       headers["X-XSS-Protection"] = "0";
     }
 
-    this.emit("headers", headers);
+    this.emit("headers", headers, req);
     return headers;
   }
 }

diff --git a/lib/transports/websocket.js b/lib/transports/websocket.js
@@ -14,9 +14,6 @@ class WebSocket extends Transport {
     this.socket.on("message", this.onData.bind(this));
     this.socket.once("close", this.onClose.bind(this));
     this.socket.on("error", this.onError.bind(this));
-    this.socket.on("headers", headers => {
-      this.emit("headers", headers);
-    });
     this.writable = true;
     this.perMessageDeflate = null;
   }

diff --git a/test/server.js b/test/server.js
@@ -3137,6 +3137,143 @@ describe("server", () => {
       };
       testForHeaders(headers, done);
     });
+
+    it("should emit a 'initial_headers' event (polling)", done => {
+      const partialDone = createPartialDone(done, 2);
+
+      engine = listen({ cookie: true }, port => {
+        engine.on("initial_headers", (headers, req) => {
+          expect(req.method).to.be("GET");
+          headers["test"] = "123";
+          headers["set-cookie"] = "mycookie=456";
+          partialDone();
+        });
+
+        request
+          .get("http://localhost:%d/engine.io/".s(port))
+          .query({ transport: "polling" })
+          .end((err, res) => {
+            expect(err).to.be(null);
+            expect(res.status).to.be(200);
+            expect(res.headers["test"]).to.be("123");
+            expect(res.headers["set-cookie"].length).to.be(2);
+            expect(res.headers["set-cookie"][1]).to.be("mycookie=456");
+
+            const sid = JSON.parse(res.text.substring(4)).sid;
+
+            request
+              .post("http://localhost:%d/engine.io/".s(port))
+              .query({ transport: "polling", sid })
+              .send("1:6")
+              .end((err, res) => {
+                expect(err).to.be(null);
+                expect(res.status).to.be(200);
+                expect(res.headers["test"]).to.be(undefined);
+                expect(res.headers["set-cookie"]).to.be(undefined);
+                partialDone();
+              });
+          });
+      });
+    });
+
+    it("should emit a 'headers' event (polling)", done => {
+      const partialDone = createPartialDone(done, 3);
+
+      engine = listen({ cookie: true }, port => {
+        engine.on("headers", headers => {
+          headers["test"] = "123";
+          headers["set-cookie"] = "mycookie=456";
+          partialDone();
+        });
+
+        request
+          .get("http://localhost:%d/engine.io/".s(port))
+          .query({ transport: "polling" })
+          .end((err, res) => {
+            expect(err).to.be(null);
+            expect(res.status).to.be(200);
+            expect(res.headers["test"]).to.be("123");
+            expect(res.headers["set-cookie"].length).to.be(2);
+            expect(res.headers["set-cookie"][1]).to.be("mycookie=456");
+
+            const sid = JSON.parse(res.text.substring(4)).sid;
+
+            request
+              .post("http://localhost:%d/engine.io/".s(port))
+              .query({ transport: "polling", sid })
+              .send("1:6")
+              .end((err, res) => {
+                expect(err).to.be(null);
+                expect(res.status).to.be(200);
+                expect(res.headers["set-cookie"].length).to.be(1);
+                expect(res.headers["set-cookie"][0]).to.be("mycookie=456");
+                partialDone();
+              });
+          });
+      });
+    });
+
+    it("should emit a 'initial_headers' event (websocket)", function(done) {
+      if (process.env.EIO_WS_ENGINE === "eiows") {
+        this.skip();
+      }
+      const partialDone = createPartialDone(done, 2);
+
+      engine = listen({ cookie: true }, port => {
+        engine.on("initial_headers", (headers, req) => {
+          expect(req.method).to.be("GET");
+          headers["test"] = "123";
+          headers["set-cookie"] = "mycookie=456";
+          partialDone();
+        });
+
+        client = eioc("ws://localhost:%d".s(port), {
+          transports: ["websocket"]
+        });
+
+        client.transport.ws.on("upgrade", res => {
+          expect(res.headers["test"]).to.be("123");
+          expect(res.headers["set-cookie"].length).to.be(1);
+          expect(res.headers["set-cookie"][0]).to.be("mycookie=456");
+          partialDone();
+        });
+      });
+    });
+
+    it("should emit a single 'initial_headers' event per connection", done => {
+      const partialDone = createPartialDone(done, 2);
+
+      engine = listen(port => {
+        engine.on("initial_headers", () => {
+          partialDone();
+        });
+
+        client = eioc("ws://localhost:%d".s(port));
+
+        client.on("upgrade", () => {
+          partialDone();
+        });
+      });
+    });
+
+    it("should emit several 'headers' events per connection", function(done) {
+      if (process.env.EIO_WS_ENGINE === "eiows") {
+        this.skip();
+      }
+      const partialDone = createPartialDone(done, 4);
+
+      engine = listen(port => {
+        engine.on("headers", () => {
+          partialDone();
+        });
+
+        client = eioc("ws://localhost:%d".s(port));
+
+        client.on("upgrade", () => {
+          partialDone();
+        });
+      });
+    });
   });
 
   describe("cors", () => {