Sidebar tree socket permissions

apps/server/src/ws/ws-tree.service.ts

@@ -0,0 +1,47 @@
+import { Injectable } from '@nestjs/common';
+import { Page } from '@docmost/db/types/entity.types';
+import { WsService } from './ws.service';
+
+@Injectable()
+export class WsTreeService {
+  constructor(private readonly wsService: WsService) {}
+
+  async notifyPageRestricted(page: Page, excludeUserId: string): Promise<void> {
+    await this.wsService.emitToSpaceExceptUsers(page.spaceId, [excludeUserId], {
+      operation: 'deleteTreeNode',
+      spaceId: page.spaceId,
+      payload: {
+        node: {
+          id: page.id,
+          slugId: page.slugId,
+        },
+      },
+    });
+  }
+
+  async notifyPermissionGranted(page: Page, userIds: string[]): Promise<void> {
+    if (userIds.length === 0) return;
+
+    await this.wsService.emitToUsers(userIds, {
+      operation: 'addTreeNode',
+      spaceId: page.spaceId,
+      payload: {
+        parentId: page.parentPageId ?? null,
+        index: 0,
+        data: {
+          id: page.id,
+          slugId: page.slugId,
+          name: page.title ?? '',
+          title: page.title,
+          icon: page.icon,
+          position: page.position,
+          spaceId: page.spaceId,
+          parentPageId: page.parentPageId,
+          creatorId: page.creatorId,
+          hasChildren: false,
+          children: [],
+        },
+      },
+    });
+  }
+}

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

