From 175010cb85afc6582d6a2bd7698a21ed659980d7 Mon Sep 17 00:00:00 2001 From: carfeii Date: Mon, 5 Oct 2026 23:12:18 +0800 Subject: [PATCH] fix(api): enforce WIKI_AUTH_MODE everywhere and cap task admission Two related fixes: - WIKI_AUTH_MODE previously only protected DELETE /api/wiki_cache; every other sensitive/costly route (wiki generation, repo indexing, chat, codemap, cache reads) performed no authorization check at all, so an operator who enabled the one documented protection mechanism got no actual protection anywhere but that single endpoint. Adds a shared check_wiki_auth_code() dependency (api/routers/auth.py) and an authorization_code field on RepoRequestBase, applied to every route that generates, indexes, chats, or reads cached wiki content. - Wiki task submission (POST /wiki/tasks) and repo indexing (POST /repo/prepare) each admitted a new asyncio task into their respective tracking state unconditionally, before any concurrency limit applied. A caller varying repo_url could queue an unbounded number of pending clone/embedding jobs. Adds an admission cap on distinct non-terminal tasks (TaskRegistry.max_pending, default 10x the execution semaphore) and a parallel cap on /repo/prepare's own in-flight job set, both returning 429 once full. Adds regression tests in tests/backend/routers/test_wiki_auth_enforcement.py (new) and tests/backend/services/test_wiki_task.py, including fail-before coverage run against the unpatched code. --- api/routers/auth.py | 16 +- api/routers/chat.py | 7 + api/routers/codemap.py | 9 ++ api/routers/repo.py | 88 ++++++++--- api/routers/wiki.py | 23 +-- api/schemas/base.py | 4 + api/services/wiki/__init__.py | 2 + api/services/wiki/tasks.py | 25 ++- .../routers/test_wiki_auth_enforcement.py | 147 ++++++++++++++++++ tests/backend/services/test_wiki_task.py | 47 ++++++ 10 files changed, 334 insertions(+), 34 deletions(-) create mode 100644 tests/backend/routers/test_wiki_auth_enforcement.py diff --git a/api/routers/auth.py b/api/routers/auth.py index 9607e2317..3333433e4 100644 --- a/api/routers/auth.py +++ b/api/routers/auth.py @@ -1,4 +1,4 @@ -from fastapi import APIRouter +from fastapi import APIRouter, HTTPException from api.config import WIKI_AUTH_CODE, WIKI_AUTH_MODE from api.schemas import AuthorizationConfig @@ -6,6 +6,20 @@ router = APIRouter(prefix="/auth", tags=["auth"]) +def check_wiki_auth_code(authorization_code: str | None) -> None: + """Enforce WIKI_AUTH_MODE consistently across every protected route. + + Previously only DELETE /api/wiki_cache checked this inline; every other + sensitive/costly route (wiki generation, repo indexing, chat, codemap, + cache reads) performed no such check at all, so enabling WIKI_AUTH_MODE + gave callers no actual protection anywhere but that one endpoint. + """ + if not WIKI_AUTH_MODE: + return + if not authorization_code or WIKI_AUTH_CODE != authorization_code: + raise HTTPException(status_code=401, detail="Authorization code is invalid") + + @router.get("/status") async def get_auth_status(): """ diff --git a/api/routers/chat.py b/api/routers/chat.py index b33a92529..4a5cb91d7 100644 --- a/api/routers/chat.py +++ b/api/routers/chat.py @@ -3,6 +3,7 @@ from fastapi.websockets import WebSocketState from api.logger import get_logger +from api.routers.auth import check_wiki_auth_code from api.schemas import ChatCompletionRequest from api.services.research import RepoNotIndexedError, research_chat @@ -26,6 +27,11 @@ async def handle_websocket_chat(websocket: WebSocket): await websocket.accept() try: request = ChatCompletionRequest(**await websocket.receive_json()) + try: + check_wiki_auth_code(request.authorization_code) + except HTTPException as e: + await websocket.send_text(f"Error: {e.detail}") + return if not request.messages or len(request.messages) == 0: await websocket.send_text("Error: No messages provided") return @@ -69,6 +75,7 @@ async def handle_websocket_chat(websocket: WebSocket): @router.post("/chat/completions/stream") async def chat_completions_stream(request: ChatCompletionRequest): """Stream a chat completion response directly using Google Generative AI""" # Validate request + check_wiki_auth_code(request.authorization_code) if not request.messages or len(request.messages) == 0: raise HTTPException(status_code=400, detail="No messages provided") diff --git a/api/routers/codemap.py b/api/routers/codemap.py index ca4321db6..1d7638845 100644 --- a/api/routers/codemap.py +++ b/api/routers/codemap.py @@ -3,6 +3,7 @@ from fastapi.websockets import WebSocketState from api.logger import get_logger +from api.routers.auth import check_wiki_auth_code from api.schemas import CodeMapRequest from api.services.codemap import generate_codemap, read_repo_file @@ -17,6 +18,11 @@ async def handle_websocket_codemap(websocket: WebSocket): await websocket.accept() try: request = CodeMapRequest(**await websocket.receive_json()) + try: + check_wiki_auth_code(request.authorization_code) + except HTTPException as e: + await websocket.send_text(str(e.detail)) + return async for event in generate_codemap(request): if websocket.application_state != WebSocketState.CONNECTED: break @@ -39,6 +45,7 @@ async def handle_websocket_codemap(websocket: WebSocket): @router.post("/codemap/stream") async def codemap_stream(request: CodeMapRequest): """HTTP fallback: stream codemap generation events as NDJSON.""" + check_wiki_auth_code(request.authorization_code) try: return StreamingResponse( generate_codemap(request), @@ -54,8 +61,10 @@ async def codemap_file( repo_url: str = Query(..., description="Repository URL or local path"), file_path: str = Query(..., description="Repository-relative file path"), type: str = Query("github", description="Repository type"), + authorization_code: str | None = Query(None, description="Authorization code"), ): """Return the full content of a file from the cloned/local repository.""" + check_wiki_auth_code(authorization_code) try: content = read_repo_file(repo_url, type, file_path) return {"file_path": file_path, "content": content} diff --git a/api/routers/repo.py b/api/routers/repo.py index ea79e9db9..3ab511ddd 100644 --- a/api/routers/repo.py +++ b/api/routers/repo.py @@ -1,5 +1,6 @@ import asyncio import json +import os from fastapi import APIRouter, Query from fastapi.responses import StreamingResponse @@ -7,6 +8,7 @@ from api.logger import get_logger from api.rag import repo_index_exist from api.repository import Repo +from api.routers.auth import check_wiki_auth_code from api.schemas import RepoPrepareRequest from api.services.research import prepare_repo_index as prepare_index @@ -19,9 +21,27 @@ # (undici default = 300s). 10s keeps the connection alive with wide margin. _HEARTBEAT_INTERVAL_SEC = 10 +# Admission cap: distinct in-flight /repo/prepare indexing jobs allowed at +# once. This endpoint has no registry of its own (unlike wiki task +# submission); every request used to get an unconditional asyncio.create_task +# with no limit at all, so a caller varying repo_url could queue unbounded +# clone/embedding work. +try: + _MAX_CONCURRENT_PREPARE_TASKS = int( + os.environ.get("DEEPWIKI_MAX_CONCURRENT_PREPARE_TASKS", "10") + ) +except (TypeError, ValueError): + _MAX_CONCURRENT_PREPARE_TASKS = 10 + +_in_flight_prepare_keys: set[str] = set() + @router.post("/prepare") async def prepare_repo_index(request: RepoPrepareRequest): + check_wiki_auth_code(request.authorization_code) + + key = f"{request.type}_{request.repo_url}" + async def event_stream(): # Fast path: cache already warm -> return immediately. if repo_index_exist(Repo(repo_url=request.repo_url, repo_type=request.type)): @@ -29,32 +49,54 @@ async def event_stream(): yield "event: done\ndata: ok\n\n" return + if key in _in_flight_prepare_keys: + yield "event: ready\ndata: already in progress\n\n" + yield "event: done\ndata: ok\n\n" + return + + if len(_in_flight_prepare_keys) >= _MAX_CONCURRENT_PREPARE_TASKS: + yield ( + "event: error\ndata: " + + json.dumps( + { + "error": f"{_MAX_CONCURRENT_PREPARE_TASKS} indexing jobs are " + "already in flight; try again later" + } + ) + + "\n\n" + ) + return + # First byte -> response headers are flushed now (kills the timeout). yield ": indexing-start\n\n" - task = asyncio.create_task(prepare_index(request)) - elapsed = 0 - # Wait on the task, but wake up every interval to emit a heartbeat. - while not task.done(): - try: - # shield: a heartbeat timeout must NOT cancel the indexing task. - await asyncio.wait_for( - asyncio.shield(task), timeout=_HEARTBEAT_INTERVAL_SEC - ) - except asyncio.TimeoutError: - elapsed += _HEARTBEAT_INTERVAL_SEC - yield f"event: progress\ndata: {json.dumps({'elapsed_sec': elapsed})}\n\n" - except Exception: - # Task raised; break out and report via task.exception() below. - break - - exc = task.exception() - if exc is not None: - logger.error("Repo indexing failed for %s: %s", request.repo_url, exc) - yield f"event: error\ndata: {json.dumps({'error': str(exc)})}\n\n" - else: - logger.info("Repo indexing complete for %s", request.repo_url) - yield "event: done\ndata: ok\n\n" + _in_flight_prepare_keys.add(key) + try: + task = asyncio.create_task(prepare_index(request)) + elapsed = 0 + # Wait on the task, but wake up every interval to emit a heartbeat. + while not task.done(): + try: + # shield: a heartbeat timeout must NOT cancel the indexing task. + await asyncio.wait_for( + asyncio.shield(task), timeout=_HEARTBEAT_INTERVAL_SEC + ) + except asyncio.TimeoutError: + elapsed += _HEARTBEAT_INTERVAL_SEC + yield f"event: progress\ndata: {json.dumps({'elapsed_sec': elapsed})}\n\n" + except Exception: + # Task raised; break out and report via task.exception() below. + break + + exc = task.exception() + if exc is not None: + logger.error("Repo indexing failed for %s: %s", request.repo_url, exc) + yield f"event: error\ndata: {json.dumps({'error': str(exc)})}\n\n" + else: + logger.info("Repo indexing complete for %s", request.repo_url) + yield "event: done\ndata: ok\n\n" + finally: + _in_flight_prepare_keys.discard(key) return StreamingResponse( event_stream(), diff --git a/api/routers/wiki.py b/api/routers/wiki.py index fd66f4b74..5cd335091 100644 --- a/api/routers/wiki.py +++ b/api/routers/wiki.py @@ -6,8 +6,9 @@ from fastapi import APIRouter, HTTPException, Query from fastapi.responses import JSONResponse, Response, StreamingResponse -from api.config import WIKI_AUTH_CODE, WIKI_AUTH_MODE, configs +from api.config import configs from api.logger import get_logger +from api.routers.auth import check_wiki_auth_code from api.schemas import ( ProcessedProjectEntry, WikiCacheData, @@ -26,6 +27,7 @@ list_wiki_cache, read_wiki_cache, registry, + TaskQueueFullError, WikiTask, ) @@ -142,8 +144,10 @@ async def read_wiki( repo: str = Query(..., description="Repository name"), repo_type: str = Query(..., description="Repository type (e.g., github, gitlab)"), language: str = Query(..., description="Language of the wiki content"), + authorization_code: Optional[str] = Query(None, description="Authorization code"), ): """Retrieve cached wiki data (structure and generated pages) for a repository.""" + check_wiki_auth_code(authorization_code) supported_langs = configs["lang_config"]["supported_languages"] if language not in supported_langs: language = configs["lang_config"]["default"] @@ -172,16 +176,13 @@ async def delete_wiki( """ Deletes a specific wiki cache from the file system. """ + check_wiki_auth_code(authorization_code) + # Language validation supported_langs = configs["lang_config"]["supported_languages"] if language not in supported_langs: raise HTTPException(status_code=400, detail="Language is not supported") - if WIKI_AUTH_MODE: - logger.info("check the authorization code") - if not authorization_code or WIKI_AUTH_CODE != authorization_code: - raise HTTPException(status_code=401, detail="Authorization code is invalid") - logger.info( f"Attempting to delete wiki cache for {owner}/{repo} ({repo_type}), lang: {language}" ) @@ -222,10 +223,14 @@ async def submit_wiki_task(request: WikiTaskRequest): Returns one of: created (new task), joined (an active task for the repo already exists), or from_cache (this variant is already generated). """ + check_wiki_auth_code(request.authorization_code) - return await registry.submit( - WikiTask.from_wiki_request(request), async_func=generate_repo_wiki - ) + try: + return await registry.submit( + WikiTask.from_wiki_request(request), async_func=generate_repo_wiki + ) + except TaskQueueFullError as e: + raise HTTPException(status_code=429, detail=str(e)) @router.get( diff --git a/api/schemas/base.py b/api/schemas/base.py index ea99ea8e8..25a0d6464 100644 --- a/api/schemas/base.py +++ b/api/schemas/base.py @@ -29,6 +29,10 @@ class RepoRequestBase(BaseModel): default_factory=list, description="List or newline-separated string of file patterns to include exclusively", ) + authorization_code: str | None = Field( + None, + description="Authorization code, required on every request when WIKI_AUTH_MODE is enabled", + ) @field_validator( "excluded_dirs", diff --git a/api/services/wiki/__init__.py b/api/services/wiki/__init__.py index 26a91733d..4f3accc4a 100644 --- a/api/services/wiki/__init__.py +++ b/api/services/wiki/__init__.py @@ -13,6 +13,7 @@ WikiTask, registry, generate_repo_wiki, + TaskQueueFullError, ) __all__ = [ @@ -27,4 +28,5 @@ "WikiTask", "registry", "generate_repo_wiki", + "TaskQueueFullError", ] diff --git a/api/services/wiki/tasks.py b/api/services/wiki/tasks.py index 351bec138..b3fb36cdb 100644 --- a/api/services/wiki/tasks.py +++ b/api/services/wiki/tasks.py @@ -68,6 +68,19 @@ def _env_int(name, default: int) -> int: WIKI_PAGE_RETRIES = _env_int("DEEPWIKI_WIKI_PAGE_RETRIES", 2) # How long a terminal (COMPLETED/FAILED) task lingers in the registry. WIKI_TASK_TTL_SECONDS = _env_int("DEEPWIKI_WIKI_TASK_TTL_SECONDS", 300) +# Admission cap: distinct non-terminal tasks allowed in the registry at once. +# Unlike MAX_CONCURRENT_WIKI_TASKS (the execution semaphore), this bounds how +# many tasks can be WAITING for that semaphore, since each submitted repo_key +# is admitted into the registry (and gets its own asyncio.Task) regardless of +# whether a slot is free. Without this, a caller varying repo_url could queue +# unlimited pending work. +WIKI_MAX_PENDING_TASKS = _env_int( + "DEEPWIKI_WIKI_MAX_PENDING_TASKS", MAX_CONCURRENT_WIKI_TASKS * 10 +) + + +class TaskQueueFullError(Exception): + """Raised when the registry already holds WIKI_MAX_PENDING_TASKS non-terminal tasks.""" class WikiTask(BaseModel): @@ -142,10 +155,15 @@ class TaskRegistry: _lock: asyncio.Lock _semaphore: asyncio.Semaphore - def __init__(self, max_concurrent: int = MAX_CONCURRENT_WIKI_TASKS): + def __init__( + self, + max_concurrent: int = MAX_CONCURRENT_WIKI_TASKS, + max_pending: int = WIKI_MAX_PENDING_TASKS, + ): self._tasks = {} self._lock = asyncio.Lock() self._semaphore = asyncio.Semaphore(max_concurrent) + self._max_pending = max_pending def get(self, id: str) -> WikiTask | None: return self._tasks.get(id) @@ -173,6 +191,11 @@ async def submit( joined=True, ) + if len(self.active()) >= self._max_pending: + raise TaskQueueFullError( + f"{self._max_pending} tasks are already pending; try again later" + ) + if wiki_cache_exists( owner=task.request.owner, repo=task.request.repo, diff --git a/tests/backend/routers/test_wiki_auth_enforcement.py b/tests/backend/routers/test_wiki_auth_enforcement.py new file mode 100644 index 000000000..511111a62 --- /dev/null +++ b/tests/backend/routers/test_wiki_auth_enforcement.py @@ -0,0 +1,147 @@ +"""Regression coverage for WIKI_AUTH_MODE enforcement. + +Before this fix, enabling WIKI_AUTH_MODE only protected DELETE +/api/wiki_cache. Every other costly or sensitive route (wiki generation, +repo indexing, chat, codemap, cache reads) performed no check at all, so +an operator who turned on the one documented protection mechanism got no +actual protection anywhere but that single endpoint. +""" + +import pytest +from fastapi.testclient import TestClient + +import api.routers.auth as auth_router +import api.services.wiki.tasks as wt + + +@pytest.fixture(autouse=True) +def _clear_registry(): + wt.registry._tasks.clear() + yield + wt.registry._tasks.clear() + + +@pytest.fixture +def auth_enabled(monkeypatch): + monkeypatch.setattr(auth_router, "WIKI_AUTH_MODE", True) + monkeypatch.setattr(auth_router, "WIKI_AUTH_CODE", "correct-code") + yield + + +def _wiki_task_body(**overrides): + body = {"owner": "o", "repo": "r", "type": "github", "repo_url": "https://github.com/o/r"} + body.update(overrides) + return body + + +def test_submit_wiki_task_rejects_missing_code(auth_enabled): + from api.main import app + + with TestClient(app) as client: + r = client.post("/wiki/tasks", json=_wiki_task_body()) + assert r.status_code == 401 + + +def test_submit_wiki_task_rejects_wrong_code(auth_enabled): + from api.main import app + + with TestClient(app) as client: + r = client.post( + "/wiki/tasks", json=_wiki_task_body(authorization_code="wrong") + ) + assert r.status_code == 401 + + +def test_submit_wiki_task_accepts_correct_code(auth_enabled, monkeypatch): + monkeypatch.setattr(wt, "wiki_cache_exists", lambda *p, **kwargs: False) + monkeypatch.setattr(wt, "repo_index_exist", lambda repo: True) + + async def fake_determine(task): + from api.schemas import WikiStructureModel + + return WikiStructureModel(id="wiki", title="T", description="D", pages=[]) + + monkeypatch.setattr(wt, "_determine_structure", fake_determine) + + from api.main import app + + with TestClient(app) as client: + r = client.post( + "/wiki/tasks", + json=_wiki_task_body(authorization_code="correct-code"), + ) + assert r.status_code == 200, r.text + + +def test_prepare_repo_rejects_missing_code(auth_enabled): + from api.main import app + + with TestClient(app) as client: + r = client.post( + "/repo/prepare", + json={"repo_url": "https://github.com/o/r", "type": "github"}, + ) + assert r.status_code == 401 + + +def test_codemap_file_rejects_missing_code(auth_enabled): + from api.main import app + + with TestClient(app) as client: + r = client.get( + "/codemap/file", + params={"repo_url": "/", "file_path": "x", "type": "github"}, + ) + assert r.status_code == 401 + + +def test_read_wiki_cache_rejects_missing_code(auth_enabled): + from api.main import app + + with TestClient(app) as client: + r = client.get( + "/api/wiki_cache", + params={ + "owner": "o", + "repo": "r", + "repo_type": "github", + "language": "en", + }, + ) + assert r.status_code == 401 + + +def test_delete_wiki_cache_still_enforces_code(auth_enabled): + """The one route that already checked this must keep working.""" + from api.main import app + + with TestClient(app) as client: + r = client.delete( + "/api/wiki_cache", + params={ + "owner": "o", + "repo": "r", + "repo_type": "github", + "language": "en", + }, + ) + assert r.status_code == 401 + + +def test_routes_are_unprotected_when_auth_mode_disabled(monkeypatch): + """Default (WIKI_AUTH_MODE off): no regression for normal operators.""" + monkeypatch.setattr(auth_router, "WIKI_AUTH_MODE", False) + + from api.main import app + + with TestClient(app) as client: + r = client.get( + "/api/wiki_cache", + params={ + "owner": "o", + "repo": "r", + "repo_type": "github", + "language": "en", + }, + ) + assert r.status_code == 200 diff --git a/tests/backend/services/test_wiki_task.py b/tests/backend/services/test_wiki_task.py index 937926d37..23d6c3984 100644 --- a/tests/backend/services/test_wiki_task.py +++ b/tests/backend/services/test_wiki_task.py @@ -5,6 +5,7 @@ import api.services.wiki.tasks as wt from api.schemas import WikiPage, WikiStructureModel, WikiTaskRequest from api.services.wiki.tasks import ( + TaskQueueFullError, TaskRegistry, TaskStatus, WikiTask, @@ -264,3 +265,49 @@ async def test_public_dict_hides_token(): assert "token" not in d assert "SECRET" not in str(d) assert d["name"] == "o/r" and d["status"] == "pending" + + +# --------------------------------------------------------------------------- # +# admission cap (unbounded-task-admission DoS regression) +# --------------------------------------------------------------------------- # +async def test_submit_rejects_once_max_pending_reached(monkeypatch): + """Regression: varying repo_url used to admit an unbounded number of + pending tasks into the registry before any semaphore applied. A small + max_pending must cap distinct non-terminal tasks, not just concurrent + execution.""" + monkeypatch.setattr(wt, "wiki_cache_exists", lambda *p, **kwargs: False) + + async def never_finishes(task): + await asyncio.sleep(10) + return True + + registry = TaskRegistry(max_concurrent=5, max_pending=2) + + await registry.submit(_req(owner="o", repo="r1"), async_func=never_finishes) + await registry.submit(_req(owner="o", repo="r2"), async_func=never_finishes) + + with pytest.raises(TaskQueueFullError): + await registry.submit(_req(owner="o", repo="r3"), async_func=never_finishes) + + for task in list(registry._tasks.values()): + task.task.cancel() + + +async def test_submit_joining_an_existing_task_does_not_count_against_cap(): + """Re-submitting the SAME repo_key (join) must not be blocked by the cap + even when it is already at max_pending, since it does not add a new + registry entry.""" + registry = TaskRegistry(max_concurrent=5, max_pending=1) + + async def never_finishes(task): + await asyncio.sleep(10) + return True + + first = await registry.submit(_req(owner="o", repo="r1"), async_func=never_finishes) + assert first.created is True + + joined = await registry.submit(_req(owner="o", repo="r1"), async_func=never_finishes) + assert joined.joined is True + + for task in list(registry._tasks.values()): + task.task.cancel()