diff --git a/strix/runtime/caido_bootstrap.py b/strix/runtime/caido_bootstrap.py index a9c7c82a..80a5815a 100644 --- a/strix/runtime/caido_bootstrap.py +++ b/strix/runtime/caido_bootstrap.py @@ -27,6 +27,8 @@ logger = logging.getLogger(__name__) _LOGIN_AS_GUEST_BODY = ( '{"query":"mutation LoginAsGuest { loginAsGuest { token { accessToken } } }"}' ) +_PROJECT_SETUP_TIMEOUT_MS = 45_000 +_PROJECT_SETUP_ATTEMPTS = 3 async def _login_as_guest( @@ -78,6 +80,55 @@ async def _login_as_guest( raise RuntimeError(f"loginAsGuest failed after {attempts} attempts: {last_err}") +async def _setup_project(host_url: str, access_token: str) -> None: + """Connect with a longer deadline and select the sandbox project.""" + from caido_sdk_client import Client, TokenAuthOptions + from caido_sdk_client.types import CreateProjectOptions + + client = Client( + host_url, + auth=TokenAuthOptions(token=access_token), + timeout_ms=_PROJECT_SETUP_TIMEOUT_MS, + ) + try: + await client.connect() + project = None + last_exc: Exception | None = None + for i in range(1, _PROJECT_SETUP_ATTEMPTS + 1): + try: + if project is None: + project = await client.project.create( + CreateProjectOptions(name="sandbox", temporary=True), + ) + await client.project.select(project.id) + except Exception as exc: # noqa: BLE001 + last_exc = exc + if project is None: + with contextlib.suppress(Exception): + projects = await client.project.list() + sandbox_projects = [item for item in projects if item.name == "sandbox"] + if sandbox_projects: + project = max(sandbox_projects, key=lambda item: item.id) + logger.warning( + "Caido project setup attempt %d/%d failed: %s", + i, + _PROJECT_SETUP_ATTEMPTS, + exc, + ) + if i < _PROJECT_SETUP_ATTEMPTS: + await asyncio.sleep(min(2.0 * i, 8.0)) + else: + logger.info("Caido project selected: %s", project.id) + return + assert last_exc is not None + raise RuntimeError( + f"Caido project setup failed after {_PROJECT_SETUP_ATTEMPTS} attempts" + ) from last_exc + finally: + with contextlib.suppress(Exception): + await client.aclose() + + async def bootstrap_caido( session: BaseSandboxSession, *, @@ -89,27 +140,23 @@ async def bootstrap_caido( # only needed once a sandbox is actually being bootstrapped, so it is # imported here rather than at module scope. from caido_sdk_client import Client, TokenAuthOptions - from caido_sdk_client.types import CreateProjectOptions logger.info("Bootstrapping Caido client (host=%s, container=%s)", host_url, container_url) access_token = await _login_as_guest(session, container_url=container_url) + await _setup_project(host_url, access_token) + client = Client(host_url, auth=TokenAuthOptions(token=access_token)) 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 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() raise - logger.info("Caido project selected: %s", project.id) return client diff --git a/tests/test_caido_bootstrap.py b/tests/test_caido_bootstrap.py index 98e42a94..f4c68c88 100644 --- a/tests/test_caido_bootstrap.py +++ b/tests/test_caido_bootstrap.py @@ -10,6 +10,7 @@ from __future__ import annotations import asyncio import sys import types +from dataclasses import dataclass from typing import Any import pytest @@ -33,32 +34,95 @@ class _FakeSession: return _FakeExecResult('{"data":{"loginAsGuest":{"token":{"accessToken":"t"}}}}') +@dataclass +class _FakeProject: + id: str + name: str = "sandbox" + + +class _FakeProjectSDK: + def __init__( + self, + *, + create_errors: list[BaseException] | None = None, + select_errors: list[BaseException] | None = None, + projects: list[_FakeProject] | None = None, + ) -> None: + self.create_errors = list(create_errors or []) + self.select_errors = list(select_errors or []) + self.projects = projects or [] + self.create_calls = 0 + self.selected_ids: list[str] = [] + self.list_calls = 0 + + async def create(self, _options: Any) -> _FakeProject: + self.create_calls += 1 + if self.create_errors: + raise self.create_errors.pop(0) + return _FakeProject("created") + + async def select(self, project_id: str) -> _FakeProject: + self.selected_ids.append(project_id) + if self.select_errors: + raise self.select_errors.pop(0) + return _FakeProject(project_id) + + async def list(self) -> list[_FakeProject]: + self.list_calls += 1 + return self.projects + + class _FakeClient: - def __init__(self, connect_error: BaseException) -> None: + def __init__( + self, + connect_error: BaseException | None = None, + *, + project: _FakeProjectSDK | None = None, + ) -> None: self.connect_error = connect_error + self.project = project or _FakeProjectSDK() self.closed = False async def connect(self) -> None: - raise self.connect_error + if self.connect_error is not None: + raise self.connect_error async def aclose(self) -> None: self.closed = True +class _FakeClientFactory: + def __init__(self, clients: list[_FakeClient]) -> None: + self.clients = iter(clients) + self.calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] + + def __call__(self, *args: Any, **kwargs: Any) -> _FakeClient: + self.calls.append((args, kwargs)) + return next(self.clients) + + +def _install_sdk( + monkeypatch: pytest.MonkeyPatch, + clients: list[_FakeClient], +) -> _FakeClientFactory: + factory = _FakeClientFactory(clients) + sdk = types.ModuleType("caido_sdk_client") + sdk.Client = factory # 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 **kwargs: kwargs # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "caido_sdk_client", sdk) + monkeypatch.setitem(sys.modules, "caido_sdk_client.types", sdk_types) + return factory + + async def _bootstrap_expecting( monkeypatch: pytest.MonkeyPatch, error: BaseException ) -> _FakeClient: """Run a bootstrap whose ``connect()`` fails with ``error``.""" + setup_client = _FakeClient() 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) + _install_sdk(monkeypatch, [setup_client, client]) with pytest.raises(type(error)): await bootstrap_caido( @@ -66,6 +130,7 @@ async def _bootstrap_expecting( host_url="http://host", container_url="http://container", ) + assert setup_client.closed return client @@ -79,3 +144,130 @@ async def test_cancellation_during_connect_closes_the_client( async def test_failed_connect_closes_the_client(monkeypatch: pytest.MonkeyPatch) -> None: client = await _bootstrap_expecting(monkeypatch, RuntimeError("no listener")) assert client.closed + + +async def test_select_retries_without_creating_another_project( + monkeypatch: pytest.MonkeyPatch, +) -> None: + setup_project = _FakeProjectSDK(select_errors=[RuntimeError("not ready")]) + setup_client = _FakeClient(project=setup_project) + returned_client = _FakeClient() + factory = _install_sdk(monkeypatch, [setup_client, returned_client]) + sleep_calls: list[float] = [] + + async def _sleep(delay: float) -> None: + sleep_calls.append(delay) + + monkeypatch.setattr( + "strix.runtime.caido_bootstrap.asyncio.sleep", + _sleep, + ) + + result = await bootstrap_caido( + _FakeSession(), # type: ignore[arg-type] + host_url="http://host", + container_url="http://container", + ) + + assert result is returned_client + assert setup_project.create_calls == 1 + assert setup_project.selected_ids == ["created", "created"] + assert setup_client.closed + assert factory.calls[0][1]["timeout_ms"] == 45_000 + assert factory.calls[1][1].get("timeout_ms") is None + assert sleep_calls == [2.0] + + +async def test_create_failure_reuses_the_most_recent_sandbox_project( + monkeypatch: pytest.MonkeyPatch, +) -> None: + setup_project = _FakeProjectSDK( + create_errors=[RuntimeError("create timed out")], + projects=[ + _FakeProject("project-1"), + _FakeProject("project-2"), + _FakeProject("other", name="other"), + ], + ) + setup_client = _FakeClient(project=setup_project) + returned_client = _FakeClient() + _install_sdk(monkeypatch, [setup_client, returned_client]) + + async def _sleep(_delay: float) -> None: + pass + + monkeypatch.setattr("strix.runtime.caido_bootstrap.asyncio.sleep", _sleep) + + result = await bootstrap_caido( + _FakeSession(), # type: ignore[arg-type] + host_url="http://host", + container_url="http://container", + ) + + assert result is returned_client + assert setup_project.create_calls == 1 + assert setup_project.list_calls == 1 + assert setup_project.selected_ids == ["project-2"] + assert setup_client.closed + + +async def test_project_setup_failure_chains_last_error_and_closes_setup_client( + monkeypatch: pytest.MonkeyPatch, +) -> None: + errors: list[BaseException] = [ + RuntimeError("first"), + RuntimeError("second"), + RuntimeError("last"), + ] + setup_project = _FakeProjectSDK(select_errors=errors) + setup_client = _FakeClient(project=setup_project) + _install_sdk(monkeypatch, [setup_client]) + + async def _sleep(_delay: float) -> None: + pass + + monkeypatch.setattr("strix.runtime.caido_bootstrap.asyncio.sleep", _sleep) + + with pytest.raises( + RuntimeError, + match="Caido project setup failed after 3 attempts", + ) as exc_info: + await bootstrap_caido( + _FakeSession(), # type: ignore[arg-type] + host_url="http://host", + container_url="http://container", + ) + + assert exc_info.value.__cause__ is errors[-1] + assert setup_client.closed + assert setup_project.create_calls == 1 + assert setup_project.selected_ids == ["created", "created", "created"] + + +async def test_cancelled_project_setup_is_not_retried( + monkeypatch: pytest.MonkeyPatch, +) -> None: + setup_project = _FakeProjectSDK(select_errors=[asyncio.CancelledError()]) + setup_client = _FakeClient(project=setup_project) + _install_sdk(monkeypatch, [setup_client]) + sleep_calls: list[float] = [] + + async def _sleep(delay: float) -> None: + sleep_calls.append(delay) + + monkeypatch.setattr( + "strix.runtime.caido_bootstrap.asyncio.sleep", + _sleep, + ) + + with pytest.raises(asyncio.CancelledError): + await bootstrap_caido( + _FakeSession(), # type: ignore[arg-type] + host_url="http://host", + container_url="http://container", + ) + + assert setup_project.create_calls == 1 + assert setup_project.selected_ids == ["created"] + assert sleep_calls == [] + assert setup_client.closed