refac

Co-Authored-By: Classic298 <27028174+Classic298@users.noreply.github.com>

backend/open_webui/models/access_grants.py

@@ -686,8 +686,7 @@ class AccessGrantsTable:
         Get all users who have the specified permission on a resource.
         Returns a list of UserModel instances.
         """
-        from open_webui.models.groups import Groups
-        from open_webui.models.users import UserModel, Users
+        from open_webui.models.users import Users
 
         async with get_async_db_context(db) as db:
             result = await db.execute(
@@ -699,27 +698,69 @@ class AccessGrantsTable:
             )
             grants = result.scalars().all()
 
-            # Check for public access
-            for grant in grants:
-                if grant.principal_type == 'user' and grant.principal_id == '*':
-                    result = await Users.get_users(filter={'roles': ['!pending']}, db=db)
-                    return result.get('users', [])
-
-            user_ids_with_access = set()
-
-            for grant in grants:
-                if grant.principal_type == 'user':
-                    user_ids_with_access.add(grant.principal_id)
-                elif grant.principal_type == 'group':
-                    group_user_ids = await Groups.get_group_user_ids_by_id(grant.principal_id, db=db)
-                    if group_user_ids:
-                        user_ids_with_access.update(group_user_ids)
+            user_ids_with_access = await self.get_user_ids_by_access_grants(grants, permission, db=db)
 
             if not user_ids_with_access:
                 return []
 
             return await Users.get_users_by_user_ids(list(user_ids_with_access), db=db)
 
+    async def get_user_ids_by_access_grants(
+        self,
+        access_grants: list[AccessGrantModel],
+        permission: str = 'read',
+        db: AsyncSession | None = None,
+    ) -> set[str]:
+        """Get user IDs with the specified permission, including public and group grants."""
+        from open_webui.models.groups import Groups
+        from open_webui.models.users import Users
+
+        async with get_async_db_context(db) as db:
+            user_ids = set()
+            group_ids = []
+            for grant in access_grants:
+                if grant.permission != permission:
+                    continue
+                if grant.principal_type == PRINCIPAL_TYPE_USER:
+                    if grant.principal_id == WILDCARD_PRINCIPAL_ID:
+                        result = await Users.get_users(filter={'roles': ['!pending']}, db=db)
+                        return {user.id for user in result.get('users', [])}
+                    user_ids.add(grant.principal_id)
+                elif grant.principal_type == PRINCIPAL_TYPE_GROUP:
+                    group_ids.append(grant.principal_id)
+
+            if group_ids:
+                group_user_ids = await Groups.get_group_user_ids_by_ids(group_ids, db=db)
+                for members in group_user_ids.values():
+                    user_ids.update(members)
+            return user_ids
+
+    async def get_revoked_user_ids_by_resource(
+        self,
+        resource_type: str,
+        resource_id: str,
+        previous_access_grants: list[AccessGrantModel],
+        permission: str = 'read',
+        db: AsyncSession | None = None,
+    ) -> set[str]:
+        """Get user IDs that lost the specified permission after a resource's grants changed."""
+        async with get_async_db_context(db) as db:
+            access_grants = await self.get_grants_by_resource(resource_type, resource_id, db=db)
+            previous_principals = {
+                (grant.principal_type, grant.principal_id)
+                for grant in previous_access_grants
+                if grant.permission == permission
+            }
+            principals = {
+                (grant.principal_type, grant.principal_id) for grant in access_grants if grant.permission == permission
+            }
+            if previous_principals <= principals or (PRINCIPAL_TYPE_USER, WILDCARD_PRINCIPAL_ID) in principals:
+                return set()
+
+            previous_user_ids = await self.get_user_ids_by_access_grants(previous_access_grants, permission, db=db)
+            user_ids = await self.get_user_ids_by_access_grants(access_grants, permission, db=db)
+            return previous_user_ids - user_ids
+
     def has_permission_filter(
         self,
         db,

backend/open_webui/routers/channels.py

@@ -40,6 +40,7 @@ from open_webui.socket.main import (
     emit_to_users,
     enter_room_for_users,
     get_user_ids_from_room,
+    leave_room_for_users,
     sio,
 )
 from open_webui.utils.access_control import filter_allowed_access_grants, has_permission
@@ -731,8 +732,18 @@ async def update_channel_by_id(
         'sharing.public_channels',
     )
 
+    previous_access_grants = channel.access_grants
+
     try:
         channel = await Channels.update_channel_by_id(id, form_data, db=db)
+        # Group and DM channels use membership instead of access grants.
+        if form_data.access_grants is not None and channel.type not in ['group', 'dm']:
+            revoked_user_ids = await AccessGrants.get_revoked_user_ids_by_resource(
+                'channel', id, previous_access_grants, db=db
+            )
+            revoked_user_ids.discard(channel.user_id)
+            await leave_room_for_users(f'channel:{id}', list(revoked_user_ids))
+
         await publish_event(
             request,
             EVENTS.CHANNEL_UPDATED,
@@ -769,6 +780,7 @@ async def delete_channel_by_id(
 
     try:
         await Channels.delete_channel_by_id(id, db=db)
+        await sio.close_room(f'channel:{id}')
         await publish_event(
             request,
             EVENTS.CHANNEL_DELETED,

backend/open_webui/routers/notes.py

@@ -11,7 +11,7 @@ from open_webui.config import (
 from open_webui.constants import ERROR_MESSAGES
 from open_webui.events import EVENTS, publish_event
 from open_webui.internal.db import get_async_session
-from open_webui.models.access_grants import AccessGrants
+from open_webui.models.access_grants import AccessGrantModel, AccessGrants
 from open_webui.models.chats import ChatForm, ChatResponse, Chats
 from open_webui.models.config import Config
 from open_webui.models.groups import Groups
@@ -23,7 +23,7 @@ from open_webui.models.notes import (
     NoteUserResponse,
 )
 from open_webui.models.users import UserResponse, Users
-from open_webui.socket.main import sio
+from open_webui.socket.main import leave_room_for_users, sio
 from open_webui.utils.access_control import (
     filter_allowed_access_grants,
     has_permission,
@@ -46,6 +46,23 @@ def _truncate_note_data(data: Optional[dict], max_length: int = 1000) -> Optiona
     return {'content': {'md': md[:max_length]}}
 
 
+async def leave_note_rooms_for_revoked_users(
+    note: NoteModel, previous_access_grants: list[AccessGrantModel], db: AsyncSession | None = None
+):
+    revoked_user_ids = await AccessGrants.get_revoked_user_ids_by_resource(
+        'note', note.id, previous_access_grants, db=db
+    )
+    revoked_user_ids.discard(note.user_id)
+    if not revoked_user_ids:
+        return
+
+    users = await Users.get_users_by_user_ids(list(revoked_user_ids), db=db)
+    # Admins retain access to notes regardless of grants.
+    user_ids = [user.id for user in users if user.role != 'admin']
+    for room in [f'note:{note.id}', f'doc_note:{note.id}']:
+        await leave_room_for_users(room, user_ids)
+
+
 ############################
 # GetNotes
 ############################
@@ -564,8 +581,13 @@ async def update_note_by_id(
             db=db,
         )
 
+    previous_access_grants = note.access_grants
+
     try:
         note = await Notes.update_note_by_id(id, form_data, db=db)
+        if form_data.access_grants is not None:
+            await leave_note_rooms_for_revoked_users(note, previous_access_grants, db=db)
+
         pinned_note_ids = await Notes.get_pinned_note_ids(user.id, db=db)
         note.is_pinned = note.id in pinned_note_ids
 
@@ -644,6 +666,7 @@ async def update_note_access_by_id(
     )
 
     await AccessGrants.set_access_grants('note', id, form_data.access_grants, db=db)
+    await leave_note_rooms_for_revoked_users(note, note.access_grants, db=db)
 
     note = await Notes.get_note_by_id(id, db=db)
     pinned_note_ids = await Notes.get_pinned_note_ids(user.id, db=db)
@@ -744,6 +767,9 @@ async def delete_note_by_id(
 
     try:
         note = await Notes.delete_note_by_id(id, db=db)
+        for room in [f'note:{id}', f'doc_note:{id}']:
+            await sio.close_room(room)
+
         await publish_event(
             request,
             EVENTS.NOTE_DELETED,

backend/open_webui/socket/main.py

@@ -343,8 +343,20 @@ def get_session_ids_from_room(room):
 
 def get_session_ids_by_user_id(user_id: str) -> list[str]:
     """Get known session IDs for a user across the local rooms and shared session pool."""
-    session_ids = set(get_session_ids_from_room(f'user:{user_id}'))
-    session_ids.update(sid for sid, entry in SESSION_POOL.items() if entry and entry.get('id') == user_id)
+    return get_session_ids_by_user_ids([user_id])
+
+
+def get_session_ids_by_user_ids(user_ids: list[str]) -> list[str]:
+    """Get known session IDs for users across the local rooms and shared session pool."""
+    if not user_ids:
+        return []
+
+    user_ids = set(user_ids)
+    session_ids = set()
+    for user_id in user_ids:
+        session_ids.update(get_session_ids_from_room(f'user:{user_id}'))
+    for batch in get_session_pool_batches():
+        session_ids.update(sid for sid, user in batch if user and user.get('id') in user_ids)
     return list(session_ids)
 
 
@@ -385,6 +397,15 @@ async def enter_room_for_users(room: str, user_ids: list[str]):
         log.debug('Failed to make users %s join room %s: %s', user_ids, room, e)
 
 
+async def leave_room_for_users(room: str, user_ids: list[str]):
+    """Make all sessions of each user leave a room, including sessions on other workers."""
+    for sid in get_session_ids_by_user_ids(user_ids):
+        try:
+            await sio.leave_room(sid, room)
+        except Exception as e:
+            log.debug('Failed to make session %s leave room %s: %s', sid, room, e)
+
+
 async def disconnect_user_sessions(user_id: str):
     """Disconnect all Socket.IO sessions belonging to a user.