Remove avatars.ts and avatars.test.ts files and update teamCreator.ts

* fix: Collaboration server inaccurately counts connections

* Add integration test

* 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();

server/commands/teamCreator.ts

@@ -3,7 +3,6 @@ import slugify from "slugify";
 import { RESERVED_SUBDOMAINS } from "@shared/utils/domains";
 import { traceFunction } from "@server/logging/tracing";
 import { Team, Event } from "@server/models";
-import { generateAvatarUrl } from "@server/utils/avatars";
 
 type Props = {
   /** The displayed name of the team */
@@ -36,13 +35,10 @@ async function teamCreator({
   ip,
   transaction,
 }: Props): Promise<Team> {
-  // If the service did not provide a logo/avatar then we attempt to generate
-  // one via ClearBit, or fallback to colored initials in worst case scenario
+  // If the service did not provide a logo/avatar then we'll use the default
+  // avatar generation mechanism (colored initials)
   if (!avatarUrl || !avatarUrl.startsWith("http")) {
-    avatarUrl = await generateAvatarUrl({
-      domain,
-      id: subdomain,
-    });
+    avatarUrl = null;
   }
 
   const team = await Team.create(

server/queues/processors/ImportsProcessor.ts

@@ -3,7 +3,12 @@ import chunk from "lodash/chunk";
 import keyBy from "lodash/keyBy";
 import truncate from "lodash/truncate";
 import { Fragment, Node } from "prosemirror-model";
-import { CreateOptions, CreationAttributes, Transaction } from "sequelize";
+import {
+  CreateOptions,
+  CreationAttributes,
+  Transaction,
+  UniqueConstraintError,
+} from "sequelize";
 import { v4 as uuidv4 } from "uuid";
 import { randomElement } from "@shared/random";
 import { ImportInput, ImportTaskInput } from "@shared/schema";
@@ -154,45 +159,59 @@ export default abstract class ImportsProcessor<
    * @returns Promise that resolves when mapping and persistence is completed.
    */
   private async onProcessed(importModel: Import<T>, transaction: Transaction) {
-    const { collections } = await this.createCollectionsAndDocuments({
-      importModel,
-      transaction,
-    });
-
-    // Once all collections and documents are created, update collection's document structure.
-    // This ensures the root documents have the whole subtree available in the structure.
-    for (const collection of collections) {
-      await Document.unscoped().findAllInBatches<Document>(
-        {
-          where: { parentDocumentId: null, collectionId: collection.id },
-          order: [
-            ["createdAt", "DESC"],
-            ["id", "ASC"],
-          ],
-          transaction,
-        },
-        async (documents) => {
-          for (const document of documents) {
-            await collection.addDocumentToStructure(document, 0, {
-              save: false,
-              silent: true,
-              transaction,
-            });
-          }
-        }
-      );
-
-      await collection.save({ silent: true, transaction });
-    }
-
-    importModel.state = ImportState.Completed;
-    importModel.error = null; // unset any error from previous attempts.
-    await importModel.saveWithCtx(
-      createContext({
-        user: importModel.createdBy,
+    try {
+      const { collections } = await this.createCollectionsAndDocuments({
+        importModel,
         transaction,
-      })
-    );
+      });
+
+      // Once all collections and documents are created, update collection's document structure.
+      // This ensures the root documents have the whole subtree available in the structure.
+      for (const collection of collections) {
+        await Document.unscoped().findAllInBatches<Document>(
+          {
+            where: { parentDocumentId: null, collectionId: collection.id },
+            order: [
+              ["createdAt", "DESC"],
+              ["id", "ASC"],
+            ],
+            transaction,
+          },
+          async (documents) => {
+            for (const document of documents) {
+              await collection.addDocumentToStructure(document, 0, {
+                save: false,
+                silent: true,
+                transaction,
+              });
+            }
+          }
+        );
+
+        await collection.save({ silent: true, transaction });
+      }
+
+      importModel.state = ImportState.Completed;
+      importModel.error = null; // unset any error from previous attempts.
+      await importModel.saveWithCtx(
+        createContext({
+          user: importModel.createdBy,
+          transaction,
+        })
+      );
+    } catch (err) {
+      if (err instanceof UniqueConstraintError) {
+        Logger.error(
+          "ImportsProcessor persistence failed due to unique constraint error",
+          err,
+          {
+            fields: err.fields,
+          }
+        );
+      }
+
+      throw err;
+    }
   }
 
   /**

server/utils/avatars.test.ts

@@ -1,9 +0,0 @@
-import { generateAvatarUrl } from "./avatars";
-
-it("should return clearbit url if available", async () => {
-  const url = await generateAvatarUrl({
-    id: "google",
-    domain: "google.com",
-  });
-  expect(url).toBe("https://logo.clearbit.com/google.com");
-});

server/utils/avatars.ts

@@ -1,28 +0,0 @@
-import crypto from "crypto";
-import fetch from "./fetch";
-
-export async function generateAvatarUrl({
-  id,
-  domain,
-}: {
-  id: string;
-  domain?: string;
-}) {
-  // attempt to get logo from Clearbit API. If one doesn't exist then
-  // fall back to using tiley to generate a placeholder logo
-  const hash = crypto.createHash("sha256");
-  hash.update(id);
-  let cbResponse, cbUrl;
-
-  if (domain) {
-    cbUrl = `https://logo.clearbit.com/${domain}`;
-
-    try {
-      cbResponse = await fetch(cbUrl);
-    } catch (err) {
-      // okay
-    }
-  }
-
-  return cbUrl && cbResponse && cbResponse.status === 200 ? cbUrl : null;
-}