@@ -1,6 +1,7 @@
 import {
   MessageBody,
   OnGatewayConnection,
+  OnGatewayInit,
   SubscribeMessage,
   WebSocketGateway,
   WebSocketServer,
@@ -10,20 +11,30 @@ import { TokenService } from '../core/auth/services/token.service';
 import { JwtPayload, JwtType } from '../core/auth/dto/jwt-payload';
 import { OnModuleDestroy } from '@nestjs/common';
 import { SpaceMemberRepo } from '@docmost/db/repos/space/space-member.repo';
+import { WsService } from './ws.service';
+import { getSpaceRoomName, getUserRoomName } from './ws.utils';
 import * as cookie from 'cookie';
 
 @WebSocketGateway({
   cors: { origin: '*' },
   transports: ['websocket'],
 })
-export class WsGateway implements OnGatewayConnection, OnModuleDestroy {
+export class WsGateway
+  implements OnGatewayConnection, OnGatewayInit, OnModuleDestroy
+{
   @WebSocketServer()
   server: Server;
+
   constructor(
     private tokenService: TokenService,
     private spaceMemberRepo: SpaceMemberRepo,
+    private wsService: WsService,
   ) {}
 
+  afterInit(server: Server): void {
+    this.wsService.setServer(server);
+  }
+
   async handleConnection(client: Socket, ...args: any[]): Promise<void> {
     try {
       const cookies = cookie.parse(client.handshake.headers.cookie);
@@ -35,12 +46,15 @@ export class WsGateway implements OnGatewayConnection, OnModuleDestroy {
       const userId = token.sub;
       const workspaceId = token.workspaceId;
 
+      client.data.userId = userId;
+
       const userSpaceIds = await this.spaceMemberRepo.getUserSpaceIds(userId);
 
+      const userRoom = getUserRoomName(userId);
       const workspaceRoom = `workspace-${workspaceId}`;
-      const spaceRooms = userSpaceIds.map((id) => this.getSpaceRoomName(id));
+      const spaceRooms = userSpaceIds.map((id) => getSpaceRoomName(id));
 
-      client.join([workspaceRoom, ...spaceRooms]);
+      client.join([userRoom, workspaceRoom, ...spaceRooms]);
     } catch (err) {
       client.emit('Unauthorized');
       client.disconnect();
@@ -48,17 +62,9 @@ export class WsGateway implements OnGatewayConnection, OnModuleDestroy {
   }
 
   @SubscribeMessage('message')
-  handleMessage(client: Socket, data: any): void {
-    const spaceEvents = [
-      'updateOne',
-      'addTreeNode',
-      'moveTreeNode',
-      'deleteTreeNode',
-    ];
-
-    if (spaceEvents.includes(data?.operation) && data?.spaceId) {
-      const room = this.getSpaceRoomName(data.spaceId);
-      client.broadcast.to(room).emit('message', data);
+  async handleMessage(client: Socket, data: any): Promise<void> {
+    if (this.wsService.isTreeEvent(data)) {
+      await this.wsService.handleTreeEvent(client, data);
       return;
     }
 
@@ -81,8 +87,4 @@ export class WsGateway implements OnGatewayConnection, OnModuleDestroy {
       this.server.close();
     }
   }
-
-  getSpaceRoomName(spaceId: string): string {
-    return `space-${spaceId}`;
-  }
 }

apps/server/src/ws/ws.module.ts

@@ -1,9 +1,13 @@
-import { Module } from '@nestjs/common';
+import { Global, Module } from '@nestjs/common';
 import { WsGateway } from './ws.gateway';
+import { WsService } from './ws.service';
+import { WsTreeService } from './ws-tree.service';
 import { TokenModule } from '../core/auth/token.module';
 
+@Global()
 @Module({
   imports: [TokenModule],
-  providers: [WsGateway],
+  providers: [WsGateway, WsService, WsTreeService],
+  exports: [WsService, WsTreeService],
 })
 export class WsModule {}

apps/server/src/ws/ws.service.ts

@@ -0,0 +1,157 @@
+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 { PagePermissionRepo } from '@docmost/db/repos/page/page-permission.repo';
+import {
+  TREE_EVENTS,
+  WS_SPACE_RESTRICTION_CACHE_PREFIX,
+  WS_CACHE_TTL_MS,
+  getSpaceRoomName,
+  getUserRoomName,
+} from './ws.utils';
+
+@Injectable()
+export class WsService {
+  private server: Server;
+
+  constructor(
+    private readonly pagePermissionRepo: PagePermissionRepo,
+    @Inject(CACHE_MANAGER) private readonly cacheManager: Cache,
+  ) {}
+
+  setServer(server: Server): void {
+    this.server = server;
+  }
+
+  async handleTreeEvent(client: Socket, data: any): Promise<void> {
+    const room = getSpaceRoomName(data.spaceId);
+
+    const hasRestrictions = await this.spaceHasRestrictions(data.spaceId);
+    if (!hasRestrictions) {
+      client.broadcast.to(room).emit('message', data);
+      return;
+    }
+
+    const pageId = this.extractPageId(data);
+    if (!pageId) {
+      client.broadcast.to(room).emit('message', data);
+      return;
+    }
+
+    const isRestricted =
+      await this.pagePermissionRepo.hasRestrictedAncestor(pageId);
+    if (!isRestricted) {
+      client.broadcast.to(room).emit('message', data);
+      return;
+    }
+
+    await this.broadcastToAuthorizedUsers(client, room, pageId, data);
+  }
+
+  async invalidateSpaceRestrictionCache(spaceId: string): Promise<void> {
+    await this.cacheManager.del(
+      `${WS_SPACE_RESTRICTION_CACHE_PREFIX}${spaceId}`,
+    );
+  }
+
+  async emitToUsers(userIds: string[], data: any): Promise<void> {
+    if (userIds.length === 0) return;
+    const rooms = userIds.map((id) => getUserRoomName(id));
+    this.server.to(rooms).emit('message', data);
+  }
+
+  async emitToSpaceExceptUsers(
+    spaceId: string,
+    excludeUserIds: string[],
+    data: any,
+  ): Promise<void> {
+    const room = getSpaceRoomName(spaceId);
+    const sockets = await this.server.in(room).fetchSockets();
+    const excludeSet = new Set(excludeUserIds);
+
+    for (const socket of sockets) {
+      const userId = socket.data.userId as string;
+      if (userId && !excludeSet.has(userId)) {
+        socket.emit('message', data);
+      }
+    }
+  }
+
+  isTreeEvent(data: any): boolean {
+    return TREE_EVENTS.has(data?.operation) && !!data?.spaceId;
+  }
+
+  private async broadcastToAuthorizedUsers(
+    sender: Socket,
+    room: string,
+    pageId: string,
+    data: any,
+  ): Promise<void> {
+    const sockets = await this.server.in(room).fetchSockets();
+
+    const otherSockets = sockets.filter((s) => s.id !== sender.id);
+    if (otherSockets.length === 0) return;
+
+    const userSocketMap = new Map<string, typeof otherSockets>();
+    for (const socket of otherSockets) {
+      const userId = socket.data.userId as string;
+      if (!userId) continue;
+      const existing = userSocketMap.get(userId);
+      if (existing) {
+        existing.push(socket);
+      } else {
+        userSocketMap.set(userId, [socket]);
+      }
+    }
+
+    const candidateUserIds = Array.from(userSocketMap.keys());
+    if (candidateUserIds.length === 0) return;
+
+    const authorizedUserIds =
+      await this.pagePermissionRepo.getUserIdsWithPageAccess(
+        pageId,
+        candidateUserIds,
+      );
+
+    const authorizedSet = new Set(authorizedUserIds);
+    for (const [userId, userSockets] of userSocketMap) {
+      if (authorizedSet.has(userId)) {
+        for (const socket of userSockets) {
+          socket.emit('message', data);
+        }
+      }
+    }
+  }
+
+  private async spaceHasRestrictions(spaceId: string): Promise<boolean> {
+    const cacheKey = `${WS_SPACE_RESTRICTION_CACHE_PREFIX}${spaceId}`;
+
+    const cached = await this.cacheManager.get<boolean>(cacheKey);
+    if (cached !== undefined && cached !== null) {
+      return cached;
+    }
+
+    const hasRestrictions =
+      await this.pagePermissionRepo.hasRestrictedPagesInSpace(spaceId);
+
+    await this.cacheManager.set(cacheKey, hasRestrictions, WS_CACHE_TTL_MS);
+
+    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;
+    }
+  }
+}

apps/server/src/ws/ws.utils.ts

@@ -0,0 +1,17 @@
+export const WS_CACHE_TTL_MS = 30_000;
+export const WS_SPACE_RESTRICTION_CACHE_PREFIX = 'ws:space-restrictions:';
+
+export function getSpaceRoomName(spaceId: string): string {
+  return `space-${spaceId}`;
+}
+
+export function getUserRoomName(userId: string): string {
+  return `user-${userId}`;
+}
+
+export const TREE_EVENTS = new Set([
+  'updateOne',
+  'addTreeNode',
+  'moveTreeNode',
+  'deleteTreeNode',
+]);