WIP - POC

apps/client/src/features/editor/components/excalidraw/excalidraw-utils.ts

@@ -1,3 +1,5 @@
+import { ENCRYPTION_KEY_BITS } from "@excalidraw/common";
+
 type LibraryItems = any;
 
 type LibraryPersistedData = {
@@ -25,7 +27,10 @@ export const localStorageLibraryAdapter: LibraryPersistenceAdapter = {
         return JSON.parse(data);
       }
     } catch (e) {
-      console.error("Error downloading Excalidraw library from localStorage", e);
+      console.error(
+        "Error downloading Excalidraw library from localStorage",
+        e,
+      );
     }
     return null;
   },
@@ -40,3 +45,124 @@ export const localStorageLibraryAdapter: LibraryPersistenceAdapter = {
     }
   },
 };
+
+export const blobToArrayBuffer = (blob: Blob): Promise<ArrayBuffer> => {
+  if ("arrayBuffer" in blob) {
+    return blob.arrayBuffer();
+  }
+  // Safari
+  return new Promise((resolve, reject) => {
+    const reader = new FileReader();
+    reader.onload = (event) => {
+      if (!event.target?.result) {
+        return reject(new Error("Couldn't convert blob to ArrayBuffer"));
+      }
+      resolve(event.target.result as ArrayBuffer);
+    };
+    reader.readAsArrayBuffer(blob);
+  });
+};
+
+export const IV_LENGTH_BYTES = 12;
+
+// Pre-transform error: No known conditions for "./data/encryption" specifier in "@excalidraw/excalidraw" package
+//   Plugin: vite:import-analysis
+//   File: /Users/lite/WebstormProjects/docmost-ee/apps/client/src/features/editor/components/excalidraw/use-excalidraw-collab.ts:11:7
+//   7  |    decryptData,
+//   8  |    encryptData
+//   9  |  } from "@excalidraw/excalidraw/data/encryption";
+
+//@ts-ignore
+export const createIV = (): Uint8Array<ArrayBuffer> => {
+  const arr = new Uint8Array(IV_LENGTH_BYTES);
+  return window.crypto.getRandomValues(arr);
+};
+
+export const generateEncryptionKey = async <
+  T extends "string" | "cryptoKey" = "string",
+>(
+  returnAs?: T,
+): Promise<T extends "cryptoKey" ? CryptoKey : string> => {
+  const key = await window.crypto.subtle.generateKey(
+    {
+      name: "AES-GCM",
+      length: ENCRYPTION_KEY_BITS,
+    },
+    true, // extractable
+    ["encrypt", "decrypt"],
+  );
+  return (
+    returnAs === "cryptoKey"
+      ? key
+      : (await window.crypto.subtle.exportKey("jwk", key)).k
+  ) as T extends "cryptoKey" ? CryptoKey : string;
+};
+
+export const getCryptoKey = (key: string, usage: KeyUsage) =>
+  window.crypto.subtle.importKey(
+    "jwk",
+    {
+      alg: "A128GCM",
+      ext: true,
+      k: key,
+      key_ops: ["encrypt", "decrypt"],
+      kty: "oct",
+    },
+    {
+      name: "AES-GCM",
+      length: ENCRYPTION_KEY_BITS,
+    },
+    false, // extractable
+    [usage],
+  );
+
+export const encryptData = async (
+  key: string | CryptoKey,
+  //@ts-ignore
+  data: Uint8Array<ArrayBuffer> | ArrayBuffer | Blob | File | string,
+  //@ts-ignore
+): Promise<{ encryptedBuffer: ArrayBuffer; iv: Uint8Array<ArrayBuffer> }> => {
+  const importedKey =
+    typeof key === "string" ? await getCryptoKey(key, "encrypt") : key;
+  const iv = createIV();
+  //@ts-ignore
+  const buffer: ArrayBuffer | Uint8Array<ArrayBuffer> =
+    typeof data === "string"
+      ? new TextEncoder().encode(data)
+      : data instanceof Uint8Array
+        ? data
+        : data instanceof Blob
+          ? await blobToArrayBuffer(data)
+          : data;
+
+  // We use symmetric encryption. AES-GCM is the recommended algorithm and
+  // includes checks that the ciphertext has not been modified by an attacker.
+  const encryptedBuffer = await window.crypto.subtle.encrypt(
+    {
+      name: "AES-GCM",
+      iv,
+    },
+    importedKey,
+    buffer,
+  );
+
+  return { encryptedBuffer, iv };
+};
+
+export const decryptData = async (
+  //@ts-ignore
+  iv: Uint8Array<ArrayBuffer>,
+  //@ts-ignore
+  encrypted: Uint8Array<ArrayBuffer> | ArrayBuffer,
+  privateKey: string,
+): Promise<ArrayBuffer> => {
+  const key = await getCryptoKey(privateKey, "decrypt");
+  return window.crypto.subtle.decrypt(
+    {
+      name: "AES-GCM",
+      iv,
+    },
+    key,
+    encrypted,
+  );
+};

