diff --git a/strix/runtime/caido_bootstrap.py b/strix/runtime/caido_bootstrap.py index a9c7c82ae..c1b96c7f9 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 +_BOOTSTRAP_ATTEMPTS = 3 async def _login_as_guest( @@ -78,6 +80,69 @@ async def _login_as_guest( raise RuntimeError(f"loginAsGuest failed after {attempts} attempts: {last_err}") +async def _find_sandbox_project(client: Client) -> str | None: + """Look for a project a create that timed out client-side may have left behind.""" + with contextlib.suppress(Exception): + projects = [item for item in await client.project.list() if item.name == "sandbox"] + if projects: + return str(max(projects, key=lambda item: item.id).id) + return None + + +async def _setup_project(host_url: str, access_token: str) -> None: + """Select the sandbox project, retrying the whole connect/create/select sequence. + + Each attempt gets a fresh client with a deadline well past the SDK default: + these mutations are slow on a cold Caido, while the long-lived client the + scan uses keeps the short default so a traffic poll cannot stall on it. + Until a project is selected Caido answers every proxied request with a 500, + so giving up here costs the whole run, not just the traffic capture. + """ + from caido_sdk_client import Client, TokenAuthOptions + from caido_sdk_client.types import CreateProjectOptions + + project_id: str | None = None + last_exc: Exception | None = None + for attempt in range(1, _BOOTSTRAP_ATTEMPTS + 1): + client = Client( + host_url, + auth=TokenAuthOptions(token=access_token), + timeout_ms=_PROJECT_SETUP_TIMEOUT_MS, + ) + try: + await client.connect() + if project_id is None: + try: + created = await client.project.create( + CreateProjectOptions(name="sandbox", temporary=True), + ) + except Exception: + # A create that timed out client-side may still have landed. + project_id = await _find_sandbox_project(client) + raise + project_id = created.id + await client.project.select(project_id) + except Exception as exc: # noqa: BLE001 + last_exc = exc + logger.warning( + "Caido project setup attempt %d/%d failed: %s", + attempt, + _BOOTSTRAP_ATTEMPTS, + exc, + ) + if attempt < _BOOTSTRAP_ATTEMPTS: + await asyncio.sleep(min(2.0 * attempt, 8.0)) + else: + logger.info("Caido project selected: %s", project_id) + return + finally: + with contextlib.suppress(Exception): + await client.aclose() + raise RuntimeError( + f"Caido project setup failed after {_BOOTSTRAP_ATTEMPTS} attempts" + ) from last_exc + + async def bootstrap_caido( session: BaseSandboxSession, *, @@ -89,27 +154,39 @@ 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) - 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 + await _setup_project(host_url, access_token) + + last_exc: Exception | None = None + for attempt in range(1, _BOOTSTRAP_ATTEMPTS + 1): + client = Client(host_url, auth=TokenAuthOptions(token=access_token)) + try: + # A cancellation while connecting can leave a half-connected + # transport behind, so close the client before propagating it. + await client.connect() + except Exception as exc: # noqa: BLE001 + with contextlib.suppress(Exception): + await client.aclose() + last_exc = exc + logger.warning( + "Caido client connect attempt %d/%d failed: %s", + attempt, + _BOOTSTRAP_ATTEMPTS, + exc, + ) + if attempt < _BOOTSTRAP_ATTEMPTS: + await asyncio.sleep(min(2.0 * attempt, 8.0)) + except BaseException: + # Teardown can cancel the bootstrap at any await; do not retry. + with contextlib.suppress(Exception): + await client.aclose() + raise + else: + return client + raise RuntimeError( + f"Caido client connect failed after {_BOOTSTRAP_ATTEMPTS} attempts" + ) from last_exc diff --git a/tests/test_caido_bootstrap.py b/tests/test_caido_bootstrap.py index 98e42a94d..655ca3cef 100644 --- a/tests/test_caido_bootstrap.py +++ b/tests/test_caido_bootstrap.py @@ -10,13 +10,18 @@ from __future__ import annotations import asyncio import sys import types -from typing import Any +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any import pytest from strix.runtime.caido_bootstrap import bootstrap_caido +if TYPE_CHECKING: + from collections.abc import Sequence + + class _FakeExecResult: stderr = b"" exit_code = 0 @@ -33,49 +38,305 @@ 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 -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. +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 = lambda *_a, **_k: client # type: ignore[attr-defined] + 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 **_k: None # type: ignore[attr-defined] + 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 - with pytest.raises(type(error)): + +def _setup_clients( + project: _FakeProjectSDK, + count: int, + *, + connect_errors: list[BaseException | None] | None = None, +) -> list[_FakeClient]: + """One client per setup attempt, all sharing the same server-side project state.""" + errors: list[BaseException | None] = list(connect_errors or []) + errors += [None] * (count - len(errors)) + return [_FakeClient(errors[i], project=project) for i in range(count)] + + +async def _bootstrap_expecting( + monkeypatch: pytest.MonkeyPatch, errors: Sequence[BaseException] +) -> tuple[list[_FakeClient], list[float], BaseException]: + """Run a bootstrap whose scan-client connections fail.""" + setup_client = _FakeClient() + scan_clients = [_FakeClient(error) for error in errors] + _install_sdk(monkeypatch, [setup_client, *scan_clients]) + 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(BaseException) as exc_info: await bootstrap_caido( _FakeSession(), # type: ignore[arg-type] host_url="http://host", container_url="http://container", ) - return client + assert setup_client.closed + return [setup_client, *scan_clients], sleep_calls, exc_info.value async def test_cancellation_during_connect_closes_the_client( monkeypatch: pytest.MonkeyPatch, ) -> None: - client = await _bootstrap_expecting(monkeypatch, asyncio.CancelledError()) - assert client.closed + clients, sleep_calls, error = await _bootstrap_expecting( + monkeypatch, + [asyncio.CancelledError()], + ) + assert isinstance(error, asyncio.CancelledError) + assert len(clients) == 2 + assert all(client.closed for client in clients) + assert sleep_calls == [] async def test_failed_connect_closes_the_client(monkeypatch: pytest.MonkeyPatch) -> None: - client = await _bootstrap_expecting(monkeypatch, RuntimeError("no listener")) - assert client.closed + errors = [ + RuntimeError("first"), + RuntimeError("second"), + RuntimeError("last"), + ] + clients, sleep_calls, error = await _bootstrap_expecting(monkeypatch, errors) + assert isinstance(error, RuntimeError) + assert str(error) == "Caido client connect failed after 3 attempts" + assert error.__cause__ is errors[-1] + assert len(clients) == 4 + assert all(client.closed for client in clients) + assert sleep_calls == [2.0, 4.0] + + +async def test_select_retries_without_creating_another_project( + monkeypatch: pytest.MonkeyPatch, +) -> None: + setup_project = _FakeProjectSDK(select_errors=[RuntimeError("not ready")]) + setup_clients = _setup_clients(setup_project, 2) + returned_client = _FakeClient() + factory = _install_sdk(monkeypatch, [*setup_clients, 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 all(client.closed for client in setup_clients) + assert factory.calls[0][1]["timeout_ms"] == 45_000 + assert factory.calls[1][1]["timeout_ms"] == 45_000 + assert factory.calls[2][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_clients = _setup_clients(setup_project, 2) + returned_client = _FakeClient() + _install_sdk(monkeypatch, [*setup_clients, 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 all(client.closed for client in setup_clients) + + +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_clients = _setup_clients(setup_project, 3) + _install_sdk(monkeypatch, setup_clients) + + 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 all(client.closed for client in setup_clients) + 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_clients = _setup_clients(setup_project, 1) + _install_sdk(monkeypatch, setup_clients) + 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 all(client.closed for client in setup_clients) + + +async def test_setup_connect_failure_is_retried(monkeypatch: pytest.MonkeyPatch) -> None: + setup_project = _FakeProjectSDK() + setup_clients = _setup_clients( + setup_project, + 2, + connect_errors=[RuntimeError("gateway not ready")], + ) + returned_client = _FakeClient() + _install_sdk(monkeypatch, [*setup_clients, 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"] + assert setup_project.list_calls == 0 + assert sleep_calls == [2.0] + assert all(client.closed for client in setup_clients)