From e93a59f4dd8eea5b1df55b4e6dcc57465887b95d Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Tue, 22 Sep 2026 13:13:22 -0400 Subject: [PATCH] refac Co-Authored-By: Classic298 <27028174+Classic298@users.noreply.github.com> --- backend/open_webui/models/access_grants.py | 75 +++++++++++++++++----- backend/open_webui/routers/channels.py | 12 ++++ backend/open_webui/routers/notes.py | 30 ++++++++- backend/open_webui/socket/main.py | 25 +++++++- 4 files changed, 121 insertions(+), 21 deletions(-) diff --git a/backend/open_webui/models/access_grants.py b/backend/open_webui/models/access_grants.py index 03a7ef7202..49cb83a180 100644 --- a/backend/open_webui/models/access_grants.py +++ b/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, diff --git a/backend/open_webui/routers/channels.py b/backend/open_webui/routers/channels.py index de25739c92..80e243c8ae 100644 --- a/backend/open_webui/routers/channels.py +++ b/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, diff --git a/backend/open_webui/routers/notes.py b/backend/open_webui/routers/notes.py index a84dd3f94e..e44a311f45 100644 --- a/backend/open_webui/routers/notes.py +++ b/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, diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index 2e0969ab2e..659dc3da49 100644 --- a/backend/open_webui/socket/main.py +++ b/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.