diff --git a/backend/open_webui/routers/terminals.py b/backend/open_webui/routers/terminals.py index 1d89b4754c..441b53598c 100644 --- a/backend/open_webui/routers/terminals.py +++ b/backend/open_webui/routers/terminals.py @@ -17,7 +17,6 @@ from open_webui.events import EVENTS, publish_event from open_webui.env import AIOHTTP_CLIENT_SESSION_SSL from open_webui.models.config import Config from open_webui.models.groups import Groups -from open_webui.models.users import Users from open_webui.utils.access_control import has_connection_access from open_webui.utils.auth import get_verified_user from open_webui.utils.terminals import get_terminal_server_url @@ -208,7 +207,7 @@ async def _resolve_authenticated_connection(ws: WebSocket, server_id: str): import asyncio import json - from open_webui.utils.auth import decode_token, is_valid_token + from open_webui.utils.auth import get_verified_user_by_token # First-message authentication try: @@ -217,14 +216,9 @@ async def _resolve_authenticated_connection(ws: WebSocket, server_id: str): if payload.get('type') != 'auth': await ws.close(code=4001, reason='Expected auth message') return None - token = payload.get('token', '') - data = decode_token(token) - if data is None or 'id' not in data or not await is_valid_token(data, getattr(ws.app.state, 'redis', None)): - await ws.close(code=4001, reason='Invalid token') - return None - user = await Users.get_user_by_id(data['id']) + user = await get_verified_user_by_token(payload.get('token', ''), getattr(ws.app.state, 'redis', None)) if user is None: - await ws.close(code=4001, reason='User not found') + await ws.close(code=4001, reason='Invalid token') return None except (asyncio.TimeoutError, json.JSONDecodeError): await ws.close(code=4001, reason='Auth timeout or invalid payload') diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index 67fbd5bcdd..691b91f322 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -36,7 +36,7 @@ from open_webui.models.users import UserNameResponse, Users from open_webui.socket.utils import RedisDict, RedisLock, YdocManager from open_webui.tasks import create_task, stop_item_tasks from open_webui.utils.access_control import has_permission -from open_webui.utils.auth import decode_token, is_valid_token +from open_webui.utils.auth import get_verified_user_by_token from open_webui.utils.redis import ( build_sentinel_url, get_redis_connection, @@ -368,10 +368,7 @@ async def connect(sid, environ, auth): scope = (environ or {}).get('asgi.scope') or {} fastapi_app = scope.get('app') redis = getattr(getattr(fastapi_app, 'state', None), 'redis', None) or REDIS - data = decode_token(auth['token']) - - if data is not None and 'id' in data and await is_valid_token(data, redis): - user = await Users.get_user_by_id(data['id']) + user = await get_verified_user_by_token(auth['token'], redis) if user: SESSION_POOL[sid] = { @@ -399,19 +396,14 @@ async def user_join(sid, data): scope = environ.get('asgi.scope') or {} fastapi_app = scope.get('app') redis = getattr(getattr(fastapi_app, 'state', None), 'redis', None) or REDIS - token_data = decode_token(auth['token']) - if token_data is None or 'id' not in token_data or not await is_valid_token(token_data, redis): + user = await get_verified_user_by_token(auth['token'], redis) + if not user: return existing = SESSION_POOL.get(sid) - if existing and existing.get('id') == token_data['id']: + if existing and existing.get('id') == user.id: SESSION_POOL[sid] = {**existing, 'last_seen_at': int(time.time())} - user_id, user_name, user_role = existing['id'], existing['name'], existing['role'] else: - user = await Users.get_user_by_id(token_data['id']) - if not user: - return - SESSION_POOL[sid] = { **user.model_dump( exclude=[ @@ -425,16 +417,15 @@ async def user_join(sid, data): 'last_seen_at': int(time.time()), } await sio.enter_room(sid, f'user:{user.id}') - user_id, user_name, user_role = user.id, user.name, user.role # Join all the channels only if user has channels permission - if user_role == 'admin' or await has_permission(user_id, 'features.channels'): - channels = await Channels.get_channels_by_user_id(user_id) + if user.role == 'admin' or await has_permission(user.id, 'features.channels'): + channels = await Channels.get_channels_by_user_id(user.id) log.debug(f'{channels=}') for channel in channels: await sio.enter_room(sid, f'channel:{channel.id}') - return {'id': user_id, 'name': user_name} + return {'id': user.id, 'name': user.name} @sio.on('heartbeat') @@ -455,11 +446,7 @@ async def join_channel(sid, data): scope = environ.get('asgi.scope') or {} fastapi_app = scope.get('app') redis = getattr(getattr(fastapi_app, 'state', None), 'redis', None) or REDIS - data = decode_token(auth['token']) - if data is None or 'id' not in data or not await is_valid_token(data, redis): - return - - user = await Users.get_user_by_id(data['id']) + user = await get_verified_user_by_token(auth['token'], redis) if not user: return @@ -481,11 +468,7 @@ async def join_note(sid, data): scope = environ.get('asgi.scope') or {} fastapi_app = scope.get('app') redis = getattr(getattr(fastapi_app, 'state', None), 'redis', None) or REDIS - token_data = decode_token(auth['token']) - if token_data is None or 'id' not in token_data or not await is_valid_token(token_data, redis): - return - - user = await Users.get_user_by_id(token_data['id']) + user = await get_verified_user_by_token(auth['token'], redis) if not user: return diff --git a/backend/open_webui/utils/auth.py b/backend/open_webui/utils/auth.py index f9d596e615..0384d6fcb9 100644 --- a/backend/open_webui/utils/auth.py +++ b/backend/open_webui/utils/auth.py @@ -488,8 +488,11 @@ async def get_current_user_by_api_key(request, api_key: str): return user +VERIFIED_USER_ROLES = {'user', 'admin'} + + def get_verified_user(user=Depends(get_current_user)): - if user.role not in {'user', 'admin'}: + if user.role not in VERIFIED_USER_ROLES: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.ACCESS_PROHIBITED, @@ -497,6 +500,19 @@ def get_verified_user(user=Depends(get_current_user)): return user +async def get_verified_user_by_token(token: str, redis=None): + """Resolve a verified user from a raw token, for WebSocket handshakes that run outside the HTTP dependency chain.""" + decoded = decode_token(token) + if decoded is None or 'id' not in decoded or not await is_valid_token(decoded, redis): + return None + + user = await Users.get_user_by_id(decoded['id']) + if user is None or user.role not in VERIFIED_USER_ROLES: + return None + + return user + + def get_admin_user(user=Depends(get_current_user)): if user.role != 'admin': raise HTTPException(