mirror of
https://github.com/usestrix/strix.git
synced 2026-08-25 20:32:38 +02:00
revert(proxy): drop overfit Caido reconnect/HTTPQL band-aids, keep serialization lock (#799)
Co-authored-by: Ahmed Allam <ahmed39652003@gmail.com>
This commit is contained in:
co-authored by
Ahmed Allam
parent
e4548cb28c
commit
96ca7e544d
@@ -3,9 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import urllib.request
|
||||
@@ -28,9 +26,6 @@ if TYPE_CHECKING:
|
||||
from caido_sdk_client import Client as CaidoClient
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
RequestPart = Literal["request", "response"]
|
||||
SortBy = Literal[
|
||||
"timestamp",
|
||||
@@ -50,18 +45,6 @@ _SITEMAP_PAGE_SIZE = 30
|
||||
_DEFAULT_CAIDO_URL = "http://127.0.0.1:48080"
|
||||
_CLIENT_CACHE: dict[str, Client] = {}
|
||||
_CLIENT_LOCK = asyncio.Lock()
|
||||
|
||||
# Substrings that mean the shared client's transport has died or is being used
|
||||
# concurrently — recoverable by rebuilding the client and retrying once.
|
||||
_CONNECTION_ERROR_MARKERS = (
|
||||
"transport is already connected",
|
||||
"connector is closed",
|
||||
"server disconnected",
|
||||
"session is closed",
|
||||
"cannot write to closing transport",
|
||||
"connection reset",
|
||||
"connection closed",
|
||||
)
|
||||
_REQ_FIELD_MAP: dict[SortBy, tuple[str, str]] = {
|
||||
"timestamp": ("req", "created_at"),
|
||||
"host": ("req", "host"),
|
||||
@@ -108,22 +91,6 @@ async def _new_client() -> Client:
|
||||
return client
|
||||
|
||||
|
||||
async def _safe_aclose(client: Client | None) -> None:
|
||||
"""Close a (possibly dead) client without letting teardown errors escape."""
|
||||
if client is None:
|
||||
return
|
||||
with contextlib.suppress(Exception):
|
||||
await client.aclose()
|
||||
|
||||
|
||||
def _is_connection_error(exc: BaseException) -> bool:
|
||||
message = str(exc).lower()
|
||||
if any(marker in message for marker in _CONNECTION_ERROR_MARKERS):
|
||||
return True
|
||||
cause = exc.__cause__ or exc.__context__
|
||||
return cause is not None and cause is not exc and _is_connection_error(cause)
|
||||
|
||||
|
||||
async def get_client() -> Client:
|
||||
"""Return the shared Caido client, creating it under a lock if needed.
|
||||
|
||||
@@ -139,73 +106,19 @@ async def get_client() -> Client:
|
||||
return client
|
||||
|
||||
|
||||
async def call_with_client[T](
|
||||
fn: Callable[[Client], Awaitable[T]], *, idempotent: bool = True
|
||||
) -> T:
|
||||
"""Run ``fn`` against the shared client, serialized and reconnect-safe.
|
||||
async def call_with_client[T](fn: Callable[[Client], Awaitable[T]]) -> T:
|
||||
"""Run ``fn`` against the shared client, serialized through ``_CLIENT_LOCK``.
|
||||
|
||||
The Caido GraphQL transport is not safe for concurrent use: two in-flight
|
||||
requests race and raise "Transport is already connected". All proxy calls
|
||||
are therefore serialized through ``_CLIENT_LOCK``. If the cached client's
|
||||
transport has since died ("Connector is closed" / "Server disconnected"),
|
||||
the stale client is closed and rebuilt so subsequent calls stop failing
|
||||
against a dead client.
|
||||
|
||||
``fn`` is only re-run automatically when ``idempotent`` is true. For
|
||||
mutations (replay, scope create/update/delete) a connection error may
|
||||
arrive *after* Caido applied the change, so we heal the client for future
|
||||
calls but re-raise instead of risking a double-apply.
|
||||
requests race and raise "Transport is already connected". Serializing every
|
||||
proxy call through the lock prevents that.
|
||||
"""
|
||||
async with _CLIENT_LOCK:
|
||||
client = _CLIENT_CACHE.get("default")
|
||||
if client is None:
|
||||
client = await _new_client()
|
||||
_CLIENT_CACHE["default"] = client
|
||||
try:
|
||||
return await fn(client)
|
||||
except Exception as exc:
|
||||
if not _is_connection_error(exc):
|
||||
raise
|
||||
new_client = await _new_client()
|
||||
_CLIENT_CACHE["default"] = new_client
|
||||
await _safe_aclose(client)
|
||||
if not idempotent:
|
||||
raise
|
||||
return await fn(new_client)
|
||||
|
||||
|
||||
class SharedCaidoClient:
|
||||
"""Serialized, reconnect-safe wrapper around one host-side Caido client.
|
||||
|
||||
Every agent in a scan shares a single instance (propagated through the
|
||||
shallow-copied run context). ``call`` serializes access — the SDK transport
|
||||
is not concurrency-safe — and, when the transport dies, rebuilds the client
|
||||
via ``reconnect`` (which preserves the Caido project) and closes the dead
|
||||
one, so a transient Caido restart no longer disables proxy tools for the
|
||||
rest of the scan.
|
||||
"""
|
||||
|
||||
def __init__(self, client: Client, reconnect: Callable[[], Awaitable[Client]]) -> None:
|
||||
self._client = client
|
||||
self._reconnect = reconnect
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def call[T](self, fn: Callable[[Client], Awaitable[T]], *, idempotent: bool = True) -> T:
|
||||
async with self._lock:
|
||||
try:
|
||||
return await fn(self._client)
|
||||
except Exception as exc:
|
||||
if not _is_connection_error(exc):
|
||||
raise
|
||||
dead, self._client = self._client, await self._reconnect()
|
||||
await _safe_aclose(dead)
|
||||
if not idempotent:
|
||||
raise
|
||||
return await fn(self._client)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
async with self._lock:
|
||||
await _safe_aclose(self._client)
|
||||
return await fn(client)
|
||||
|
||||
|
||||
async def close_client() -> None:
|
||||
@@ -546,9 +459,7 @@ async def repeat_request(
|
||||
)
|
||||
return await replay_send_raw(client, raw=raw, connection=connection)
|
||||
|
||||
# A replay mutates server state; don't auto-retry if the transport dies
|
||||
# mid-send (the request may already have been sent).
|
||||
return await call_with_client(_run, idempotent=False)
|
||||
return await call_with_client(_run)
|
||||
|
||||
|
||||
async def scope_rules(
|
||||
@@ -569,8 +480,7 @@ async def scope_rules(
|
||||
scope_name=scope_name,
|
||||
)
|
||||
|
||||
# get/list are read-only and safe to retry; create/update/delete mutate.
|
||||
return await call_with_client(_run, idempotent=action in {"get", "list"})
|
||||
return await call_with_client(_run)
|
||||
|
||||
|
||||
async def _scope_rules_with_client(
|
||||
@@ -819,11 +729,9 @@ async def view_sitemap_entry(entry_id: str) -> dict[str, Any]:
|
||||
__all__ = [
|
||||
"RequestPart",
|
||||
"ScopeAction",
|
||||
"SharedCaidoClient",
|
||||
"SitemapDepth",
|
||||
"SortBy",
|
||||
"SortOrder",
|
||||
"call_with_client",
|
||||
"close_client",
|
||||
"get_client",
|
||||
"list_requests",
|
||||
|
||||
+48
-77
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import dataclasses
|
||||
import json
|
||||
import logging
|
||||
@@ -13,13 +14,14 @@ from typing import TYPE_CHECKING, Any, Literal
|
||||
from agents import RunContextWrapper, function_tool
|
||||
|
||||
from strix.tools.proxy import caido_api
|
||||
from strix.tools.proxy.caido_api import SharedCaidoClient
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
from caido_sdk_client import Client
|
||||
|
||||
from strix.tools.proxy.caido_api import (
|
||||
@@ -29,7 +31,7 @@ if TYPE_CHECKING:
|
||||
SortOrder,
|
||||
)
|
||||
else:
|
||||
from strix.tools.proxy.caido_api import (
|
||||
from strix.tools.proxy.caido_api import ( # noqa: TC001
|
||||
RequestPart,
|
||||
SitemapDepth,
|
||||
SortBy,
|
||||
@@ -39,19 +41,21 @@ else:
|
||||
|
||||
ScopeAction = Literal["get", "list", "create", "update", "delete"]
|
||||
|
||||
# All agents in a scan share one host-side Caido client whose GraphQL transport
|
||||
# is not concurrency-safe (parallel calls raise "Transport is already
|
||||
# connected"). Serialize every host-side proxy call through this lock.
|
||||
_CAIDO_CALL_LOCK = asyncio.Lock()
|
||||
|
||||
def _ctx_proxy(ctx: RunContextWrapper) -> SharedCaidoClient | None:
|
||||
"""Return the scan-wide serialized, reconnect-safe Caido client holder.
|
||||
|
||||
All agents in a scan share one :class:`SharedCaidoClient` whose GraphQL
|
||||
transport is not concurrency-safe (parallel calls raise "Transport is
|
||||
already connected"). ``SharedCaidoClient.call`` serializes access and
|
||||
rebuilds the transport if it dies mid-scan. Returns ``None`` when no holder
|
||||
is present (e.g. standalone tool invocation outside a scan run).
|
||||
"""
|
||||
def _ctx_client(ctx: RunContextWrapper) -> Client | None:
|
||||
inner = ctx.context if isinstance(ctx.context, dict) else {}
|
||||
proxy = inner.get("caido_client")
|
||||
return proxy if isinstance(proxy, SharedCaidoClient) else None
|
||||
return inner.get("caido_client")
|
||||
|
||||
|
||||
async def _call[T](client: Client, fn: Callable[[Client], Awaitable[T]]) -> T:
|
||||
"""Run ``fn`` against the shared client, serialized under ``_CAIDO_CALL_LOCK``."""
|
||||
async with _CAIDO_CALL_LOCK:
|
||||
return await fn(client)
|
||||
|
||||
|
||||
def _to_tool_json(value: Any) -> Any:
|
||||
@@ -93,39 +97,6 @@ def _err(name: str, exc: Exception) -> str:
|
||||
)
|
||||
|
||||
|
||||
_HTTPQL_HINT = (
|
||||
"HTTPQL syntax: quote string values and leave integers unquoted; combine "
|
||||
"terms with AND / OR (there is no NOT). Numeric fields (resp.code, req.port, "
|
||||
"id, roundtrip) use eq/ne/gt/gte/lt/lte; text/byte fields (req.host, req.path, "
|
||||
"req.method, req.raw, resp.raw) use cont/ncont/eq/ne/like/nlike/regex/nregex. "
|
||||
"Example: 'resp.code.gte:200 AND resp.code.lt:300 AND req.host.cont:\"api\"'."
|
||||
)
|
||||
|
||||
|
||||
def _is_httpql_error(exc: Exception) -> bool:
|
||||
message = str(exc).lower()
|
||||
return "httpql" in message or ("filter" in message and "pars" in message)
|
||||
|
||||
|
||||
def _httpql_error(exc: Exception, httpql_filter: str | None) -> str:
|
||||
"""Return an actionable error for a rejected HTTPQL filter.
|
||||
|
||||
Preserves Caido's exact parser message and echoes the offending query so
|
||||
the agent can self-correct instead of retrying the same broken filter.
|
||||
"""
|
||||
logger.info("list_requests rejected HTTPQL filter %r: %s", httpql_filter, exc)
|
||||
return json.dumps(
|
||||
{
|
||||
"success": False,
|
||||
"error": f"Invalid HTTPQL filter: {exc}",
|
||||
"httpql_filter": httpql_filter,
|
||||
"hint": _HTTPQL_HINT,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
default=str,
|
||||
)
|
||||
|
||||
|
||||
@function_tool(timeout=120)
|
||||
async def list_requests(
|
||||
ctx: RunContextWrapper,
|
||||
@@ -184,12 +155,13 @@ async def list_requests(
|
||||
sort_order: ``asc`` or ``desc``.
|
||||
scope_id: Restrict to a Caido scope (managed via ``scope_rules``).
|
||||
"""
|
||||
proxy = _ctx_proxy(ctx)
|
||||
if proxy is None:
|
||||
client = _ctx_client(ctx)
|
||||
if client is None:
|
||||
return _no_client()
|
||||
|
||||
try:
|
||||
connection = await proxy.call(
|
||||
connection = await _call(
|
||||
client,
|
||||
lambda client: caido_api.list_requests_with_client(
|
||||
client,
|
||||
httpql_filter=httpql_filter,
|
||||
@@ -198,7 +170,7 @@ async def list_requests(
|
||||
sort_by=sort_by,
|
||||
sort_order=sort_order,
|
||||
scope_id=scope_id,
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
entries = []
|
||||
@@ -252,8 +224,6 @@ async def list_requests(
|
||||
default=str,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
if httpql_filter and _is_httpql_error(exc):
|
||||
return _httpql_error(exc, httpql_filter)
|
||||
return _err("list_requests", exc)
|
||||
|
||||
|
||||
@@ -291,13 +261,14 @@ async def view_request(
|
||||
page: 1-indexed page number (only when no ``search_pattern``).
|
||||
page_size: Lines per page.
|
||||
"""
|
||||
proxy = _ctx_proxy(ctx)
|
||||
if proxy is None:
|
||||
client = _ctx_client(ctx)
|
||||
if client is None:
|
||||
return _no_client()
|
||||
|
||||
try:
|
||||
result = await proxy.call(
|
||||
lambda client: caido_api.get_request_with_client(client, request_id, part=part)
|
||||
result = await _call(
|
||||
client,
|
||||
lambda client: caido_api.get_request_with_client(client, request_id, part=part),
|
||||
)
|
||||
if result is None:
|
||||
return json.dumps(
|
||||
@@ -408,8 +379,8 @@ async def repeat_request(
|
||||
- ``body`` — replace the body string entirely.
|
||||
- ``cookies`` — dict of cookies to add/update.
|
||||
"""
|
||||
proxy = _ctx_proxy(ctx)
|
||||
if proxy is None:
|
||||
client = _ctx_client(ctx)
|
||||
if client is None:
|
||||
return _no_client()
|
||||
mods = modifications or {}
|
||||
|
||||
@@ -431,9 +402,7 @@ async def repeat_request(
|
||||
return await caido_api.replay_send_raw(client, raw=raw, connection=connection)
|
||||
|
||||
try:
|
||||
# A replay mutates target state, so don't auto-retry on a mid-send
|
||||
# transport failure (the request may already have been sent).
|
||||
replay = await proxy.call(_do, idempotent=False)
|
||||
replay = await _call(client, _do)
|
||||
if replay is None:
|
||||
return json.dumps(
|
||||
{"success": False, "error": f"Request {request_id} not found"},
|
||||
@@ -492,18 +461,19 @@ async def list_sitemap(
|
||||
(recursive subtree). Only meaningful with ``parent_id``.
|
||||
page: 1-indexed page (30 entries per page).
|
||||
"""
|
||||
proxy = _ctx_proxy(ctx)
|
||||
if proxy is None:
|
||||
client = _ctx_client(ctx)
|
||||
if client is None:
|
||||
return _no_client()
|
||||
try:
|
||||
payload = await proxy.call(
|
||||
payload = await _call(
|
||||
client,
|
||||
lambda client: caido_api.list_sitemap_with_client(
|
||||
client,
|
||||
scope_id=scope_id,
|
||||
parent_id=parent_id,
|
||||
depth=depth,
|
||||
page=page,
|
||||
)
|
||||
),
|
||||
)
|
||||
return json.dumps(payload, ensure_ascii=False, default=str)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
@@ -525,12 +495,13 @@ async def view_sitemap_entry(
|
||||
Args:
|
||||
entry_id: ID from ``list_sitemap`` (or any nested entry).
|
||||
"""
|
||||
proxy = _ctx_proxy(ctx)
|
||||
if proxy is None:
|
||||
client = _ctx_client(ctx)
|
||||
if client is None:
|
||||
return _no_client()
|
||||
try:
|
||||
payload = await proxy.call(
|
||||
lambda client: caido_api.view_sitemap_entry_with_client(client, entry_id)
|
||||
payload = await _call(
|
||||
client,
|
||||
lambda client: caido_api.view_sitemap_entry_with_client(client, entry_id),
|
||||
)
|
||||
return json.dumps(payload, ensure_ascii=False, default=str)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
@@ -583,13 +554,13 @@ async def scope_rules(
|
||||
scope_id: Required for ``get`` / ``update`` / ``delete``.
|
||||
scope_name: Required for ``create`` / ``update``.
|
||||
"""
|
||||
proxy = _ctx_proxy(ctx)
|
||||
if proxy is None:
|
||||
client = _ctx_client(ctx)
|
||||
if client is None:
|
||||
return _no_client()
|
||||
|
||||
try:
|
||||
if action == "list":
|
||||
scopes = await proxy.call(caido_api.scope_list)
|
||||
scopes = await _call(client, caido_api.scope_list)
|
||||
return json.dumps(
|
||||
{"success": True, "scopes": [_to_tool_json(s) for s in scopes]},
|
||||
ensure_ascii=False,
|
||||
@@ -602,7 +573,7 @@ async def scope_rules(
|
||||
ensure_ascii=False,
|
||||
default=str,
|
||||
)
|
||||
scope = await proxy.call(lambda client: caido_api.scope_get(client, scope_id))
|
||||
scope = await _call(client, lambda client: caido_api.scope_get(client, scope_id))
|
||||
return json.dumps(
|
||||
{"success": True, "scope": _to_tool_json(scope)},
|
||||
ensure_ascii=False,
|
||||
@@ -615,11 +586,11 @@ async def scope_rules(
|
||||
ensure_ascii=False,
|
||||
default=str,
|
||||
)
|
||||
scope = await proxy.call(
|
||||
scope = await _call(
|
||||
client,
|
||||
lambda client: caido_api.scope_create(
|
||||
client, name=scope_name, allowlist=allowlist, denylist=denylist
|
||||
),
|
||||
idempotent=False,
|
||||
)
|
||||
return json.dumps(
|
||||
{"success": True, "scope": _to_tool_json(scope)},
|
||||
@@ -636,11 +607,11 @@ async def scope_rules(
|
||||
ensure_ascii=False,
|
||||
default=str,
|
||||
)
|
||||
scope = await proxy.call(
|
||||
scope = await _call(
|
||||
client,
|
||||
lambda client: caido_api.scope_update(
|
||||
client, scope_id, name=scope_name, allowlist=allowlist, denylist=denylist
|
||||
),
|
||||
idempotent=False,
|
||||
)
|
||||
return json.dumps(
|
||||
{"success": True, "scope": _to_tool_json(scope)},
|
||||
@@ -653,7 +624,7 @@ async def scope_rules(
|
||||
ensure_ascii=False,
|
||||
default=str,
|
||||
)
|
||||
await proxy.call(lambda client: caido_api.scope_delete(client, scope_id), idempotent=False)
|
||||
await _call(client, lambda client: caido_api.scope_delete(client, scope_id))
|
||||
return json.dumps(
|
||||
{
|
||||
"success": True,
|
||||
|
||||
Reference in New Issue
Block a user