docs

server/collaboration/ConnectionLimitExtension.test.ts

@@ -0,0 +1,109 @@
+import { Server } from "@hocuspocus/server";
+import WebSocket from "ws";
+import EDITOR_VERSION from "@shared/editor/version";
+import { sleep } from "@server/utils/timers";
+import { ConnectionLimitExtension } from "./ConnectionLimitExtension";
+import { EditorVersionExtension } from "./EditorVersionExtension";
+
+jest.mock("@server/env", () => ({
+  COLLABORATION_MAX_CLIENTS_PER_DOCUMENT: 2,
+}));
+
+describe("ConnectionLimitExtension", () => {
+  let server: typeof Server;
+  let extension: ConnectionLimitExtension;
+  const port = 12345;
+  const url = `ws://localhost:${port}`;
+  const documentName = "test";
+
+  beforeEach(async () => {
+    extension = new ConnectionLimitExtension();
+    server = Server.configure({
+      port,
+      extensions: [extension, new EditorVersionExtension()],
+    });
+    await server.listen();
+  });
+
+  afterEach(async () => {
+    await server.destroy();
+  });
+
+  const getConnections = () =>
+    extension.connectionsByDocument.get(documentName)?.size ?? 0;
+
+  const createWebSocket = (editorVersion = EDITOR_VERSION) =>
+    new Promise<WebSocket>((resolve, reject) => {
+      const ws = new WebSocket(
+        `${url}/${documentName}?editorVersion=${editorVersion}`
+      );
+      ws.on("open", () => resolve(ws));
+      ws.on("error", reject);
+    });
+
+  it("should allow connections within limit", async () => {
+    const ws1 = await createWebSocket();
+    const ws2 = await createWebSocket();
+
+    expect(ws1.readyState).toBe(WebSocket.OPEN);
+    expect(ws2.readyState).toBe(WebSocket.OPEN);
+    expect(getConnections()).toBe(2);
+
+    ws1.close();
+    ws2.close();
+
+    await sleep(250);
+    expect(getConnections()).toBe(0);
+  });
+
+  it("should close connections exceeding limit", async () => {
+    const ws1 = await createWebSocket();
+    const ws2 = await createWebSocket();
+
+    const ws3 = await createWebSocket();
+    await sleep(250);
+
+    expect(ws3.readyState).toBe(WebSocket.CLOSED);
+    expect(ws2.readyState).toBe(WebSocket.OPEN);
+    expect(ws1.readyState).toBe(WebSocket.OPEN);
+    expect(getConnections()).toBe(2);
+
+    ws1.close();
+    ws2.close();
+
+    await sleep(250);
+    expect(getConnections()).toBe(0);
+  });
+
+  it("should handle connections closed by other extensions", async () => {
+    const ws1 = await createWebSocket();
+
+    // Create a connection that will be closed by the EditorVersionExtension
+    const ws2 = await createWebSocket("1.0.0");
+
+    ws1.close();
+    ws2.close();
+
+    await sleep(250);
+    expect(getConnections()).toBe(0);
+  });
+
+  it("should allow new connection after disconnect", async () => {
+    const ws1 = await createWebSocket();
+    const ws2 = await createWebSocket();
+
+    ws1.close();
+    await sleep(250);
+    expect(getConnections()).toBe(1);
+
+    const ws3 = await createWebSocket();
+    expect(ws3.readyState).toBe(WebSocket.OPEN);
+    expect(getConnections()).toBe(2);
+
+    ws2.close();
+    ws3.close();
+
+    await sleep(250);
+    expect(getConnections()).toBe(0);
+  });
+});

server/collaboration/ConnectionLimitExtension.ts

@@ -1,8 +1,10 @@
 import {
   Extension,
+  connectedPayload,
   onConnectPayload,
   onDisconnectPayload,
 } from "@hocuspocus/server";
+import pluralize from "pluralize";
 import { TooManyConnections } from "@shared/collaboration/CloseEvents";
 import env from "@server/env";
 import Logger from "@server/logging/Logger";
@@ -14,7 +16,7 @@ export class ConnectionLimitExtension implements Extension {
   /**
    * Map of documentId -> connection count
    */
-  connectionsByDocument: Map<string, Set<string>> = new Map();
+  public connectionsByDocument: Map<string, Set<string>> = new Map();
 
   /**
    * On disconnect hook
@@ -34,23 +36,30 @@ export class ConnectionLimitExtension implements Extension {
       }
     }
 
+    const connectionCount = connections?.size ?? 0;
     Logger.debug(
       "multiplayer",
-      `${connections?.size} connections to "${documentName}"`
+      `${connectionCount} ${pluralize(
+        "connection",
+        connectionCount
+      )} to "${documentName}"`
     );
 
     return Promise.resolve();
   }
 
   /**
-   * On connect hook
+   * onConnect hook is called when a new connection has been established.
+   * This is where we can check if the document has reached the maximum number of
+   * connections and reject the connection if it has.
    *
-   * @param data The connect payload
-   * @returns Promise, resolving will allow the connection, rejecting will drop it
+   * @param data The onConnect payload
+   * @returns Promise, resolving will allow the connection, rejecting will drop.
    */
-  onConnect({ documentName, socketId }: withContext<onConnectPayload>) {
+  onConnect({ documentName }: withContext<onConnectPayload>) {
     const connections =
       this.connectionsByDocument.get(documentName) || new Set();
+
     if (connections?.size >= env.COLLABORATION_MAX_CLIENTS_PER_DOCUMENT) {
       Logger.info(
         "multiplayer",
@@ -61,12 +70,30 @@ export class ConnectionLimitExtension implements Extension {
       return Promise.reject(TooManyConnections);
     }
 
+    return Promise.resolve();
+  }
+
+  /**
+   * Connected hook is called after a new connection has been established.
+   * We can safely update the connection count for the document.
+   *
+   * @param data The onConnect payload
+   * @returns Promise
+   */
+  connected({ documentName, socketId }: withContext<connectedPayload>) {
+    const connections =
+      this.connectionsByDocument.get(documentName) || new Set();
+
     connections.add(socketId);
     this.connectionsByDocument.set(documentName, connections);
+    const connectionCount = connections.size ?? 0;
 
     Logger.debug(
       "multiplayer",
-      `${connections.size} connections to "${documentName}"`
+      `${connectionCount} ${pluralize(
+        "connection",
+        connectionCount
+      )} to "${documentName}"`
     );
 
     return Promise.resolve();