feat: implement cookie management for the Node.js client

darrachequesne · darrachequesne · commit 5fc88a62d401 · 2023-06-13T17:31:00.000+02:00
When setting the `withCredentials` option to `true`, the Node.js client will now include the cookies in the HTTP requests, making it easier to use it with cookie-based sticky sessions. Related: socketio/socket.io#3812
diff --git a/lib/transports/polling.ts b/lib/transports/polling.ts
@@ -3,7 +3,11 @@ import debugModule from "debug"; // debug()
 import { yeast } from "../contrib/yeast.js";
 import { encode } from "../contrib/parseqs.js";
 import { encodePayload, decodePayload, RawData } from "engine.io-parser";
-import { XHR as XMLHttpRequest } from "./xmlhttprequest.js";
+import {
+  CookieJar,
+  createCookieJar,
+  XHR as XMLHttpRequest,
+} from "./xmlhttprequest.js";
 import { Emitter } from "@socket.io/component-emitter";
 import { SocketOptions } from "../socket.js";
 import { installTimerFunctions, pick } from "../util.js";
@@ -26,6 +30,7 @@ export class Polling extends Transport {
 
   private polling: boolean = false;
   private pollXhr: any;
+  private cookieJar?: CookieJar;
 
   /**
    * XHR Polling constructor.
@@ -56,6 +61,10 @@ export class Polling extends Transport {
      */
     const forceBase64 = opts && opts.forceBase64;
     this.supportsBinary = hasXHR2 && !forceBase64;
+
+    if (this.opts.withCredentials) {
+      this.cookieJar = createCookieJar();
+    }
   }
 
   override get name() {
@@ -251,7 +260,11 @@ export class Polling extends Transport {
    * @private
    */
   request(opts = {}) {
-    Object.assign(opts, { xd: this.xd, xs: this.xs }, this.opts);
+    Object.assign(
+      opts,
+      { xd: this.xd, xs: this.xs, cookieJar: this.cookieJar },
+      this.opts
+    );
     return new Request(this.uri(), opts);
   }
 
@@ -296,7 +309,7 @@ interface RequestReservedEvents {
 }
 
 export class Request extends Emitter<{}, {}, RequestReservedEvents> {
-  private readonly opts: { xd; xs } & SocketOptions;
+  private readonly opts: { xd; xs; cookieJar: CookieJar } & SocketOptions;
   private readonly method: string;
   private readonly uri: string;
   private readonly async: boolean;
@@ -375,6 +388,8 @@ export class Request extends Emitter<{}, {}, RequestReservedEvents> {
         xhr.setRequestHeader("Accept", "*/*");
       } catch (e) {}
 
+      this.opts.cookieJar?.addCookies(xhr);
+
       // ie6 check
       if ("withCredentials" in xhr) {
         xhr.withCredentials = this.opts.withCredentials;
@@ -385,6 +400,10 @@ export class Request extends Emitter<{}, {}, RequestReservedEvents> {
       }
 
       xhr.onreadystatechange = () => {
+        if (xhr.readyState === 3) {
+          this.opts.cookieJar?.parseCookies(xhr);
+        }
+
         if (4 !== xhr.readyState) return;
         if (200 === xhr.status || 1223 === xhr.status) {
           this.onLoad();
diff --git a/lib/transports/xmlhttprequest.browser.ts b/lib/transports/xmlhttprequest.browser.ts
@@ -21,3 +21,5 @@ export function XHR(opts) {
     } catch (e) {}
   }
 }
+
+export function createCookieJar() {}
diff --git a/lib/transports/xmlhttprequest.ts b/lib/transports/xmlhttprequest.ts
@@ -1,3 +1,102 @@
 import * as XMLHttpRequestModule from "xmlhttprequest-ssl";
 
 export const XHR = XMLHttpRequestModule.default || XMLHttpRequestModule;
+
+export function createCookieJar() {
+  return new CookieJar();
+}
+
+interface Cookie {
+  name: string;
+  value: string;
+  expires?: Date;
+}
+
+/**
+ * @see https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Set-Cookie
+ */
+export function parse(setCookieString: string): Cookie {
+  const parts = setCookieString.split("; ");
+  const i = parts[0].indexOf("=");
+
+  if (i === -1) {
+    return;
+  }
+
+  const name = parts[0].substring(0, i).trim();
+
+  if (!name.length) {
+    return;
+  }
+
+  let value = parts[0].substring(i + 1).trim();
+
+  if (value.charCodeAt(0) === 0x22) {
+    // remove double quotes
+    value = value.slice(1, -1);
+  }
+
+  const cookie: Cookie = {
+    name,
+    value,
+  };
+
+  for (let j = 1; j < parts.length; j++) {
+    const subParts = parts[j].split("=");
+    if (subParts.length !== 2) {
+      continue;
+    }
+    const key = subParts[0].trim();
+    const value = subParts[1].trim();
+    switch (key) {
+      case "Expires":
+        cookie.expires = new Date(value);
+        break;
+      case "Max-Age":
+        const expiration = new Date();
+        expiration.setUTCSeconds(
+          expiration.getUTCSeconds() + parseInt(value, 10)
+        );
+        cookie.expires = expiration;
+        break;
+      default:
+      // ignore other keys
+    }
+  }
+
+  return cookie;
+}
+
+export class CookieJar {
+  private cookies = new Map<string, Cookie>();
+
+  public parseCookies(xhr: any) {
+    const values = xhr.getResponseHeader("set-cookie");
+    if (!values) {
+      return;
+    }
+    values.forEach((value) => {
+      const parsed = parse(value);
+      if (parsed) {
+        this.cookies.set(parsed.name, parsed);
+      }
+    });
+  }
+
+  public addCookies(xhr: any) {
+    const cookies = [];
+
+    this.cookies.forEach((cookie, name) => {
+      if (cookie.expires?.getTime() < Date.now()) {
+        this.cookies.delete(name);
+      } else {
+        cookies.push(`${name}=${cookie.value}`);
+      }
+    });
+
+    if (cookies.length) {
+      xhr.setDisableHeaderCheck(true);
+      xhr.setRequestHeader("cookie", cookies.join("; "));
+    }
+  }
+}
diff --git a/test/node.js b/test/node.js
@@ -2,6 +2,8 @@ const path = require("path");
 const { exec } = require("child_process");
 const { Socket } = require("../");
 const { repeat } = require("./util");
+const expect = require("expect.js");
+const { parse } = require("../build/cjs/transports/xmlhttprequest.js");
 
 describe("node.js", () => {
   describe("autoRef option", () => {
@@ -55,4 +57,80 @@ describe("node.js", () => {
       });
     });
   });
+
+  it("should send cookies with withCredentials: true", (done) => {
+    const socket = new Socket("http://localhost:3000", {
+      transports: ["polling"],
+      withCredentials: true,
+    });
+
+    socket.on("open", () => {
+      socket.send("sendHeaders");
+    });
+
+    socket.on("message", (data) => {
+      if (data === "hi") {
+        return;
+      }
+      const headers = JSON.parse(data);
+      expect(headers.cookie).to.eql("1=1; 2=2");
+
+      socket.close();
+      done();
+    });
+  });
+
+  it("should not send cookies with withCredentials: false", (done) => {
+    const socket = new Socket("http://localhost:3000", {
+      transports: ["polling"],
+      withCredentials: false,
+    });
+
+    socket.on("open", () => {
+      socket.send("sendHeaders");
+    });
+
+    socket.on("message", (data) => {
+      if (data === "hi") {
+        return;
+      }
+      const headers = JSON.parse(data);
+      expect(headers.cookie).to.eql(undefined);
+
+      socket.close();
+      done();
+    });
+  });
+});
+
+describe("cookie parsing", () => {
+  it("should parse a simple set-cookie header", () => {
+    const cookieStr = "foo=bar";
+
+    expect(parse(cookieStr)).to.eql({
+      name: "foo",
+      value: "bar",
+    });
+  });
+
+  it("should parse a complex set-cookie header", () => {
+    const cookieStr =
+      "foo=bar; Max-Age=1000; Domain=.example.com; Path=/; Expires=Tue, 01 Jul 2025 10:01:11 GMT; HttpOnly; Secure; SameSite=strict";
+
+    expect(parse(cookieStr)).to.eql({
+      name: "foo",
+      value: "bar",
+      expires: new Date("Tue Jul 01 2025 06:01:11 GMT-0400 (EDT)"),
+    });
+  });
+
+  it("should parse a weird but valid cookie", () => {
+    const cookieStr =
+      "foo=bar=bar&foo=foo&John=Doe&Doe=John; Domain=.example.com; Path=/; HttpOnly; Secure";
+
+    expect(parse(cookieStr)).to.eql({
+      name: "foo",
+      value: "bar=bar&foo=foo&John=Doe&Doe=John",
+    });
+  });
 });
diff --git a/test/support/hooks.js b/test/support/hooks.js
@@ -7,6 +7,7 @@ const { attach } = require("engine.io");
 const { rollup } = require("rollup");
 
 const rollupConfig = require("../../support/rollup.config.umd.js");
+const { serialize } = require("cookie");
 
 let httpServer, engine;
 
@@ -50,11 +51,28 @@ exports.mochaHooks = {
         } else if (data === "give utf8") {
           socket.send("пойду спать всем спокойной ночи");
           return;
+        } else if (data === "sendHeaders") {
+          const headers = socket.transport?.dataReq?.headers;
+          return socket.send(JSON.stringify(headers));
         }
 
         socket.send(data);
       });
     });
+
+    engine.on("initial_headers", (headers) => {
+      headers["set-cookie"] = [
+        serialize("1", "1", { maxAge: 86400 }),
+        serialize("2", "2", {
+          sameSite: true,
+          path: "/",
+          httpOnly: true,
+          secure: true,
+        }),
+        serialize("3", "3", { maxAge: 0 }),
+        serialize("4", "4", { expires: new Date() }),
+      ];
+    });
   },
 
   afterAll() {

Original file line number	Diff line number	Diff line change
`@@ -21,3 +21,5 @@ export function XHR(opts) {`
`21`	`21`	`} catch (e) {}`
`22`	`22`	`}`
`23`	`23`	`}`
	`24`	`+`
	`25`	`+export function createCookieJar() {}`