diff --git a/strix/runtime/caido_bootstrap.py b/strix/runtime/caido_bootstrap.py index 8398350b..a9c7c82a 100644 --- a/strix/runtime/caido_bootstrap.py +++ b/strix/runtime/caido_bootstrap.py @@ -96,15 +96,17 @@ async def bootstrap_caido( access_token = await _login_as_guest(session, container_url=container_url) client = Client(host_url, auth=TokenAuthOptions(token=access_token)) - await client.connect() - try: + # connect() is inside the guard as well: a cancellation there (scan + # teardown while the bootstrap is still in flight) would otherwise + # leave the half-connected transport behind. + await client.connect() project = await client.project.create( CreateProjectOptions(name="sandbox", temporary=True), ) await client.project.select(project.id) except BaseException: - # The connected client never reaches the session bundle if project + # The client never reaches the session bundle if connect or project # setup fails, so close it here to avoid leaking the transport. with contextlib.suppress(Exception): await client.aclose() diff --git a/strix/runtime/caido_handle.py b/strix/runtime/caido_handle.py new file mode 100644 index 00000000..5b1d1c74 --- /dev/null +++ b/strix/runtime/caido_handle.py @@ -0,0 +1,60 @@ +"""Handle for a Caido bootstrap running concurrently with the scan start. + +The Caido sidecar login + project setup costs a couple of seconds of +guest-side polling, and nothing needs the client until the first proxy +tool call (or the first traffic poll). :class:`CaidoBootstrapHandle` +wraps the in-flight bootstrap task so session bring-up can return as +soon as the container is up; consumers resolve the client at first use. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +from typing import TYPE_CHECKING + + +if TYPE_CHECKING: + from caido_sdk_client import Client + + +logger = logging.getLogger(__name__) + + +class CaidoBootstrapHandle: + """Resolves to the connected Caido client once the bootstrap finishes. + + A failed bootstrap is surfaced (once) to every ``get()`` caller as the + original exception; proxy tools degrade to their "client unavailable" + result instead of the failure killing the scan at bring-up. + """ + + def __init__(self, task: asyncio.Task[Client]) -> None: + self._task = task + + async def get(self) -> Client: + """Wait for the bootstrap and return the client. + + Shielded so one caller's cancellation (e.g. a tool timeout) does not + cancel the shared bootstrap for everyone else. + """ + return await asyncio.shield(self._task) + + def peek(self) -> Client | None: + """Return the client if the bootstrap already finished cleanly.""" + if self._task.done() and not self._task.cancelled() and self._task.exception() is None: + return self._task.result() + return None + + async def aclose(self) -> None: + """Cancel an in-flight bootstrap or close the finished client.""" + if not self._task.done(): + self._task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await self._task + return + client = self.peek() + if client is not None: + with contextlib.suppress(Exception): + await client.aclose() diff --git a/strix/runtime/session_manager.py b/strix/runtime/session_manager.py index 62204385..e8b3a279 100644 --- a/strix/runtime/session_manager.py +++ b/strix/runtime/session_manager.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import logging import os import sys @@ -15,6 +16,7 @@ from strix.config import load_settings from strix.core.paths import run_dir_for, runtime_state_dir from strix.runtime.backends import backend_supports_bind_mounts, get_backend from strix.runtime.caido_bootstrap import bootstrap_caido +from strix.runtime.caido_handle import CaidoBootstrapHandle if TYPE_CHECKING: @@ -333,10 +335,19 @@ async def create_or_reuse( host_caido_url = f"{scheme}://{caido_endpoint.host}:{caido_endpoint.port}" logger.debug("Caido host endpoint resolved: %s", host_caido_url) - caido_client = await bootstrap_caido( - session, - host_url=host_caido_url, - container_url=container_caido_url, + # The Caido login + project setup polls the guest for a couple of seconds + # and nothing needs the client before the first proxy tool call, so it + # runs concurrently with the rest of scan start; consumers resolve the + # handle at first use (see CaidoBootstrapHandle). + caido_client = CaidoBootstrapHandle( + asyncio.create_task( + bootstrap_caido( + session, + host_url=host_caido_url, + container_url=container_caido_url, + ), + name=f"caido-bootstrap-{scan_id}", + ) ) bundle = { diff --git a/strix/tools/proxy/tools.py b/strix/tools/proxy/tools.py index 091c489f..fabcc7ff 100644 --- a/strix/tools/proxy/tools.py +++ b/strix/tools/proxy/tools.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Literal from agents import RunContextWrapper, function_tool +from strix.runtime.caido_handle import CaidoBootstrapHandle from strix.tools.proxy import caido_api @@ -47,9 +48,16 @@ ScopeAction = Literal["get", "list", "create", "update", "delete"] _CAIDO_CALL_LOCK = asyncio.Lock() -def _ctx_client(ctx: RunContextWrapper) -> Client | None: - inner = ctx.context if isinstance(ctx.context, dict) else {} - return inner.get("caido_client") +async def _ctx_client(ctx: RunContextWrapper) -> Client | None: + inner: dict[str, Any] = ctx.context if isinstance(ctx.context, dict) else {} + client: Client | CaidoBootstrapHandle | None = inner.get("caido_client") + if isinstance(client, CaidoBootstrapHandle): + try: + return await client.get() + except Exception: # noqa: BLE001 + logger.warning("Caido bootstrap failed; proxy tools unavailable", exc_info=True) + return None + return client async def _call[T](client: Client, fn: Callable[[Client], Awaitable[T]]) -> T: @@ -155,7 +163,7 @@ async def list_requests( sort_order: ``asc`` or ``desc``. scope_id: Restrict to a Caido scope (managed via ``scope_rules``). """ - client = _ctx_client(ctx) + client = await _ctx_client(ctx) if client is None: return _no_client() @@ -261,7 +269,7 @@ async def view_request( page: 1-indexed page number (only when no ``search_pattern``). page_size: Lines per page. """ - client = _ctx_client(ctx) + client = await _ctx_client(ctx) if client is None: return _no_client() @@ -379,7 +387,7 @@ async def repeat_request( - ``body`` — replace the body string entirely. - ``cookies`` — dict of cookies to add/update. """ - client = _ctx_client(ctx) + client = await _ctx_client(ctx) if client is None: return _no_client() mods = modifications or {} @@ -461,7 +469,7 @@ async def list_sitemap( (recursive subtree). Only meaningful with ``parent_id``. page: 1-indexed page (30 entries per page). """ - client = _ctx_client(ctx) + client = await _ctx_client(ctx) if client is None: return _no_client() try: @@ -495,7 +503,7 @@ async def view_sitemap_entry( Args: entry_id: ID from ``list_sitemap`` (or any nested entry). """ - client = _ctx_client(ctx) + client = await _ctx_client(ctx) if client is None: return _no_client() try: @@ -554,7 +562,7 @@ async def scope_rules( scope_id: Required for ``get`` / ``update`` / ``delete``. scope_name: Required for ``create`` / ``update``. """ - client = _ctx_client(ctx) + client = await _ctx_client(ctx) if client is None: return _no_client() diff --git a/tests/test_caido_bootstrap.py b/tests/test_caido_bootstrap.py new file mode 100644 index 00000000..98e42a94 --- /dev/null +++ b/tests/test_caido_bootstrap.py @@ -0,0 +1,81 @@ +"""A bootstrap that dies mid-setup must not leave its transport behind. + +The bootstrap now runs concurrently with the scan start, so teardown can +cancel it at any await — including inside ``Client.connect()``, where the +client exists but no caller will ever see it to close it. +""" + +from __future__ import annotations + +import asyncio +import sys +import types +from typing import Any + +import pytest + +from strix.runtime.caido_bootstrap import bootstrap_caido + + +class _FakeExecResult: + stderr = b"" + exit_code = 0 + + def __init__(self, stdout: str) -> None: + self.stdout = stdout + + def ok(self) -> bool: + return True + + +class _FakeSession: + async def exec(self, *_args: Any, **_kwargs: Any) -> _FakeExecResult: + return _FakeExecResult('{"data":{"loginAsGuest":{"token":{"accessToken":"t"}}}}') + + +class _FakeClient: + def __init__(self, connect_error: BaseException) -> None: + self.connect_error = connect_error + self.closed = False + + async def connect(self) -> None: + raise self.connect_error + + async def aclose(self) -> None: + self.closed = True + + +async def _bootstrap_expecting( + monkeypatch: pytest.MonkeyPatch, error: BaseException +) -> _FakeClient: + """Run a bootstrap whose ``connect()`` fails with ``error``.""" + client = _FakeClient(error) + # The SDK is imported inside bootstrap_caido (it is slow to import), so the + # fakes are injected as the modules it imports. + sdk = types.ModuleType("caido_sdk_client") + sdk.Client = lambda *_a, **_k: client # type: ignore[attr-defined] + sdk.TokenAuthOptions = lambda token: token # type: ignore[attr-defined] + sdk_types = types.ModuleType("caido_sdk_client.types") + sdk_types.CreateProjectOptions = lambda **_k: None # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "caido_sdk_client", sdk) + monkeypatch.setitem(sys.modules, "caido_sdk_client.types", sdk_types) + + with pytest.raises(type(error)): + await bootstrap_caido( + _FakeSession(), # type: ignore[arg-type] + host_url="http://host", + container_url="http://container", + ) + return client + + +async def test_cancellation_during_connect_closes_the_client( + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = await _bootstrap_expecting(monkeypatch, asyncio.CancelledError()) + assert client.closed + + +async def test_failed_connect_closes_the_client(monkeypatch: pytest.MonkeyPatch) -> None: + client = await _bootstrap_expecting(monkeypatch, RuntimeError("no listener")) + assert client.closed diff --git a/tests/test_caido_handle.py b/tests/test_caido_handle.py new file mode 100644 index 00000000..84d8ce8a --- /dev/null +++ b/tests/test_caido_handle.py @@ -0,0 +1,103 @@ +"""Tests for the concurrent Caido bootstrap handle.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest + +from strix.runtime.caido_handle import CaidoBootstrapHandle + + +class _FakeClient: + def __init__(self) -> None: + self.closed = False + + async def aclose(self) -> None: + self.closed = True + + +def _handle(coro: Any) -> CaidoBootstrapHandle: + return CaidoBootstrapHandle(asyncio.ensure_future(coro)) + + +async def test_get_waits_for_the_bootstrap() -> None: + client = _FakeClient() + started = asyncio.Event() + + async def _bootstrap() -> Any: + started.set() + await asyncio.sleep(0.01) + return client + + handle = _handle(_bootstrap()) + await started.wait() + assert handle.peek() is None + assert await handle.get() is client + assert handle.peek() is client + + +async def test_get_reraises_bootstrap_failure_to_every_caller() -> None: + async def _bootstrap() -> Any: + raise RuntimeError("caido never came up") + + handle = _handle(_bootstrap()) + for _ in range(2): + with pytest.raises(RuntimeError, match="caido never came up"): + await handle.get() + assert handle.peek() is None + + +async def test_caller_cancellation_does_not_cancel_the_shared_bootstrap() -> None: + client = _FakeClient() + + async def _bootstrap() -> Any: + await asyncio.sleep(0.05) + return client + + handle = _handle(_bootstrap()) + + with pytest.raises(TimeoutError): + await asyncio.wait_for(handle.get(), timeout=0.01) + + assert await handle.get() is client + + +async def test_aclose_closes_a_finished_client() -> None: + client = _FakeClient() + + async def _bootstrap() -> Any: + return client + + handle = _handle(_bootstrap()) + await handle.get() + await handle.aclose() + assert client.closed is True + + +async def test_aclose_cancels_an_in_flight_bootstrap() -> None: + cancelled = asyncio.Event() + + async def _bootstrap() -> Any: + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + cancelled.set() + raise + return _FakeClient() + + handle = _handle(_bootstrap()) + await asyncio.sleep(0) + await handle.aclose() + assert cancelled.is_set() + + +async def test_aclose_swallows_a_failed_bootstrap() -> None: + async def _bootstrap() -> Any: + raise RuntimeError("boom") + + handle = _handle(_bootstrap()) + with pytest.raises(RuntimeError, match="boom"): + await handle.get() + await handle.aclose() diff --git a/tests/test_proxy_client.py b/tests/test_proxy_client.py index da54af66..1a3aa849 100644 --- a/tests/test_proxy_client.py +++ b/tests/test_proxy_client.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, cast import pytest +from strix.runtime.caido_handle import CaidoBootstrapHandle from strix.tools.proxy import caido_api, tools @@ -198,12 +199,31 @@ class _Ctx: self.context = context -def test_ctx_client_returns_client_when_present() -> None: +async def test_ctx_client_returns_client_when_present() -> None: client = _FakeClient("host") - got = tools._ctx_client(cast("Any", _Ctx({"caido_client": client}))) + got = await tools._ctx_client(cast("Any", _Ctx({"caido_client": client}))) assert got is client -def test_ctx_client_returns_none_without_client() -> None: - assert tools._ctx_client(cast("Any", _Ctx({}))) is None - assert tools._ctx_client(cast("Any", _Ctx(None))) is None +async def test_ctx_client_returns_none_without_client() -> None: + assert await tools._ctx_client(cast("Any", _Ctx({}))) is None + assert await tools._ctx_client(cast("Any", _Ctx(None))) is None + + +async def test_ctx_client_resolves_bootstrap_handle() -> None: + client = _FakeClient("host") + + async def _bootstrap() -> Any: + return client + + handle = CaidoBootstrapHandle(asyncio.ensure_future(_bootstrap())) + got = await tools._ctx_client(cast("Any", _Ctx({"caido_client": handle}))) + assert got is client + + +async def test_ctx_client_degrades_when_bootstrap_failed() -> None: + async def _bootstrap() -> Any: + raise RuntimeError("caido never came up") + + handle = CaidoBootstrapHandle(asyncio.ensure_future(_bootstrap())) + assert await tools._ctx_client(cast("Any", _Ctx({"caido_client": handle}))) is None