From f650d2591b4624138806d6a48e9045a441eac38d Mon Sep 17 00:00:00 2001 From: vvzvlad Date: Sat, 20 Jun 2026 19:48:06 +0300 Subject: [PATCH] fix(tree): address realtime-tree-server review findings - make addTreeNode receivers idempotent (invalidateOnCreatePage guard + buildTree dedup) so the author's self-echo no longer duplicates the node - broadcast realtime tree updates for bulk copy/duplicate and import via a root refetch: PAGE_CREATED now carries spaceId and the WS listener falls back to refetchRootTreeNodeEvent when no per-node snapshot is present - remove the now-dead client-relay inbound path (isTreeEvent/handleTreeEvent) that remained a stale-restriction-cache attack surface - honest string|null cast for a root move's parent id - add tests: buildTree dedup; onPageCreated per-node vs refetch branching Co-Authored-By: Claude Opus 4.8 --- .../src/features/page/queries/page-query.ts | 10 ++ .../features/page/tree/utils/utils.test.ts | 40 ++++++++ .../src/features/page/tree/utils/utils.ts | 7 ++ .../src/features/websocket/use-tree-socket.ts | 4 +- .../src/core/page/services/page.service.ts | 5 + .../services/file-import-task.service.ts | 4 + .../src/ws/listeners/page-ws.listener.spec.ts | 95 +++++++++++++++++++ .../src/ws/listeners/page-ws.listener.ts | 18 +++- apps/server/src/ws/ws.gateway.ts | 8 +- apps/server/src/ws/ws.service.ts | 62 ++---------- apps/server/src/ws/ws.utils.ts | 8 -- 11 files changed, 190 insertions(+), 71 deletions(-) create mode 100644 apps/client/src/features/page/tree/utils/utils.test.ts create mode 100644 apps/server/src/ws/listeners/page-ws.listener.spec.ts diff --git a/apps/client/src/features/page/queries/page-query.ts b/apps/client/src/features/page/queries/page-query.ts index 11ba7f32..c631f892 100644 --- a/apps/client/src/features/page/queries/page-query.ts +++ b/apps/client/src/features/page/queries/page-query.ts @@ -360,6 +360,16 @@ export function invalidateOnCreatePage(data: Partial) { queryKey, (old) => { if (!old) return old; + + // Idempotency guard: the server now self-echoes addTreeNode back to the + // author, so this writer can run twice for one create (mutation onSuccess + // + socket echo). Skip the append if the page is already in the cache to + // avoid a duplicate node / duplicate React key. + const exists = old.pages.some((page) => + page.items.some((item) => item.id === newPage.id), + ); + if (exists) return old; + return { ...old, pages: old.pages.map((page, index) => { diff --git a/apps/client/src/features/page/tree/utils/utils.test.ts b/apps/client/src/features/page/tree/utils/utils.test.ts new file mode 100644 index 00000000..7de43da1 --- /dev/null +++ b/apps/client/src/features/page/tree/utils/utils.test.ts @@ -0,0 +1,40 @@ +import { describe, it, expect } from "vitest"; +import { buildTree } from "./utils"; +import type { IPage } from "@/features/page/types/page.types.ts"; + +function page(id: string, position: string): IPage { + return { + id, + slugId: `slug-${id}`, + title: id.toUpperCase(), + icon: "", + position, + hasChildren: false, + spaceId: "space-1", + parentPageId: null as unknown as string, + } as IPage; +} + +describe("buildTree", () => { + it("builds one node per unique page", () => { + const tree = buildTree([page("a", "a1"), page("b", "a2")]); + expect(tree.map((n) => n.id)).toEqual(["a", "b"]); + }); + + it("dedups a duplicate id so the tree has no duplicate node", () => { + // A realtime cache write could append a page twice; buildTree must not emit + // two references to the same node (which would crash the sidebar render with + // a duplicate React key). + const tree = buildTree([ + page("a", "a1"), + page("b", "a2"), + page("a", "a1"), // duplicate id + ]); + + expect(tree).toHaveLength(2); + expect(tree.map((n) => n.id).sort()).toEqual(["a", "b"]); + // No id appears more than once. + const ids = tree.map((n) => n.id); + expect(new Set(ids).size).toBe(ids.length); + }); +}); diff --git a/apps/client/src/features/page/tree/utils/utils.ts b/apps/client/src/features/page/tree/utils/utils.ts index 0c42f9b9..47246108 100644 --- a/apps/client/src/features/page/tree/utils/utils.ts +++ b/apps/client/src/features/page/tree/utils/utils.ts @@ -29,7 +29,14 @@ export function buildTree(pages: IPage[]): SpaceTreeNode[] { }; }); + // Defense-in-depth: a duplicate id in `pages` would push two references to the + // same node, producing a duplicate React key that crashes the sidebar render. + // Track ids we've already pushed and skip repeats so a stray duplicate from a + // realtime cache write can never break the tree. + const seen = new Set(); pages.forEach((page) => { + if (seen.has(page.id)) return; + seen.add(page.id); tree.push(pageMap[page.id]); }); diff --git a/apps/client/src/features/websocket/use-tree-socket.ts b/apps/client/src/features/websocket/use-tree-socket.ts index 89506a16..317392e7 100644 --- a/apps/client/src/features/websocket/use-tree-socket.ts +++ b/apps/client/src/features/websocket/use-tree-socket.ts @@ -115,7 +115,9 @@ export const useTreeSocket = () => { | undefined; const patch: Partial = { position: event.payload.position, - parentPageId: newParentId as string, + // Honest type: a root move has a null parent, so this is + // `string | null`, not always `string`. + parentPageId: newParentId as string | null, }; if (pageData) { // The tree node stores the title as `name`. diff --git a/apps/server/src/core/page/services/page.service.ts b/apps/server/src/core/page/services/page.service.ts index 672ff3fb..615433fd 100644 --- a/apps/server/src/core/page/services/page.service.ts +++ b/apps/server/src/core/page/services/page.service.ts @@ -758,9 +758,14 @@ export class PageService { } const insertedPageIds = insertablePages.map((page) => page.id); + // `spaceId` is the single destination space for the whole copy/duplicate + // (every inserted page above gets `spaceId: spaceId`). It lets the WS + // listener trigger a root refetch for the bulk subtree (no `pages` snapshot + // here on purpose — we want the refetch fallback, not per-node addTreeNode). this.eventEmitter.emit(EventName.PAGE_CREATED, { pageIds: insertedPageIds, workspaceId: authUser.workspaceId, + spaceId, }); //TODO: best to handle this in a queue diff --git a/apps/server/src/integrations/import/services/file-import-task.service.ts b/apps/server/src/integrations/import/services/file-import-task.service.ts index 40525ddf..218c75ca 100644 --- a/apps/server/src/integrations/import/services/file-import-task.service.ts +++ b/apps/server/src/integrations/import/services/file-import-task.service.ts @@ -552,9 +552,13 @@ export class FileImportTaskService { } if (validPageIds.size > 0) { + // Carry the destination spaceId so the WS listener can trigger a root + // refetch for the imported subtree (no `pages` snapshot -> refetch + // fallback rather than per-node addTreeNode). this.eventEmitter.emit(EventName.PAGE_CREATED, { pageIds: Array.from(validPageIds), workspaceId: fileTask.workspaceId, + spaceId: fileTask.spaceId, }); } diff --git a/apps/server/src/ws/listeners/page-ws.listener.spec.ts b/apps/server/src/ws/listeners/page-ws.listener.spec.ts new file mode 100644 index 00000000..734e8228 --- /dev/null +++ b/apps/server/src/ws/listeners/page-ws.listener.spec.ts @@ -0,0 +1,95 @@ +import { Test, TestingModule } from '@nestjs/testing'; +import { PageWsListener } from './page-ws.listener'; +import { WsTreeService } from '../ws-tree.service'; +import { + PageEvent, + TreeNodeSnapshot, +} from '../../database/listeners/page.listener'; + +const snapshot: TreeNodeSnapshot = { + id: 'page-1', + slugId: 'slug-1', + title: 'Hello', + icon: '📄', + position: 'a1', + spaceId: 'space-1', + parentPageId: null, +}; + +describe('PageWsListener.onPageCreated', () => { + let listener: PageWsListener; + let wsTree: { + broadcastPageCreated: jest.Mock; + broadcastRefetchRoot: jest.Mock; + }; + + beforeEach(async () => { + wsTree = { + broadcastPageCreated: jest.fn().mockResolvedValue(undefined), + broadcastRefetchRoot: jest.fn().mockResolvedValue(undefined), + }; + + const module: TestingModule = await Test.createTestingModule({ + providers: [ + PageWsListener, + { provide: WsTreeService, useValue: wsTree }, + ], + }).compile(); + + listener = module.get(PageWsListener); + }); + + it('with `pages`: broadcasts a per-node addTreeNode and does NOT refetch root', async () => { + const event: PageEvent = { + pageIds: ['page-1'], + workspaceId: 'ws-1', + pages: [snapshot], + }; + + await listener.onPageCreated(event); + + expect(wsTree.broadcastPageCreated).toHaveBeenCalledTimes(1); + expect(wsTree.broadcastPageCreated).toHaveBeenCalledWith(snapshot); + expect(wsTree.broadcastRefetchRoot).not.toHaveBeenCalled(); + }); + + it('without `pages` but WITH `spaceId` (bulk create): falls back to a root refetch', async () => { + const event: PageEvent = { + pageIds: ['page-1', 'page-2'], + workspaceId: 'ws-1', + spaceId: 'space-9', + }; + + await listener.onPageCreated(event); + + expect(wsTree.broadcastPageCreated).not.toHaveBeenCalled(); + expect(wsTree.broadcastRefetchRoot).toHaveBeenCalledTimes(1); + expect(wsTree.broadcastRefetchRoot).toHaveBeenCalledWith('space-9'); + }); + + it('with an EMPTY `pages` array but WITH `spaceId`: still falls back to a root refetch', async () => { + const event: PageEvent = { + pageIds: ['page-1'], + workspaceId: 'ws-1', + pages: [], + spaceId: 'space-9', + }; + + await listener.onPageCreated(event); + + expect(wsTree.broadcastPageCreated).not.toHaveBeenCalled(); + expect(wsTree.broadcastRefetchRoot).toHaveBeenCalledWith('space-9'); + }); + + it('without `pages` and without `spaceId`: does nothing (no broadcast)', async () => { + const event: PageEvent = { + pageIds: ['page-1'], + workspaceId: 'ws-1', + }; + + await listener.onPageCreated(event); + + expect(wsTree.broadcastPageCreated).not.toHaveBeenCalled(); + expect(wsTree.broadcastRefetchRoot).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/server/src/ws/listeners/page-ws.listener.ts b/apps/server/src/ws/listeners/page-ws.listener.ts index c2afc2c6..100d0f85 100644 --- a/apps/server/src/ws/listeners/page-ws.listener.ts +++ b/apps/server/src/ws/listeners/page-ws.listener.ts @@ -33,8 +33,22 @@ export class PageWsListener { @OnEvent(EventName.PAGE_CREATED) async onPageCreated(event: PageEvent): Promise { - for (const page of event.pages ?? []) { - await this.wsTree.broadcastPageCreated(page); + // Two creation shapes: + // - Single-page create carries precise node snapshots (`pages`), so we + // broadcast a pointwise addTreeNode per node. + // - Bulk create (copy/duplicate, import) produces whole subtrees and omits + // `pages`; per-node placement would be fragile, so we fall back to a root + // refetch (carries no page data, clients re-fetch via the permission- + // checked API). Same mechanism PAGE_RESTORED uses. + if (event.pages?.length) { + for (const page of event.pages) { + await this.wsTree.broadcastPageCreated(page); + } + return; + } + + if (event.spaceId) { + await this.wsTree.broadcastRefetchRoot(event.spaceId); } } diff --git a/apps/server/src/ws/ws.gateway.ts b/apps/server/src/ws/ws.gateway.ts index fd6810c8..a4f66257 100644 --- a/apps/server/src/ws/ws.gateway.ts +++ b/apps/server/src/ws/ws.gateway.ts @@ -62,10 +62,10 @@ export class WsGateway } @SubscribeMessage('message') - async handleMessage(client: Socket, data: any): Promise { - if (this.wsService.isTreeEvent(data)) { - await this.wsService.handleTreeEvent(client, data); - } + handleMessage(_client: Socket, _data: any): void { + // Inbound tree events from clients are no longer accepted: tree updates are + // now server-authoritative (broadcast by PageWsListener from domain events). + // The old client-relay path was removed to close that attack surface. } /* diff --git a/apps/server/src/ws/ws.service.ts b/apps/server/src/ws/ws.service.ts index 3a8c8bc3..75cf598e 100644 --- a/apps/server/src/ws/ws.service.ts +++ b/apps/server/src/ws/ws.service.ts @@ -1,10 +1,9 @@ import { Inject, Injectable } from '@nestjs/common'; import { CACHE_MANAGER } from '@nestjs/cache-manager'; import { Cache } from 'cache-manager'; -import { Server, Socket } from 'socket.io'; +import { Server } from 'socket.io'; import { PagePermissionRepo } from '@docmost/db/repos/page/page-permission.repo'; import { - TREE_EVENTS, WS_SPACE_RESTRICTION_CACHE_PREFIX, WS_CACHE_TTL_MS, getSpaceRoomName, @@ -23,39 +22,6 @@ export class WsService { this.server = server; } - async handleTreeEvent(client: Socket, data: any): Promise { - const room = getSpaceRoomName(data.spaceId); - - if (!client.rooms.has(room)) { - return; - } - - if (data.operation === 'refetchRootTreeNodeEvent') { - client.broadcast.to(room).emit('message', data); - return; - } - - const hasRestrictions = await this.spaceHasRestrictions(data.spaceId); - if (!hasRestrictions) { - client.broadcast.to(room).emit('message', data); - return; - } - - const pageId = this.extractPageId(data); - if (!pageId) { - return; - } - - const isRestricted = - await this.pagePermissionRepo.hasRestrictedAncestor(pageId); - if (!isRestricted) { - client.broadcast.to(room).emit('message', data); - return; - } - - await this.broadcastToAuthorizedUsers(room, client.id, pageId, data); - } - // Drop the cached spaceHasRestrictions verdict for a space. spaceHasRestrictions // caches "does this space have ANY restricted page" for WS_CACHE_TTL_MS (30s), // and emitTreeEvent / emitCommentEvent take a room-wide fast path when it is @@ -72,6 +38,9 @@ export class WsService { // have ZERO callers in apps/server/src; PageAccessService only validates access. // This primitive is kept (and tested) so that flow, when it lands, has the // correct hook to invalidate the cache. + // + // TODO: the future restriction-mutation endpoint (restrict/grant/revoke page + // access) MUST call this with the affected page's spaceId. async invalidateSpaceRestrictionCache(spaceId: string): Promise { await this.cacheManager.del( `${WS_SPACE_RESTRICTION_CACHE_PREFIX}${spaceId}`, @@ -136,8 +105,8 @@ export class WsService { // Unconditional broadcast to everyone in the space room. Used for space-wide // signals that carry no page payload (e.g. refetchRootTreeNodeEvent on // restore): there is no per-page data to leak, and each client refetches the - // root tree through its own authorized query. Mirrors handleTreeEvent's - // special-casing of refetchRootTreeNodeEvent (no restriction check). + // root tree through its own authorized query (refetchRootTreeNodeEvent carries + // no per-page data, so no restriction check is needed). emitToSpaceRoom(spaceId: string, data: any): void { this.server.to(getSpaceRoomName(spaceId)).emit('message', data); } @@ -198,10 +167,6 @@ export class WsService { await this.broadcastToAuthorizedUsers(room, null, pageId, data); } - isTreeEvent(data: any): boolean { - return TREE_EVENTS.has(data?.operation) && !!data?.spaceId; - } - private async broadcastToAuthorizedUsers( room: string, excludeSocketId: string | null, @@ -264,19 +229,4 @@ export class WsService { return hasRestrictions; } - - private extractPageId(data: any): string | null { - switch (data.operation) { - case 'addTreeNode': - return data.payload?.data?.id ?? null; - case 'moveTreeNode': - return data.payload?.id ?? null; - case 'deleteTreeNode': - return data.payload?.node?.id ?? null; - case 'updateOne': - return data.id ?? null; - default: - return null; - } - } } diff --git a/apps/server/src/ws/ws.utils.ts b/apps/server/src/ws/ws.utils.ts index c0d48b54..a1eb1569 100644 --- a/apps/server/src/ws/ws.utils.ts +++ b/apps/server/src/ws/ws.utils.ts @@ -8,11 +8,3 @@ export function getSpaceRoomName(spaceId: string): string { export function getUserRoomName(userId: string): string { return `user-${userId}`; } - -export const TREE_EVENTS = new Set([ - 'updateOne', - 'addTreeNode', - 'moveTreeNode', - 'deleteTreeNode', - 'refetchRootTreeNodeEvent', -]);