From 5b561e98f777808dbe8e686a6f1265b01ea4fa7e Mon Sep 17 00:00:00 2001
From: Tom Moor <tom@getoutline.com>
Date: Sun, 22 Jan 2023 07:50:32 -0800
Subject: [PATCH] chore: Add configurable per-document connection limit
 extension (#4717)

* chore: Add configurable per-document connection limit extension

* docs
---
 .../collaboration/ConnectionLimitExtension.ts | 64 +++++++++++++++++++
 server/env.ts                                 | 11 ++++
 server/services/collaboration.ts              |  2 +
 3 files changed, 77 insertions(+)
 create mode 100644 server/collaboration/ConnectionLimitExtension.ts

diff --git a/server/collaboration/ConnectionLimitExtension.ts b/server/collaboration/ConnectionLimitExtension.ts
new file mode 100644
index 000000000..8e4405666
--- /dev/null
+++ b/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();
+  }
+}
diff --git a/server/env.ts b/server/env.ts
index 1d7b72471..713f1ea34 100644
--- a/server/env.ts
+++ b/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.
    */
diff --git a/server/services/collaboration.ts b/server/services/collaboration.ts
index d446ab6cb..ea2d68fe9 100644
--- a/server/services/collaboration.ts
+++ b/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(),