chore: Add configurable per-document connection limit extension (#4717)

* chore: Add configurable per-document connection limit extension

* docs

server/collaboration/ConnectionLimitExtension.ts

@@ -0,0 +1,64 @@
+import {
+  Extension,
+  onConnectPayload,
+  onDisconnectPayload,
+} from "@hocuspocus/server";
+import env from "@server/env";
+import Logger from "@server/logging/Logger";
+import { trace } from "@server/logging/tracing";
+
+@trace()
+export class ConnectionLimitExtension implements Extension {
+  /**
+   * Map of documentId -> connection count
+   */
+  connectionsByDocument: Map<string, number> = new Map();
+
+  /**
+   * onDisconnect hook
+   * @param data The disconnect payload
+   */
+  onDisconnect(data: onDisconnectPayload) {
+    const { documentName } = data;
+
+    const currConnections = this.connectionsByDocument.get(documentName) || 0;
+    const newConnections = currConnections - 1;
+    this.connectionsByDocument.set(documentName, newConnections);
+
+    Logger.debug(
+      "multiplayer",
+      `${newConnections} connections to "${documentName}"`
+    );
+
+    return Promise.resolve();
+  }
+
+  /**
+   * onConnect hook
+   * @param data The connect payload
+   */
+  onConnect(data: onConnectPayload) {
+    const { documentName } = data;
+
+    const currConnections = this.connectionsByDocument.get(documentName) || 0;
+    if (currConnections >= env.COLLABORATION_MAX_CLIENTS_PER_DOCUMENT) {
+      Logger.info(
+        "multiplayer",
+        `Rejected connection to "${documentName}" because it has reached the maximum number of connections`
+      );
+
+      // Rejecting the promise will cause the connection to be dropped.
+      return Promise.reject();
+    }
+
+    const newConnections = currConnections + 1;
+    this.connectionsByDocument.set(documentName, newConnections);
+
+    Logger.debug(
+      "multiplayer",
+      `${newConnections} connections to "${documentName}"`
+    );
+
+    return Promise.resolve();
+  }
+}

server/env.ts

@@ -147,6 +147,17 @@ export class Environment {
     process.env.COLLABORATION_URL
   );
 
+  /**
+   * The maximum number of network clients that can be connected to a single
+   * document at once. Defaults to 100.
+   */
+  @IsOptional()
+  @IsNumber()
+  public COLLABORATION_MAX_CLIENTS_PER_DOCUMENT = parseInt(
+    process.env.COLLABORATION_MAX_CLIENTS_PER_DOCUMENT || "100",
+    10
+  );
+
   /**
    * The port that the server will listen on, defaults to 3000.
    */

server/services/collaboration.ts

@@ -5,6 +5,7 @@ import { Server } from "@hocuspocus/server";
 import Koa from "koa";
 import WebSocket from "ws";
 import { DocumentValidation } from "@shared/validations";
+import { ConnectionLimitExtension } from "@server/collaboration/ConnectionLimitExtension";
 import Logger from "@server/logging/Logger";
 import ShutdownHelper, { ShutdownOrder } from "@server/utils/ShutdownHelper";
 import AuthenticationExtension from "../collaboration/AuthenticationExtension";
@@ -28,6 +29,7 @@ export default function init(
     timeout: 30000,
     maxDebounce: 10000,
     extensions: [
+      new ConnectionLimitExtension(),
       new AuthenticationExtension(),
       new PersistenceExtension(),
       new LoggerExtension(),