From cc4e68c8472725b452d4170f484a9bb8c6f6c8f0 Mon Sep 17 00:00:00 2001 From: Henry Dowling Date: Sun, 9 Aug 2026 18:05:34 -0400 Subject: [PATCH] Protect VFS execution from concurrent overload --- backend/routers/vfs.py | 4 +- backend/services/vfs_service.py | 74 +++++++++++++++++++++++++------- backend/tests/test_vfs.py | 76 ++++++++++++++++++++++++++++++++- 3 files changed, 137 insertions(+), 17 deletions(-) diff --git a/backend/routers/vfs.py b/backend/routers/vfs.py index 1f776ac5..c46a44d1 100644 --- a/backend/routers/vfs.py +++ b/backend/routers/vfs.py @@ -11,7 +11,7 @@ from ..auth import get_current_user from ..services import security_audit_service, vfs_service -from ..services.vfs_service import VfsBudgetExceeded +from ..services.vfs_service import VfsBudgetExceeded, VfsBusy router = APIRouter(prefix="/api/v1/me/vfs", tags=["vfs"]) @@ -53,6 +53,8 @@ async def run_vfs( raise HTTPException(status_code=400, detail=str(e)) from e except VfsBudgetExceeded as e: raise HTTPException(status_code=413, detail=str(e)) from e + except VfsBusy as e: + raise HTTPException(status_code=429, detail=str(e), headers={"Retry-After": "2"}) from e @router.post("/searches", status_code=204) diff --git a/backend/services/vfs_service.py b/backend/services/vfs_service.py index 06cd17ac..8f8f2368 100644 --- a/backend/services/vfs_service.py +++ b/backend/services/vfs_service.py @@ -17,6 +17,7 @@ import asyncio import functools import threading +import time from contextlib import contextmanager import anyio @@ -32,16 +33,29 @@ # on an open connection until the client's own timeout fires. MAX_DOCUMENT_READS = 400 +# Each active script retains a full VFS model and the documents it reads. This +# cap keeps one caller's burst below the API process's memory limit. +MAX_CONCURRENT_SCRIPTS = 4 + +# Nested ASGI requests must finish inside the script's remaining budget. The +# worker itself stays attached so a timed-out request never frees a slot while +# memory-heavy work is still running. +MAX_SCRIPT_SECONDS = 60 + SOURCE_ENTRIES_PAGE = 1000 class VfsBudgetExceeded(Exception): - """More direct document reads than one shell invocation is allowed. + """More document reads or time than one shell invocation is allowed. Deliberately not a VfsClientError: the shell downgrades those to per-file warnings, and this must abort the whole command. Reads inside a grep sweep raise VfsScanBudget instead, which the shell turns into a partial result.""" +class VfsBusy(Exception): + """This API process is already running its safe number of VFS scripts.""" + + class InProcessVfsClient: """`VfsClient` served by the running app over nested ASGI calls. @@ -56,6 +70,24 @@ def __init__(self, http: httpx.AsyncClient, loop: asyncio.AbstractEventLoop) -> self._scan = False self._document_reads = 0 self._reads_lock = threading.Lock() + self._started = time.monotonic() + + def _result_before_deadline(self, future): + seconds_left = MAX_SCRIPT_SECONDS - (time.monotonic() - self._started) + if seconds_left <= 0: + future.cancel() + raise VfsBudgetExceeded( + f"command ran longer than {MAX_SCRIPT_SECONDS}s; " + "scope it to a smaller subtree or use search" + ) + try: + return future.result(timeout=seconds_left) + except TimeoutError: + future.cancel() + raise VfsBudgetExceeded( + f"command ran longer than {MAX_SCRIPT_SECONDS}s; " + "scope it to a smaller subtree or use search" + ) from None @contextmanager def internal_calls(self): @@ -91,7 +123,7 @@ def record_search(self, pattern: str, roots: list[str], docs_scanned: int) -> No ), self._loop, ) - response = future.result() + response = self._result_before_deadline(future) if response.status_code >= 400: raise VfsClientError(_error_detail(response)) @@ -110,7 +142,7 @@ def _request(self, method: str, endpoint: str, **params) -> httpx.Response: self._http.request(method, endpoint, params=params or None, headers=headers), self._loop, ) - response = future.result() + response = self._result_before_deadline(future) if response.status_code >= 400: raise VfsClientError(_error_detail(response)) return response @@ -224,25 +256,37 @@ def _run_script( } +_scripts_running = 0 + + async def run_vfs_script(app, authorization: str, script: str, cwd: str) -> dict: """Execute one read-only shell script against the caller's Stash. `authorization` is forwarded verbatim onto every nested request, so the VFS sees precisely what that credential sees anywhere else in the API. """ - loop = asyncio.get_running_loop() - transport = httpx.ASGITransport(app=app) - async with httpx.AsyncClient( - transport=transport, - base_url="http://vfs.internal", - # X-Stash-Via tags nested reads as ask-the-stash traffic in the audit - # trail (see auth._set_request_via). - headers={"Authorization": authorization, "X-Stash-Via": "ask"}, - timeout=None, - ) as http: - return await anyio.to_thread.run_sync( - functools.partial(_run_script, http, loop, script, cwd) + global _scripts_running + if _scripts_running >= MAX_CONCURRENT_SCRIPTS: + raise VfsBusy( + f"{MAX_CONCURRENT_SCRIPTS} VFS commands are already running; retry in a moment" ) + _scripts_running += 1 + try: + loop = asyncio.get_running_loop() + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, + base_url="http://vfs.internal", + # X-Stash-Via tags nested reads as ask-the-stash traffic in the audit + # trail (see auth._set_request_via). + headers={"Authorization": authorization, "X-Stash-Via": "ask"}, + timeout=None, + ) as http: + return await anyio.to_thread.run_sync( + functools.partial(_run_script, http, loop, script, cwd) + ) + finally: + _scripts_running -= 1 def _resolve_node(http: httpx.AsyncClient, loop: asyncio.AbstractEventLoop, path: str) -> dict: diff --git a/backend/tests/test_vfs.py b/backend/tests/test_vfs.py index fa23d9b2..3410f79b 100644 --- a/backend/tests/test_vfs.py +++ b/backend/tests/test_vfs.py @@ -7,12 +7,14 @@ not part of the tree. """ +import asyncio +import threading from uuid import UUID import pytest from httpx import AsyncClient -from backend.services import source_service +from backend.services import source_service, vfs_service from .conftest import unique_name @@ -215,6 +217,78 @@ async def test_document_read_budget_still_aborts_direct_reads(client: AsyncClien assert resp.status_code == 413 +async def test_concurrency_cap_rejects_an_overlapping_request(client: AsyncClient, monkeypatch): + """A burst must reject excess work before it builds another VFS model.""" + monkeypatch.setattr("backend.services.vfs_service.MAX_CONCURRENT_SCRIPTS", 1) + started = threading.Event() + release = threading.Event() + + def run_script(*_args): + started.set() + release.wait() + return {"stdout": "", "stderr": "", "exit_code": 0, "cwd": "/"} + + monkeypatch.setattr("backend.services.vfs_service._run_script", run_script) + api_key, _ = await _register(client) + first = asyncio.create_task(_vfs(client, api_key, "ls /")) + + try: + assert await asyncio.to_thread(started.wait, 1) + second = await _vfs(client, api_key, "ls /") + + assert second.status_code == 429 + assert second.headers["retry-after"] == "2" + finally: + release.set() + + assert (await first).status_code == 200 + + +async def test_concurrency_slot_is_released_after_failure(client: AsyncClient, monkeypatch): + """A crashed command must not permanently reduce the process capacity.""" + monkeypatch.setattr("backend.services.vfs_service.MAX_CONCURRENT_SCRIPTS", 1) + calls = 0 + + def run_script(*_args): + nonlocal calls + calls += 1 + if calls == 1: + raise RuntimeError("script crashed") + return {"stdout": "", "stderr": "", "exit_code": 0, "cwd": "/"} + + monkeypatch.setattr("backend.services.vfs_service._run_script", run_script) + api_key, _ = await _register(client) + + first = await _vfs(client, api_key, "ls /") + + assert first.status_code == 500 + assert (await _vfs(client, api_key, "ls /")).status_code == 200 + + +async def test_wall_clock_budget_aborts_the_command(client: AsyncClient, monkeypatch): + monkeypatch.setattr("backend.services.vfs_service.MAX_SCRIPT_SECONDS", 0) + api_key, _ = await _register(client) + + resp = await _vfs(client, api_key, "ls /") + + assert resp.status_code == 413 + assert "longer than" in resp.json()["detail"] + + +async def test_wall_clock_budget_cancels_a_slow_nested_request(monkeypatch): + """A stalled content route must not hold a VFS slot past the deadline.""" + monkeypatch.setattr("backend.services.vfs_service.MAX_SCRIPT_SECONDS", 0.05) + + class SlowHttp: + async def request(self, *_args, **_kwargs): + await asyncio.sleep(10) + + client = vfs_service.InProcessVfsClient(SlowHttp(), asyncio.get_running_loop()) + + with pytest.raises(vfs_service.VfsBudgetExceeded, match="longer than"): + await asyncio.to_thread(client._request, "GET", "/slow") + + async def test_machine_fs_404s_without_provisioned_computer(client: AsyncClient, monkeypatch): """Browsing must never conjure a VM: a user who never ran a cloud agent gets a 404 from the machine fs, not a freshly provisioned sprite."""