From 3919fb4fc52cfd005ca366974b9fbbd880637b0a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 19 May 2026 11:36:05 +0000 Subject: [PATCH] fix: lock down task stop and trusted email checks Co-authored-by: mkatwi --- backend/open_webui/main.py | 14 +++- backend/open_webui/tasks.py | 14 ++++ backend/open_webui/test/utils/test_auth.py | 75 +++++++++++++++++++++ backend/open_webui/test/utils/test_tasks.py | 37 ++++++++++ backend/open_webui/utils/auth.py | 2 +- 5 files changed, 140 insertions(+), 2 deletions(-) create mode 100644 backend/open_webui/test/utils/test_auth.py create mode 100644 backend/open_webui/test/utils/test_tasks.py diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index 544756a6e8f4..6ba199970a0e 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -450,6 +450,7 @@ from open_webui.tasks import ( redis_task_command_listener, list_task_ids_by_chat_id, + get_task_chat_id, stop_task, list_tasks, ) # Import from tasks.py @@ -1447,6 +1448,17 @@ async def stop_task_endpoint( request: Request, task_id: str, user=Depends(get_verified_user) ): try: + chat_id = await get_task_chat_id(request, task_id) + if chat_id is None: + raise ValueError(f"Task with ID {task_id} not found.") + + chat = Chats.get_chat_by_id(chat_id) + if chat is None or (chat.user_id != user.id and user.role != "admin"): + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Task with ID {task_id} not found.", + ) + result = await stop_task(request, task_id) return result except ValueError as e: @@ -1454,7 +1466,7 @@ async def stop_task_endpoint( @app.get("/api/tasks") -async def list_tasks_endpoint(request: Request, user=Depends(get_verified_user)): +async def list_tasks_endpoint(request: Request, user=Depends(get_admin_user)): return {"tasks": await list_tasks(request)} diff --git a/backend/open_webui/tasks.py b/backend/open_webui/tasks.py index 2d3955f0a2b7..3ff8b04d4da8 100644 --- a/backend/open_webui/tasks.py +++ b/backend/open_webui/tasks.py @@ -135,6 +135,20 @@ async def list_task_ids_by_chat_id(request, id): return chat_tasks.get(id, []) +async def get_task_chat_id(request, task_id: str) -> Optional[str]: + """ + Return the chat ID associated with a task, if the task is chat-scoped. + """ + if is_redis(request): + return await request.app.state.redis.hget(REDIS_TASKS_KEY, task_id) + + for chat_id, task_ids in chat_tasks.items(): + if chat_id and task_id in task_ids: + return chat_id + + return None + + async def stop_task(request, task_id: str): """ Cancel a running task and remove it from the global task list. diff --git a/backend/open_webui/test/utils/test_auth.py b/backend/open_webui/test/utils/test_auth.py new file mode 100644 index 000000000000..f86331ded8ca --- /dev/null +++ b/backend/open_webui/test/utils/test_auth.py @@ -0,0 +1,75 @@ +from types import SimpleNamespace + +import pytest +from fastapi import BackgroundTasks, HTTPException, Response +from fastapi.security import HTTPAuthorizationCredentials +from starlette.requests import Request + +from open_webui.utils import auth + + +def _request_with_headers(headers: dict[str, str]) -> Request: + scope = { + "type": "http", + "method": "GET", + "path": "/api/config", + "headers": [ + (key.lower().encode(), value.encode()) for key, value in headers.items() + ], + "query_string": b"", + "server": ("testserver", 80), + "scheme": "http", + "client": ("testclient", 50000), + } + return Request(scope) + + +def _credentials() -> HTTPAuthorizationCredentials: + return HTTPAuthorizationCredentials(scheme="Bearer", credentials="jwt-token") + + +def test_trusted_email_header_allows_case_insensitive_user_email(monkeypatch): + monkeypatch.setattr(auth, "WEBUI_AUTH_TRUSTED_EMAIL_HEADER", "X-Forwarded-Email") + monkeypatch.setattr(auth, "decode_token", lambda token: {"id": "user-1"}) + monkeypatch.setattr( + auth.Users, + "get_user_by_id", + lambda user_id: SimpleNamespace( + id=user_id, + email="User@Example.com", + role="user", + ), + ) + + user = auth.get_current_user( + _request_with_headers({"X-Forwarded-Email": "user@example.com"}), + Response(), + BackgroundTasks(), + _credentials(), + ) + + assert user.email == "User@Example.com" + + +def test_trusted_email_header_rejects_different_email(monkeypatch): + monkeypatch.setattr(auth, "WEBUI_AUTH_TRUSTED_EMAIL_HEADER", "X-Forwarded-Email") + monkeypatch.setattr(auth, "decode_token", lambda token: {"id": "user-1"}) + monkeypatch.setattr( + auth.Users, + "get_user_by_id", + lambda user_id: SimpleNamespace( + id=user_id, + email="other@example.com", + role="user", + ), + ) + + with pytest.raises(HTTPException) as exc_info: + auth.get_current_user( + _request_with_headers({"X-Forwarded-Email": "user@example.com"}), + Response(), + BackgroundTasks(), + _credentials(), + ) + + assert exc_info.value.status_code == 401 diff --git a/backend/open_webui/test/utils/test_tasks.py b/backend/open_webui/test/utils/test_tasks.py new file mode 100644 index 000000000000..c8f34562e8b8 --- /dev/null +++ b/backend/open_webui/test/utils/test_tasks.py @@ -0,0 +1,37 @@ +import asyncio +from types import SimpleNamespace + +from open_webui import tasks + + +def _request(redis=None): + return SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(redis=redis))) + + +def test_get_task_chat_id_ignores_unscoped_local_tasks(monkeypatch): + monkeypatch.setattr( + tasks, + "chat_tasks", + { + "chat-1": ["task-1"], + None: ["task-without-chat"], + "chat-2": ["task-2"], + }, + ) + + assert asyncio.run(tasks.get_task_chat_id(_request(), "task-1")) == "chat-1" + assert asyncio.run(tasks.get_task_chat_id(_request(), "task-without-chat")) is None + assert asyncio.run(tasks.get_task_chat_id(_request(), "missing-task")) is None + + +def test_get_task_chat_id_reads_redis_mapping(): + class FakeRedis: + async def hget(self, key, task_id): + assert key == tasks.REDIS_TASKS_KEY + assert task_id == "task-redis" + return "chat-redis" + + assert ( + asyncio.run(tasks.get_task_chat_id(_request(redis=FakeRedis()), "task-redis")) + == "chat-redis" + ) diff --git a/backend/open_webui/utils/auth.py b/backend/open_webui/utils/auth.py index 9befaf2a910b..dceece9a6712 100644 --- a/backend/open_webui/utils/auth.py +++ b/backend/open_webui/utils/auth.py @@ -231,7 +231,7 @@ def get_current_user( trusted_email = request.headers.get( WEBUI_AUTH_TRUSTED_EMAIL_HEADER, "" ).lower() - if trusted_email and user.email != trusted_email: + if trusted_email and user.email.lower() != trusted_email: # Delete the token cookie response.delete_cookie("token") # Delete OAuth token if present