Move socket management from client -> server (#13120)

* Move socket management from client -> server

* feedback

* fix: Reconcile collection rooms on visibility change

Handle collections.update transitions between private and team-visible by
joining or rebuilding the collection channel, route collection events through
the members channel so guests are excluded, and use removeData for the
collection ids cache invalidation.

---------

Co-authored-by: Claude <noreply@anthropic.com>

app/components/WebsocketProvider.tsx

@@ -120,16 +120,6 @@ function useConnectionHandlers() {
         data,
       });
     });
-
-    // received a message from the API server that we should request
-    // to join or leave a specific room. Forward that to the ws server.
-    socket.on("join", (event) => {
-      socket.emit("join", event);
-    });
-
-    socket.on("leave", (event) => {
-      socket.emit("leave", event);
-    });
   };
 }
 
@@ -641,19 +631,24 @@ function useUserHandlers() {
       users.add(event);
     });
 
-    socket.on("users.demote", async (event: PartialExcept<User, "id">) => {
+    // the current user's role changed, so their policies are invalid and the
+    // set of accessible collections may have changed.
+    const handleRoleChange = async (event: PartialExcept<User, "id">) => {
       if (event.id === auth.user?.id) {
         documents.all.forEach((document) => policies.remove(document.id));
         try {
           await collections.fetchAll();
         } catch (err) {
           Logger.error(
-            "Failed to fetch collections after demote",
+            "Failed to fetch collections after role change",
             toError(err)
           );
         }
       }
-    });
+    };
+
+    socket.on("users.promote", handleRoleChange);
+    socket.on("users.demote", handleRoleChange);
 
     socket.on("users.delete", (event: WebsocketEntityDeletedEvent) => {
       users.remove(event.modelId);

server/models/User.test.ts

@@ -229,6 +229,22 @@ describe("user model", () => {
       expect(response[0]).toEqual(collection.id);
     });
 
+    it("should not return a cached response after a role change", async () => {
+      const team = await buildTeam();
+      const user = await buildGuestUser({
+        teamId: team.id,
+      });
+      const collection = await buildCollection({
+        teamId: team.id,
+        permission: CollectionPermission.ReadWrite,
+      });
+      expect(await user.collectionIds()).toEqual([]);
+
+      await user.update({ role: UserRole.Member });
+
+      expect(await user.collectionIds()).toEqual([collection.id]);
+    });
+
     it("should not return private collections", async () => {
       const team = await buildTeam();
       const user = await buildUser({

server/models/User.ts

@@ -482,12 +482,17 @@ class User extends ParanoidModel<
    * Returns the user's active collection ids. This includes collections the user
    * has access to through group memberships.
    *
-   * @param options Additional options to pass to the find
+   * @param options Additional options to pass to the find, set `skipCache` to bypass the cached response
    * @returns An array of collection ids
    */
-  public collectionIds = async (options: FindOptions<Collection> = {}) => {
+  public collectionIds = async (
+    options: FindOptions<Collection> & { skipCache?: boolean } = {}
+  ) => {
+    const { skipCache, ...findOptions } = options;
     const hasOptions =
-      options.transaction || options.paranoid === false || options.lock;
+      findOptions.transaction ||
+      findOptions.paranoid === false ||
+      findOptions.lock;
 
     const fetchCollectionIds = async () => {
       const collectionStubs = await Collection.findAll({
@@ -545,13 +550,13 @@ class User extends ParanoidModel<
           },
         ],
         paranoid: true,
-        ...options,
+        ...findOptions,
       });
 
       return Array.from(new Set(collectionStubs.map((c) => c.id)));
     };
 
-    if (hasOptions) {
+    if (hasOptions || skipCache) {
       return fetchCollectionIds();
     }
 
@@ -863,6 +868,30 @@ class User extends ParanoidModel<
     }
   }
 
+  // When a user's role changes their set of accessible collections may also
+  // change, so invalidate the cached collection ids.
+  @AfterUpdate
+  static async invalidateCollectionIdsAfterRoleChange(
+    model: User,
+    options: InstanceUpdateOptions<InferAttributes<User>>
+  ) {
+    if (!model.changed("role")) {
+      return;
+    }
+
+    const invalidate = () =>
+      CacheHelper.removeData(
+        RedisPrefixHelper.getUserCollectionIdsKey(model.id)
+      );
+
+    if (options.transaction) {
+      const transaction = options.transaction.parent || options.transaction;
+      transaction.afterCommit(invalidate);
+    } else {
+      await invalidate();
+    }
+  }
+
   // When a user's suspension state changes, invalidate the cached member count
   // for every group they belong to so the count reflects only active members.
   @AfterUpdate

server/queues/processors/WebsocketsProcessor.ts

@@ -332,16 +332,17 @@ export default class WebsocketsProcessor {
           return;
         }
 
+        // guests are excluded as they cannot read team-visible collections
+        // without an explicit membership.
+        const rooms = collection.isPrivate
+          ? [`user-${event.actorId}`]
+          : [`user-${event.actorId}`, `team-${collection.teamId}.members`];
+
         socketio
-          .to(this.getCollectionEventChannels(event, collection))
+          .to(rooms)
           .emit(event.name, await presentCollection(undefined, collection));
 
-        return socketio
-          .to(this.getCollectionEventChannels(event, collection))
-          .emit("join", {
-            event: event.name,
-            collectionId: collection.id,
-          });
+        return socketio.in(rooms).socketsJoin(`collection-${collection.id}`);
       }
 
       case "collections.update": {
@@ -352,9 +353,44 @@ export default class WebsocketsProcessor {
           return;
         }
 
-        return socketio
+        socketio
           .to(this.getCollectionEventChannels(event, collection))
           .emit(event.name, await presentCollection(undefined, collection));
+
+        const { attributes, previous } = event.changes ?? {};
+
+        // the collection became team-visible, all team members gain access.
+        if (attributes?.permission && previous?.permission === null) {
+          return socketio
+            .in(`team-${collection.teamId}.members`)
+            .socketsJoin(`collection-${collection.id}`);
+        }
+
+        // the collection became private, rebuild the channel from explicit
+        // memberships only.
+        if (attributes?.permission === null && previous?.permission) {
+          socketio
+            .in(`collection-${collection.id}`)
+            .socketsLeave(`collection-${collection.id}`);
+
+          const [memberships, groupMemberships] = await Promise.all([
+            UserMembership.findAll({
+              where: { collectionId: collection.id },
+            }),
+            GroupMembership.findAll({
+              where: { collectionId: collection.id },
+            }),
+          ]);
+          const rooms = [
+            ...memberships.map((m) => `user-${m.userId}`),
+            ...groupMemberships.map((m) => `group-${m.groupId}`),
+          ];
+          if (rooms.length) {
+            socketio.in(rooms).socketsJoin(`collection-${collection.id}`);
+          }
+        }
+
+        return;
       }
 
       case "collections.delete": {
@@ -412,12 +448,9 @@ export default class WebsocketsProcessor {
           .to(`user-${membership.userId}`)
           .to(`collection-${membership.collectionId}`)
           .emit(event.name, presentMembership(membership));
-
-        // tell any user clients to connect to the websocket channel for the collection
-        socketio.to(`user-${event.userId}`).emit("join", {
-          event: event.name,
-          collectionId: event.collectionId,
-        });
+        socketio
+          .in(`user-${event.userId}`)
+          .socketsJoin(`collection-${event.collectionId}`);
         return;
       }
 
@@ -444,11 +477,9 @@ export default class WebsocketsProcessor {
           .emit("collections.remove_user", membership);
 
         if (cannot(user, "read", collection)) {
-          // tell any user clients to disconnect from the websocket channel for the collection
-          socketio.to(`user-${event.userId}`).emit("leave", {
-            event: event.name,
-            collectionId: event.collectionId,
-          });
+          socketio
+            .in(`user-${event.userId}`)
+            .socketsLeave(`collection-${event.collectionId}`);
         }
 
         return;
@@ -466,11 +497,9 @@ export default class WebsocketsProcessor {
           .to(`group-${membership.groupId}`)
           .to(`collection-${membership.collectionId}`)
           .emit(event.name, presentGroupMembership(membership));
-
-        socketio.to(`group-${membership.groupId}`).emit("join", {
-          event: event.name,
-          collectionId: event.collectionId,
-        });
+        socketio
+          .in(`group-${membership.groupId}`)
+          .socketsJoin(`collection-${event.collectionId}`);
 
         return;
       }
@@ -504,11 +533,9 @@ export default class WebsocketsProcessor {
               }
 
               if (cannot(user, "read", collection)) {
-                // tell any user clients to disconnect from the websocket channel for the collection
-                socketio.to(`user-${groupUser.userId}`).emit("leave", {
-                  event: event.name,
-                  collectionId: event.collectionId,
-                });
+                socketio
+                  .in(`user-${groupUser.userId}`)
+                  .socketsLeave(`collection-${event.collectionId}`);
               }
             }
           }
@@ -700,11 +727,9 @@ export default class WebsocketsProcessor {
         socketio
           .to(`team-${event.teamId}`)
           .emit("groups.add_user", presentGroupUser(groupUser));
-
-        socketio.to(`user-${event.userId}`).emit("join", {
-          event: event.name,
-          groupId: event.modelId,
-        });
+        socketio
+          .in(`user-${event.userId}`)
+          .socketsJoin(`group-${event.modelId}`);
 
         await GroupMembership.findAllInBatches<GroupMembership>(
           {
@@ -722,12 +747,9 @@ export default class WebsocketsProcessor {
                     "collections.add_group",
                     presentGroupMembership(groupMembership)
                   );
-
-                // tell any user clients to connect to the websocket channel for the collection
-                socketio.to(`user-${event.userId}`).emit("join", {
-                  event: event.name,
-                  collectionId: groupMembership.collectionId,
-                });
+                socketio
+                  .in(`user-${event.userId}`)
+                  .socketsJoin(`collection-${groupMembership.collectionId}`);
               }
               if (groupMembership.documentId) {
                 socketio
@@ -755,11 +777,9 @@ export default class WebsocketsProcessor {
         socketio
           .to(`team-${event.teamId}`)
           .emit("groups.remove_user", membership);
-
-        socketio.to(`user-${event.userId}`).emit("leave", {
-          event: event.name,
-          groupId: event.modelId,
-        });
+        socketio
+          .in(`user-${event.userId}`)
+          .socketsLeave(`group-${event.modelId}`);
 
         const user = await User.findByPk(event.userId);
         if (!user) {
@@ -794,11 +814,9 @@ export default class WebsocketsProcessor {
               );
 
               if (cannot(user, "read", collection)) {
-                // tell any user clients to disconnect from the websocket channel for the collection
-                socketio.to(`user-${event.userId}`).emit("leave", {
-                  event: event.name,
-                  collectionId: groupMembership.collectionId,
-                });
+                socketio
+                  .in(`user-${event.userId}`)
+                  .socketsLeave(`collection-${groupMembership.collectionId}`);
               }
             }
           }
@@ -811,10 +829,9 @@ export default class WebsocketsProcessor {
         socketio.to(`team-${event.teamId}`).emit(event.name, {
           modelId: event.modelId,
         });
-        socketio.to(`group-${event.modelId}`).emit("leave", {
-          event: event.name,
-          groupId: event.modelId,
-        });
+        socketio
+          .in(`group-${event.modelId}`)
+          .socketsLeave(`group-${event.modelId}`);
 
         const groupMemberships = await GroupMembership.findAll({
           where: {
@@ -853,11 +870,11 @@ export default class WebsocketsProcessor {
                   );
 
                   if (cannot(groupUser.user, "read", collection)) {
-                    // tell any user clients to disconnect from the websocket channel for the collection
-                    socketio.to(`user-${groupUser.userId}`).emit("leave", {
-                      event: event.name,
-                      collectionId: groupMembership.collectionId,
-                    });
+                    socketio
+                      .in(`user-${groupUser.userId}`)
+                      .socketsLeave(
+                        `collection-${groupMembership.collectionId}`
+                      );
                   }
                 }
               }
@@ -915,16 +932,60 @@ export default class WebsocketsProcessor {
         return;
       }
 
+      case "users.promote":
       case "users.demote": {
-        return socketio
+        socketio
           .to(`user-${event.userId}`)
           .emit(event.name, { id: event.userId });
+
+        // the user's accessible collections may have changed with their role.
+        const user = await User.findByPk(event.userId);
+        if (!user) {
+          return;
+        }
+
+        const membersRoom = `team-${user.teamId}.members`;
+        const accessibleCollectionIds = await user.collectionIds({
+          skipCache: true,
+        });
+
+        if (user.isGuest) {
+          const teamCollections = await Collection.findAll({
+            attributes: ["id"],
+            where: { teamId: user.teamId },
+          });
+          const accessibleCollectionIdsSet = new Set(accessibleCollectionIds);
+          const inaccessibleRooms = teamCollections
+            .map((collection) => collection.id)
+            .filter((id) => !accessibleCollectionIdsSet.has(id))
+            .map((id) => `collection-${id}`);
+
+          return socketio
+            .in(`user-${user.id}`)
+            .socketsLeave([membersRoom, ...inaccessibleRooms]);
+        }
+
+        return socketio
+          .in(`user-${user.id}`)
+          .socketsJoin([
+            membersRoom,
+            ...accessibleCollectionIds.map((id) => `collection-${id}`),
+          ]);
+      }
+
+      case "users.signout":
+      case "users.suspend": {
+        // authentication is no longer valid, disconnect all of the user's clients.
+        return socketio.in(`user-${event.userId}`).disconnectSockets(true);
       }
 
       case "users.delete": {
-        return socketio
+        socketio
           .to(`team-${event.teamId}`)
           .emit(event.name, { modelId: event.userId });
+
+        // authentication is no longer valid, disconnect all of the user's clients.
+        return socketio.in(`user-${event.userId}`).disconnectSockets(true);
       }
 
       case "userMemberships.update": {
@@ -948,10 +1009,11 @@ export default class WebsocketsProcessor {
       channels.push(`user-${event.actorId}`);
     }
 
-    if (collection.isPrivate) {
-      channels.push(`collection-${collection.id}`);
-    } else {
-      channels.push(`team-${collection.teamId}`);
+    // guests are never included in the team members channel, they receive
+    // events via the collection channel when they have an explicit membership.
+    channels.push(`collection-${collection.id}`);
+    if (!collection.isPrivate) {
+      channels.push(`team-${collection.teamId}.members`);
     }
 
     return channels;

server/services/websockets.ts

@@ -12,8 +12,7 @@ import Logger from "@server/logging/Logger";
 import Metrics from "@server/logging/Metrics";
 import * as Tracing from "@server/logging/tracer";
 import { traceFunction } from "@server/logging/tracing";
-import { Collection, Group, User } from "@server/models";
-import { can } from "@server/policies";
+import type { User } from "@server/models";
 import Redis from "@server/storage/redis";
 import ShutdownHelper, { ShutdownOrder } from "@server/utils/ShutdownHelper";
 import { getUserForJWT } from "@server/utils/jwt";
@@ -183,9 +182,14 @@ async function authenticated(
   // and user so we can send authenticated events
   const rooms = [`team-${user.teamId}`, `user-${user.id}`];
 
-  // the rooms associated with collections this user has access to on
-  // connection. New collection and group subscriptions are managed
-  // from the client as needed through the 'join' event.
+  // the room of non-guest team members, who have access to all
+  // team-visible collections.
+  if (!user.isGuest) {
+    rooms.push(`team-${user.teamId}.members`);
+  }
+
+  // the rooms associated with collections and groups this user
+  // has access to on connection.
   const [collectionIds, groupIds] = await Promise.all([
     user.collectionIds(),
     user.groupIds(),
@@ -194,49 +198,6 @@ async function authenticated(
   collectionIds.forEach((colId) => rooms.push(`collection-${colId}`));
   groupIds.forEach((groupId) => rooms.push(`group-${groupId}`));
 
-  const { id: userId } = user;
-
-  // allow the client to request to join rooms
-  socket.on("join", async (event) => {
-    // The user is reloaded here rather than captured by this handler so the
-    // full model is not retained for the lifetime of the connection.
-    const authUser = await User.findByPk(userId);
-    if (!authUser) {
-      return;
-    }
-
-    // user is joining a collection channel, because their permissions have
-    // changed, granting them access.
-    if (event.collectionId) {
-      const collection = await Collection.findByPk(event.collectionId, {
-        userId: authUser.id,
-      });
-
-      if (can(authUser, "read", collection)) {
-        await socket.join(`collection-${event.collectionId}`);
-      }
-    }
-    if (event.groupId) {
-      const group = await Group.scope({
-        method: ["withMembership", authUser.id],
-      }).findByPk(event.groupId);
-
-      if (can(authUser, "read", group)) {
-        await socket.join(`group-${event.groupId}`);
-      }
-    }
-  });
-
-  // allow the client to request to leave rooms
-  socket.on("leave", async (event) => {
-    if (event.collectionId) {
-      await socket.leave(`collection-${event.collectionId}`);
-    }
-    if (event.groupId) {
-      await socket.leave(`group-${event.groupId}`);
-    }
-  });
-
   // join all of the rooms at once
   await socket.join(rooms);
 }