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()