apps/client/src/features/editor/components/excalidraw/excalidraw-view.tsx

@@ -8,13 +8,14 @@ import {
   Text,
   useComputedColorScheme,
 } from "@mantine/core";
-import { useState } from "react";
+import { useState, useCallback } from "react";
 import { uploadFile } from "@/features/page/services/page-service.ts";
 import { svgStringToFile } from "@/lib";
 import { useDisclosure } from "@mantine/hooks";
 import { getFileUrl } from "@/lib/config.ts";
 import "@excalidraw/excalidraw/index.css";
 import type { ExcalidrawImperativeAPI } from "@excalidraw/excalidraw/types";
+import type { ExcalidrawElement } from "@excalidraw/element/types";
 import { IAttachment } from "@/features/attachments/types/attachment.types";
 import ReactClearModal from "react-clear-modal";
 import clsx from "clsx";
@@ -24,6 +25,7 @@ import { Suspense } from "react";
 import { useTranslation } from "react-i18next";
 import { useHandleLibrary } from "@excalidraw/excalidraw";
 import { localStorageLibraryAdapter } from "@/features/editor/components/excalidraw/excalidraw-utils.ts";
+import { useExcalidrawCollab } from "./use-excalidraw-collab";
 
 const Excalidraw = lazy(() =>
   import("@excalidraw/excalidraw").then((module) => ({
@@ -46,6 +48,16 @@ export default function ExcalidrawView(props: NodeViewProps) {
   const [opened, { open, close }] = useDisclosure(false);
   const computedColorScheme = useComputedColorScheme();
 
+  const pageId = editor.storage?.pageId;
+  const { broadcastScene } = useExcalidrawCollab(excalidrawAPI, pageId, opened);
+
+  const handleChange = useCallback(
+    (elements: readonly ExcalidrawElement[]) => {
+      broadcastScene(elements);
+    },
+    [broadcastScene],
+  );
+
   const handleOpen = async () => {
     if (!editor.isEditable) {
       return;
@@ -157,6 +169,7 @@ export default function ExcalidrawView(props: NodeViewProps) {
                 scrollToContent: true,
               }}
               theme={computedColorScheme}
+              onChange={handleChange}
             />
           </Suspense>
         </div>

apps/client/src/features/editor/components/excalidraw/portal.tsx

@@ -0,0 +1,257 @@
+import { CaptureUpdateAction } from "@excalidraw/excalidraw";
+import { trackEvent } from "@excalidraw/excalidraw/analytics";
+import { encryptData } from "@excalidraw/excalidraw/data/encryption";
+import { newElementWith } from "@excalidraw/element";
+import throttle from "lodash.throttle";
+
+import type { UserIdleState } from "@excalidraw/common";
+import type { OrderedExcalidrawElement } from "@excalidraw/element/types";
+import type {
+  OnUserFollowedPayload,
+  SocketId,
+} from "@excalidraw/excalidraw/types";
+
+import { WS_EVENTS, FILE_UPLOAD_TIMEOUT, WS_SUBTYPES } from "../app_constants";
+import { isSyncableElement } from "../data";
+
+import type {
+  SocketUpdateData,
+  SocketUpdateDataSource,
+  SyncableExcalidrawElement,
+} from "../data";
+import type { TCollabClass } from "./Collab";
+import type { Socket } from "socket.io-client";
+
+class Portal {
+  collab: TCollabClass;
+  socket: Socket | null = null;
+  socketInitialized: boolean = false; // we don't want the socket to emit any updates until it is fully initialized
+  roomId: string | null = null;
+  roomKey: string | null = null;
+  broadcastedElementVersions: Map<string, number> = new Map();
+
+  constructor(collab: TCollabClass) {
+    this.collab = collab;
+  }
+
+  open(socket: Socket, id: string, key: string) {
+    this.socket = socket;
+    this.roomId = id;
+    this.roomKey = key;
+
+    // Initialize socket listeners
+    this.socket.on("init-room", () => {
+      if (this.socket) {
+        this.socket.emit("join-room", this.roomId);
+        trackEvent("share", "room joined");
+      }
+    });
+    this.socket.on("new-user", async (_socketId: string) => {
+      this.broadcastScene(
+        WS_SUBTYPES.INIT,
+        this.collab.getSceneElementsIncludingDeleted(),
+        /* syncAll */ true,
+      );
+    });
+    this.socket.on("room-user-change", (clients: SocketId[]) => {
+      this.collab.setCollaborators(clients);
+    });
+
+    return socket;
+  }
+
+  close() {
+    if (!this.socket) {
+      return;
+    }
+    this.queueFileUpload.flush();
+    this.socket.close();
+    this.socket = null;
+    this.roomId = null;
+    this.roomKey = null;
+    this.socketInitialized = false;
+    this.broadcastedElementVersions = new Map();
+  }
+
+  isOpen() {
+    return !!(
+      this.socketInitialized &&
+      this.socket &&
+      this.roomId &&
+      this.roomKey
+    );
+  }
+
+  async _broadcastSocketData(
+    data: SocketUpdateData,
+    volatile: boolean = false,
+    roomId?: string,
+  ) {
+    if (this.isOpen()) {
+      const json = JSON.stringify(data);
+      const encoded = new TextEncoder().encode(json);
+      const { encryptedBuffer, iv } = await encryptData(this.roomKey!, encoded);
+
+      this.socket?.emit(
+        volatile ? WS_EVENTS.SERVER_VOLATILE : WS_EVENTS.SERVER,
+        roomId ?? this.roomId,
+        encryptedBuffer,
+        iv,
+      );
+    }
+  }
+
+  queueFileUpload = throttle(async () => {
+    try {
+      await this.collab.fileManager.saveFiles({
+        elements: this.collab.excalidrawAPI.getSceneElementsIncludingDeleted(),
+        files: this.collab.excalidrawAPI.getFiles(),
+      });
+    } catch (error: any) {
+      if (error.name !== "AbortError") {
+        this.collab.excalidrawAPI.updateScene({
+          appState: {
+            errorMessage: error.message,
+          },
+        });
+      }
+    }
+
+    let isChanged = false;
+    const newElements = this.collab.excalidrawAPI
+      .getSceneElementsIncludingDeleted()
+      .map((element) => {
+        if (this.collab.fileManager.shouldUpdateImageElementStatus(element)) {
+          isChanged = true;
+          // this will signal collaborators to pull image data from server
+          // (using mutation instead of newElementWith otherwise it'd break
+          // in-progress dragging)
+          return newElementWith(element, { status: "saved" });
+        }
+        return element;
+      });
+
+    if (isChanged) {
+      this.collab.excalidrawAPI.updateScene({
+        elements: newElements,
+        captureUpdate: CaptureUpdateAction.NEVER,
+      });
+    }
+  }, FILE_UPLOAD_TIMEOUT);
+
+  broadcastScene = async (
+    updateType: WS_SUBTYPES.INIT | WS_SUBTYPES.UPDATE,
+    elements: readonly OrderedExcalidrawElement[],
+    syncAll: boolean,
+  ) => {
+    if (updateType === WS_SUBTYPES.INIT && !syncAll) {
+      throw new Error("syncAll must be true when sending SCENE.INIT");
+    }
+
+    // sync out only the elements we think we need to to save bandwidth.
+    // periodically we'll resync the whole thing to make sure no one diverges
+    // due to a dropped message (server goes down etc).
+    const syncableElements = elements.reduce((acc, element) => {
+      if (
+        (syncAll ||
+          !this.broadcastedElementVersions.has(element.id) ||
+          element.version > this.broadcastedElementVersions.get(element.id)!) &&
+        isSyncableElement(element)
+      ) {
+        acc.push(element);
+      }
+      return acc;
+    }, [] as SyncableExcalidrawElement[]);
+
+    const data: SocketUpdateDataSource[typeof updateType] = {
+      type: updateType,
+      payload: {
+        elements: syncableElements,
+      },
+    };
+
+    for (const syncableElement of syncableElements) {
+      this.broadcastedElementVersions.set(
+        syncableElement.id,
+        syncableElement.version,
+      );
+    }
+
+    this.queueFileUpload();
+
+    await this._broadcastSocketData(data as SocketUpdateData);
+  };
+
+  broadcastIdleChange = (userState: UserIdleState) => {
+    if (this.socket?.id) {
+      const data: SocketUpdateDataSource["IDLE_STATUS"] = {
+        type: WS_SUBTYPES.IDLE_STATUS,
+        payload: {
+          socketId: this.socket.id as SocketId,
+          userState,
+          username: this.collab.state.username,
+        },
+      };
+      return this._broadcastSocketData(
+        data as SocketUpdateData,
+        true, // volatile
+      );
+    }
+  };
+
+  broadcastMouseLocation = (payload: {
+    pointer: SocketUpdateDataSource["MOUSE_LOCATION"]["payload"]["pointer"];
+    button: SocketUpdateDataSource["MOUSE_LOCATION"]["payload"]["button"];
+  }) => {
+    if (this.socket?.id) {
+      const data: SocketUpdateDataSource["MOUSE_LOCATION"] = {
+        type: WS_SUBTYPES.MOUSE_LOCATION,
+        payload: {
+          socketId: this.socket.id as SocketId,
+          pointer: payload.pointer,
+          button: payload.button || "up",
+          selectedElementIds:
+          this.collab.excalidrawAPI.getAppState().selectedElementIds,
+          username: this.collab.state.username,
+        },
+      };
+
+      return this._broadcastSocketData(
+        data as SocketUpdateData,
+        true, // volatile
+      );
+    }
+  };
+
+  broadcastVisibleSceneBounds = (
+    payload: {
+      sceneBounds: SocketUpdateDataSource["USER_VISIBLE_SCENE_BOUNDS"]["payload"]["sceneBounds"];
+    },
+    roomId: string,
+  ) => {
+    if (this.socket?.id) {
+      const data: SocketUpdateDataSource["USER_VISIBLE_SCENE_BOUNDS"] = {
+        type: WS_SUBTYPES.USER_VISIBLE_SCENE_BOUNDS,
+        payload: {
+          socketId: this.socket.id as SocketId,
+          username: this.collab.state.username,
+          sceneBounds: payload.sceneBounds,
+        },
+      };
+
+      return this._broadcastSocketData(
+        data as SocketUpdateData,
+        true, // volatile
+        roomId,
+      );
+    }
+  };
+
+  broadcastUserFollowed = (payload: OnUserFollowedPayload) => {
+    if (this.socket?.id) {
+      this.socket.emit(WS_EVENTS.USER_FOLLOW_CHANGE, payload);
+    }
+  };
+}
+
+export default Portal;

apps/client/src/features/editor/components/excalidraw/use-excalidraw-collab.ts

@@ -0,0 +1,145 @@
+import { useEffect, useRef, useCallback, useMemo } from "react";
+import { useAtom } from "jotai";
+import { socketAtom } from "@/features/websocket/atoms/socket-atom";
+import type { ExcalidrawImperativeAPI } from "@excalidraw/excalidraw/types";
+import type { ExcalidrawElement } from "@excalidraw/element/types";
+import { reconcileElements, getSceneVersion } from "@excalidraw/excalidraw";
+import throttle from "lodash.throttle";
+
+type Collaborator = {
+  socketId: string;
+  isCurrentUser?: boolean;
+};
+
+export function useExcalidrawCollab(
+  excalidrawAPI: ExcalidrawImperativeAPI | null,
+  pageId: string | undefined,
+  isOpen: boolean,
+) {
+  const [socket] = useAtom(socketAtom);
+  const lastBroadcastedVersion = useRef(-1);
+  const isInitialized = useRef(false);
+
+  const roomId = pageId ? `excalidraw-${pageId}` : null;
+
+  // Create stable throttled broadcast function
+  const broadcastScene = useMemo(
+    () =>
+      throttle((elements: readonly ExcalidrawElement[]) => {
+        if (!socket || !roomId || !isInitialized.current) {
+          console.log("broadcastScene: not ready", {
+            socket: !!socket,
+            roomId,
+            isInitialized: isInitialized.current,
+          });
+          return;
+        }
+
+        // getSceneVersion sums all element versions - increases on ANY change
+        const sceneVersion = getSceneVersion(elements);
+
+        if (sceneVersion <= lastBroadcastedVersion.current) {
+          return;
+        }
+
+        const data = {
+          type: "SCENE_UPDATE",
+          payload: { elements },
+        };
+
+        // Send as plain JSON for now (no encryption)
+        const json = JSON.stringify(data);
+        console.log("Broadcasting scene, version:", sceneVersion);
+
+        socket.emit("server-broadcast", [roomId, json, null]);
+        lastBroadcastedVersion.current = sceneVersion;
+      }, 100),
+    [socket, roomId],
+  );
+
+  // Handle incoming broadcasts
+  const handleClientBroadcast = useCallback(
+    (jsonData: string, _iv: Uint8Array | null) => {
+      if (!excalidrawAPI) return;
+
+      try {
+        const data = JSON.parse(jsonData);
+
+        if (data.type === "SCENE_UPDATE" && data.payload?.elements) {
+          const remoteElements = data.payload.elements;
+          const localElements =
+            excalidrawAPI.getSceneElementsIncludingDeleted();
+
+          const reconciledElements = reconcileElements(
+            localElements,
+            remoteElements,
+            excalidrawAPI.getAppState(),
+          );
+
+          excalidrawAPI.updateScene({
+            elements: reconciledElements,
+          });
+
+          // Update version to prevent echo
+          lastBroadcastedVersion.current = getSceneVersion(reconciledElements);
+        }
+      } catch (err) {
+        console.error("Failed to process broadcast:", err);
+      }
+    },
+    [excalidrawAPI],
+  );
+
+  // Handle room user changes
+  const handleRoomUserChange = useCallback(
+    (socketIds: string[]) => {
+      if (!excalidrawAPI || !socket) return;
+
+      const collaborators = new Map<string, Collaborator>();
+      for (const id of socketIds) {
+        collaborators.set(id, {
+          socketId: id,
+          isCurrentUser: id === socket.id,
+        });
+      }
+      // @ts-ignore
+      excalidrawAPI.updateScene({ collaborators });
+    },
+    [excalidrawAPI, socket],
+  );
+
+  // Join/leave room based on modal state
+  useEffect(() => {
+    if (!socket || !roomId || !isOpen) return;
+
+    console.log("Joining room:", roomId);
+    socket.emit("join-room", roomId);
+    isInitialized.current = true;
+
+    // Set up listeners
+    socket.on("client-broadcast", handleClientBroadcast);
+    socket.on("room-user-change", handleRoomUserChange);
+    socket.on("first-in-room", () => {
+      console.log("First in excalidraw room");
+    });
+    socket.on("new-user", (socketId: string) => {
+      console.log("New user joined:", socketId);
+      if (excalidrawAPI) {
+        broadcastScene(excalidrawAPI.getSceneElements());
+      }
+    });
+
+    return () => {
+      console.log("Leaving room:", roomId);
+      socket.emit("leave-room", roomId);
+      socket.off("client-broadcast", handleClientBroadcast);
+      socket.off("room-user-change", handleRoomUserChange);
+      socket.off("first-in-room");
+      socket.off("new-user");
+      isInitialized.current = false;
+      lastBroadcastedVersion.current = -1;
+    };
+  }, [socket, roomId, isOpen, handleClientBroadcast, handleRoomUserChange, broadcastScene, excalidrawAPI]);
+
+  return { broadcastScene };
+}

apps/server/src/ws/ws.gateway.ts

@@ -1,4 +1,5 @@
 import {
+  ConnectedSocket,
   MessageBody,
   OnGatewayConnection,
   OnGatewayDisconnect,
@@ -76,7 +77,7 @@ export class WsGateway
 
   @SubscribeMessage('join-room')
   async handleJoinRoom(
-    client: Socket,
+    @ConnectedSocket() client: Socket,
     @MessageBody() roomId: string,
   ): Promise<void> {
     await this.excalidrawCollabService.handleJoinRoom(
@@ -87,14 +88,17 @@ export class WsGateway
   }
 
   @SubscribeMessage('leave-room')
-  handleLeaveRoom(client: Socket, @MessageBody() roomName: string): void {
+  handleLeaveRoom(
+    @ConnectedSocket() client: Socket,
+    @MessageBody() roomName: string,
+  ): void {
     client.leave(roomName);
   }
 
   @SubscribeMessage('server-broadcast')
   handleServerBroadcast(
-    client: Socket,
+    @ConnectedSocket() client: Socket,
-    [roomId, encryptedData, iv]: [string, ArrayBuffer, Uint8Array],
+    @MessageBody() [roomId, encryptedData, iv]: [string, ArrayBuffer, Uint8Array],
   ): void {
     this.excalidrawCollabService.handleServerBroadcast(
       client,
@@ -106,8 +110,8 @@ export class WsGateway
 
   @SubscribeMessage('server-volatile-broadcast')
   handleServerVolatileBroadcast(
-    client: Socket,
+    @ConnectedSocket() client: Socket,
-    [roomId, encryptedData, iv]: [string, ArrayBuffer, Uint8Array],
+    @MessageBody() [roomId, encryptedData, iv]: [string, ArrayBuffer, Uint8Array],
   ): void {
     this.excalidrawCollabService.handleServerVolatileBroadcast(
       client,
@@ -119,7 +123,7 @@ export class WsGateway
 
   @SubscribeMessage('user-follow')
   async handleUserFollow(
-    client: Socket,
+    @ConnectedSocket() client: Socket,
     @MessageBody() payload: ExcalidrawFollowPayload,
   ): Promise<void> {
     await this.excalidrawCollabService.handleUserFollow(