Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| c7c3a79028 | |||
| bb1776a858 | |||
| a24fe0264d | |||
| 105538d3a6 |
@@ -30,7 +30,7 @@ from pathlib import Path
|
||||
|
||||
from ...log import info
|
||||
from .. import EnumerationError
|
||||
from . import util
|
||||
from . import lifecycle_lock, util
|
||||
from .bottle_cleanup_plan import FirecrackerBottleCleanupPlan
|
||||
|
||||
|
||||
@@ -126,12 +126,42 @@ def prepare_cleanup() -> FirecrackerBottleCleanupPlan:
|
||||
|
||||
|
||||
def cleanup(plan: FirecrackerBottleCleanupPlan) -> None:
|
||||
for pid in plan.vm_pids:
|
||||
info(f"kill firecracker VM pid {pid}")
|
||||
"""Revalidate the preview under the launch lock, then remove its survivors."""
|
||||
with lifecycle_lock.hold():
|
||||
fresh = prepare_cleanup()
|
||||
approved_pids = set(plan.vm_pids).intersection(fresh.vm_pids)
|
||||
approved_dirs = set(plan.run_dirs).intersection(fresh.run_dirs)
|
||||
for pid in sorted(approved_pids):
|
||||
_terminate_orphan(pid, _run_root())
|
||||
for path in sorted(approved_dirs):
|
||||
info(f"rm -rf {path}")
|
||||
shutil.rmtree(path, ignore_errors=True)
|
||||
|
||||
|
||||
def _terminate_orphan(pid: int, run_root: Path) -> None:
|
||||
"""Signal exactly the process identity that still owns an orphan config."""
|
||||
try:
|
||||
pidfd = os.pidfd_open(pid)
|
||||
except ProcessLookupError:
|
||||
return
|
||||
except OSError as exc:
|
||||
raise EnumerationError(
|
||||
f"could not pin Firecracker pid {pid} for cleanup: {exc}"
|
||||
) from exc
|
||||
try:
|
||||
try:
|
||||
os.kill(pid, signal.SIGTERM)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
for path in plan.run_dirs:
|
||||
info(f"rm -rf {path}")
|
||||
shutil.rmtree(path, ignore_errors=True)
|
||||
raw = Path(f"/proc/{pid}/cmdline").read_bytes()
|
||||
except FileNotFoundError:
|
||||
return
|
||||
except OSError as exc:
|
||||
raise EnumerationError(
|
||||
f"could not revalidate Firecracker pid {pid}: {exc}"
|
||||
) from exc
|
||||
command = raw.replace(b"\0", b" ").decode(errors="replace")
|
||||
run_dir = _run_dir_of(command, run_root)
|
||||
if run_dir is None or run_dir.is_dir():
|
||||
return
|
||||
info(f"kill firecracker VM pid {pid}")
|
||||
signal.pidfd_send_signal(pidfd, signal.SIGTERM)
|
||||
finally:
|
||||
os.close(pidfd)
|
||||
|
||||
@@ -46,7 +46,7 @@ from ...log import die, info, warn
|
||||
from ...supervisor.types import SUPERVISE_PORT
|
||||
from ..docker.egress import EGRESS_PORT
|
||||
from ..util import AGENT_CA_BUNDLE, AGENT_CA_PATH
|
||||
from . import firecracker_vm, image_builder, isolation_probe, netpool, util
|
||||
from . import firecracker_vm, image_builder, isolation_probe, lifecycle_lock, netpool, util
|
||||
from .bottle import FirecrackerBottle
|
||||
from .bottle_plan import FirecrackerBottlePlan
|
||||
from ...orchestrator.store.config_store import resolve_teardown_timeout
|
||||
@@ -164,25 +164,29 @@ def launch(
|
||||
)
|
||||
|
||||
# Step 6: build the per-bottle rootfs + SSH key, then boot.
|
||||
run_dir = util.cache_dir() / "run" / plan.slug
|
||||
run_dir.mkdir(parents=True, exist_ok=True)
|
||||
# Remove the run dir on teardown so the per-bottle rootfs.ext4 (~1G)
|
||||
# doesn't leak. Registered before vm.terminate below so it runs *after*
|
||||
# it (ExitStack is LIFO): the VM is gone before we rm its rootfs.
|
||||
stack.callback(lambda: shutil.rmtree(run_dir, ignore_errors=True))
|
||||
rootfs = run_dir / "rootfs.ext4"
|
||||
util.build_rootfs_ext4(agent_base, rootfs)
|
||||
private_key, pubkey = util.generate_keypair(run_dir)
|
||||
# Cleanup takes the same lock while refreshing its process snapshot.
|
||||
# Hold it until the VMM exists so a newly-created run dir can never be
|
||||
# mistaken for an orphan in the build-before-boot window.
|
||||
with lifecycle_lock.hold():
|
||||
run_dir = util.cache_dir() / "run" / plan.slug
|
||||
run_dir.mkdir(parents=True, exist_ok=True)
|
||||
# Remove the run dir on teardown so the per-bottle rootfs.ext4 (~1G)
|
||||
# doesn't leak. Registered before vm.terminate below so it runs
|
||||
# *after* it (ExitStack is LIFO).
|
||||
stack.callback(lambda: shutil.rmtree(run_dir, ignore_errors=True))
|
||||
rootfs = run_dir / "rootfs.ext4"
|
||||
util.build_rootfs_ext4(agent_base, rootfs)
|
||||
private_key, pubkey = util.generate_keypair(run_dir)
|
||||
|
||||
vm = firecracker_vm.boot(
|
||||
name=plan.container_name,
|
||||
rootfs=rootfs,
|
||||
tap=slot.iface,
|
||||
guest_ip=slot.guest_ip,
|
||||
host_ip=slot.host_ip,
|
||||
pubkey=pubkey,
|
||||
run_dir=run_dir,
|
||||
)
|
||||
vm = firecracker_vm.boot(
|
||||
name=plan.container_name,
|
||||
rootfs=rootfs,
|
||||
tap=slot.iface,
|
||||
guest_ip=slot.guest_ip,
|
||||
host_ip=slot.host_ip,
|
||||
pubkey=pubkey,
|
||||
run_dir=run_dir,
|
||||
)
|
||||
stack.callback(vm.terminate)
|
||||
firecracker_vm.wait_for_ssh(vm, private_key)
|
||||
persist_env_var_secret(private_key, slot.guest_ip, ctx.env_var_secret)
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
"""Serialize Firecracker run-directory creation with orphan cleanup."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import fcntl
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Generator
|
||||
|
||||
from . import util
|
||||
|
||||
|
||||
def _lock_path() -> Path:
|
||||
return util.cache_dir() / "run.lifecycle.lock"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def hold() -> Generator[None]:
|
||||
"""Exclude cleanup while a launch directory lacks a visible VMM."""
|
||||
path = _lock_path()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with path.open("a", encoding="utf-8") as handle:
|
||||
fcntl.flock(handle, fcntl.LOCK_EX)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
fcntl.flock(handle, fcntl.LOCK_UN)
|
||||
|
||||
|
||||
__all__ = ["hold"]
|
||||
@@ -4,7 +4,8 @@ from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
|
||||
from ...log import info, warn
|
||||
from .. import EnumerationError
|
||||
from ...log import info
|
||||
from . import util as container_mod
|
||||
from .bottle_cleanup_plan import MacosContainerBottleCleanupPlan
|
||||
|
||||
@@ -19,8 +20,8 @@ def _list_prefixed_containers() -> list[str]:
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
warn(f"container list failed: {result.stderr.strip()}")
|
||||
return []
|
||||
detail = result.stderr.strip() or f"exit {result.returncode}"
|
||||
raise EnumerationError(f"container list failed: {detail}")
|
||||
return sorted(
|
||||
name for name in (line.strip() for line in result.stdout.splitlines())
|
||||
if name.startswith(_PREFIX)
|
||||
@@ -35,7 +36,8 @@ def _list_prefixed_networks() -> list[str]:
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
return []
|
||||
detail = result.stderr.strip() or f"exit {result.returncode}"
|
||||
raise EnumerationError(f"container network list failed: {detail}")
|
||||
return sorted(
|
||||
name for name in (line.strip() for line in result.stdout.splitlines())
|
||||
if name.startswith(_PREFIX)
|
||||
|
||||
@@ -52,7 +52,11 @@ def cmd_cleanup(_argv: list[str]) -> int:
|
||||
info("cleanup: skipped")
|
||||
return 0
|
||||
|
||||
for name, backend, plan in prepared:
|
||||
# Confirmation authorizes a fresh authoritative snapshot, not blind use of
|
||||
# identities that may have changed while the operator reviewed the preview.
|
||||
refreshed = [(name, backend, backend.prepare_cleanup())
|
||||
for name, backend, _plan in prepared]
|
||||
for name, backend, plan in refreshed:
|
||||
if plan.empty:
|
||||
continue
|
||||
backend.cleanup(plan)
|
||||
|
||||
@@ -16,7 +16,7 @@ import typing
|
||||
from mitmproxy import http # type: ignore[import-not-found] # pylint: disable=import-error
|
||||
|
||||
from bot_bottle.constants import IDENTITY_HEADER
|
||||
from bot_bottle.gateway.egress.dlp_detectors import redact_tokens, strip_crlf
|
||||
from bot_bottle.gateway.egress.dlp_detectors import redact_tokens
|
||||
from bot_bottle.gateway.egress.dlp_config import (
|
||||
DEFAULT_OUTBOUND_ON_MATCH,
|
||||
ON_MATCH_BLOCK,
|
||||
@@ -25,19 +25,19 @@ from bot_bottle.gateway.egress.dlp_config import (
|
||||
from bot_bottle.gateway.egress.context import resolve_client_context
|
||||
from bot_bottle.gateway.egress.dlp import (
|
||||
build_inbound_scan_text,
|
||||
build_outbound_scan_text,
|
||||
build_token_allow_payload,
|
||||
outbound_scan_headers,
|
||||
scan_inbound,
|
||||
scan_outbound,
|
||||
)
|
||||
from bot_bottle.gateway.egress.outbound_pipeline import redact_request, scan_request
|
||||
from bot_bottle.gateway.egress.matching import (
|
||||
decide,
|
||||
decide_git_fetch,
|
||||
is_git_fetch_request,
|
||||
is_git_push_request,
|
||||
match_route,
|
||||
)
|
||||
from bot_bottle.gateway.egress.request_pipeline import (
|
||||
evaluate_route_policy,
|
||||
git_block_reason,
|
||||
)
|
||||
from bot_bottle.gateway.egress.schema import route_to_yaml_dict
|
||||
from bot_bottle.gateway.egress.types import (
|
||||
LOG_BLOCKS,
|
||||
@@ -435,21 +435,12 @@ class EgressAddon:
|
||||
request_path: str, query: str,
|
||||
) -> bool:
|
||||
"""Apply the HTTPS Git push/fetch boundary before general routing."""
|
||||
if is_git_push_request(request_path, query):
|
||||
self._block(
|
||||
flow,
|
||||
"egress: git push over HTTPS is not supported; "
|
||||
"use the bottle.git SSH path (gitleaks-scanned by "
|
||||
"git-gate's pre-receive hook).",
|
||||
ctx=self._req_ctx(flow),
|
||||
)
|
||||
return False
|
||||
if not is_git_fetch_request(request_path, query):
|
||||
reason = git_block_reason(
|
||||
config.routes, flow.request.pretty_host, request_path, query,
|
||||
)
|
||||
if not reason:
|
||||
return True
|
||||
git_decision = decide_git_fetch(config.routes, flow.request.pretty_host)
|
||||
if git_decision.action != "block":
|
||||
return True
|
||||
self._block(flow, git_decision.reason, ctx=self._req_ctx(flow))
|
||||
self._block(flow, reason, ctx=self._req_ctx(flow))
|
||||
return False
|
||||
|
||||
def _apply_route_policy(
|
||||
@@ -461,30 +452,26 @@ class EgressAddon:
|
||||
# are caught above; the route may inject gateway-owned auth below.
|
||||
# Routes with preserve_auth=True pass the header through as-is so the
|
||||
# agent's own credentials (e.g. registry bearer tokens) reach the upstream.
|
||||
if route is None or not route.preserve_auth:
|
||||
result = evaluate_route_policy(
|
||||
config,
|
||||
route,
|
||||
host=flow.request.pretty_host,
|
||||
request_path=request_path,
|
||||
method=flow.request.method,
|
||||
headers=dict(flow.request.headers),
|
||||
env=env,
|
||||
)
|
||||
if result.strip_authorization:
|
||||
flow.request.headers.pop("authorization", None)
|
||||
|
||||
# Build headers mapping for match evaluation
|
||||
req_headers = {k.lower(): v for k, v in flow.request.headers.items()}
|
||||
|
||||
decision = decide(
|
||||
config.routes,
|
||||
flow.request.pretty_host,
|
||||
request_path,
|
||||
env,
|
||||
request_method=flow.request.method,
|
||||
request_headers=req_headers,
|
||||
deny_reason=config.deny_reason,
|
||||
)
|
||||
|
||||
if decision.action == "block":
|
||||
self._block(flow, decision.reason, ctx=self._req_ctx(flow))
|
||||
if result.block_reason:
|
||||
self._block(flow, result.block_reason, ctx=self._req_ctx(flow))
|
||||
return
|
||||
|
||||
if decision.inject_authorization is not None:
|
||||
flow.request.headers["authorization"] = decision.inject_authorization
|
||||
if result.inject_authorization is not None:
|
||||
flow.request.headers["authorization"] = result.inject_authorization
|
||||
|
||||
if config.log >= LOG_FULL:
|
||||
if result.log_request:
|
||||
self._log_request(flow, env)
|
||||
|
||||
def _block_dlp(self, flow: http.HTTPFlow, result: ScanResult) -> None:
|
||||
@@ -508,20 +495,12 @@ class EgressAddon:
|
||||
Loops so the supervise policy can re-scan after each approval — a
|
||||
second, un-approved token in the same request is still caught."""
|
||||
while True:
|
||||
request_path, _, query = flow.request.path.partition("?")
|
||||
body = flow.request.get_text(strict=False) or ""
|
||||
headers = outbound_scan_headers(route, dict(flow.request.headers))
|
||||
scan_text = build_outbound_scan_text(
|
||||
flow.request.pretty_host, request_path, query, headers, body,
|
||||
)
|
||||
# CRLF is scanned only over the request line + headers, never the
|
||||
# body (see scan_outbound) — a body is not an injection vector.
|
||||
crlf_text = build_outbound_scan_text(
|
||||
flow.request.pretty_host, request_path, query, headers, "",
|
||||
)
|
||||
result = scan_outbound(
|
||||
route, scan_text, env,
|
||||
safe_tokens=self._safe_tokens_for(slug), crlf_text=crlf_text,
|
||||
request_path, _, _ = flow.request.path.partition("?")
|
||||
result = scan_request(
|
||||
flow.request,
|
||||
route,
|
||||
env,
|
||||
safe_tokens=self._safe_tokens_for(slug),
|
||||
)
|
||||
if result is None or result.severity != "block":
|
||||
return True
|
||||
@@ -531,7 +510,7 @@ class EgressAddon:
|
||||
# redact scrubs every detection (tokens and structural CRLF) and
|
||||
# forwards; it fails closed only if a match survives the scrub.
|
||||
if policy == ON_MATCH_REDACT:
|
||||
if self._redact_outbound(flow, route, env):
|
||||
if redact_request(flow.request, route, env):
|
||||
if self._flow_log(flow) >= LOG_BLOCKS:
|
||||
sys.stderr.write(json.dumps({
|
||||
"event": "egress_redacted",
|
||||
@@ -564,41 +543,6 @@ class EgressAddon:
|
||||
return False # _supervise_token_block wrote the 403 response
|
||||
# loop: the approved value is now in safe_tokens; re-scan.
|
||||
|
||||
def _redact_outbound(
|
||||
self, flow: http.HTTPFlow, route: Route, env: "typing.Mapping[str, str]",
|
||||
) -> bool:
|
||||
"""Scrub detected tokens (and CRLF injection sequences) from the mutable
|
||||
request surfaces (body, headers, path/query) and re-scan. `env` is the
|
||||
per-bottle env overlay. Returns True if the request is now clean; False
|
||||
if a block-severity match remains on a surface redaction cannot rewrite
|
||||
(the hostname) so the caller fails closed."""
|
||||
body = flow.request.get_text(strict=False)
|
||||
if body:
|
||||
redacted_body = redact_tokens(body, env=env)
|
||||
if redacted_body != body:
|
||||
flow.request.text = redacted_body
|
||||
for name, value in list(flow.request.headers.items()):
|
||||
if name.lower() == "host":
|
||||
continue # routing-critical; never a legitimate token
|
||||
redacted = strip_crlf(redact_tokens(value, env=env))
|
||||
if redacted != value:
|
||||
flow.request.headers[name] = redacted
|
||||
redacted_path = strip_crlf(redact_tokens(flow.request.path, env=env))
|
||||
if redacted_path != flow.request.path:
|
||||
flow.request.path = redacted_path
|
||||
|
||||
request_path, _, query = flow.request.path.partition("?")
|
||||
new_body = flow.request.get_text(strict=False) or ""
|
||||
headers = outbound_scan_headers(route, dict(flow.request.headers))
|
||||
scan_text = build_outbound_scan_text(
|
||||
flow.request.pretty_host, request_path, query, headers, new_body,
|
||||
)
|
||||
crlf_text = build_outbound_scan_text(
|
||||
flow.request.pretty_host, request_path, query, headers, "",
|
||||
)
|
||||
result = scan_outbound(route, scan_text, env, crlf_text=crlf_text)
|
||||
return result is None or result.severity != "block"
|
||||
|
||||
async def _supervise_token_block(
|
||||
self,
|
||||
flow: http.HTTPFlow,
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
"""Outbound DLP request scanning and redaction for the egress pipeline."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import ItemsView, Mapping, Protocol
|
||||
|
||||
from .dlp import (
|
||||
build_outbound_scan_text,
|
||||
outbound_scan_headers,
|
||||
scan_outbound,
|
||||
)
|
||||
from .dlp_detectors import redact_tokens, strip_crlf
|
||||
from .types import Route, ScanResult
|
||||
|
||||
|
||||
class MutableHeaders(Protocol):
|
||||
def items(self) -> ItemsView[str, str]: ...
|
||||
def __getitem__(self, name: str, /) -> str: ...
|
||||
def __setitem__(self, name: str, value: str, /) -> None: ...
|
||||
|
||||
|
||||
class MutableRequest(Protocol):
|
||||
pretty_host: str
|
||||
path: str
|
||||
headers: MutableHeaders
|
||||
text: str
|
||||
|
||||
def get_text(self, strict: bool = False) -> str | None: ...
|
||||
|
||||
|
||||
def scan_request(
|
||||
request: MutableRequest,
|
||||
route: Route,
|
||||
env: Mapping[str, str],
|
||||
*,
|
||||
safe_tokens: set[str] | None = None,
|
||||
) -> ScanResult | None:
|
||||
"""Scan all mutable outbound request surfaces in their canonical order."""
|
||||
request_path, _, query = request.path.partition("?")
|
||||
headers = outbound_scan_headers(route, dict(request.headers.items()))
|
||||
body = request.get_text(strict=False) or ""
|
||||
scan_text = build_outbound_scan_text(
|
||||
request.pretty_host, request_path, query, headers, body,
|
||||
)
|
||||
# Bodies cannot alter HTTP framing, so CRLF detection is deliberately
|
||||
# restricted to the request line and headers.
|
||||
crlf_text = build_outbound_scan_text(
|
||||
request.pretty_host, request_path, query, headers, "",
|
||||
)
|
||||
return scan_outbound(
|
||||
route,
|
||||
scan_text,
|
||||
env,
|
||||
safe_tokens=safe_tokens,
|
||||
crlf_text=crlf_text,
|
||||
)
|
||||
|
||||
|
||||
def redact_request(
|
||||
request: MutableRequest,
|
||||
route: Route,
|
||||
env: Mapping[str, str],
|
||||
) -> bool:
|
||||
"""Redact mutable request surfaces and return whether the result is clean."""
|
||||
body = request.get_text(strict=False)
|
||||
if body:
|
||||
redacted_body = redact_tokens(body, env=env)
|
||||
if redacted_body != body:
|
||||
request.text = redacted_body
|
||||
for name, value in list(request.headers.items()):
|
||||
if name.lower() == "host":
|
||||
continue
|
||||
redacted = strip_crlf(redact_tokens(value, env=env))
|
||||
if redacted != value:
|
||||
request.headers[name] = redacted
|
||||
redacted_path = strip_crlf(redact_tokens(request.path, env=env))
|
||||
if redacted_path != request.path:
|
||||
request.path = redacted_path
|
||||
result = scan_request(request, route, env)
|
||||
return result is None or result.severity != "block"
|
||||
|
||||
|
||||
__all__ = ["MutableRequest", "redact_request", "scan_request"]
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Framework-neutral request policy stages for the egress adapter.
|
||||
|
||||
The mitmproxy addon owns flow mutation and response construction. This module
|
||||
owns the ordered Git and route-policy decisions so those rules remain directly
|
||||
testable without a live proxy flow.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Mapping, Sequence
|
||||
|
||||
from .matching import (
|
||||
decide,
|
||||
decide_git_fetch,
|
||||
is_git_fetch_request,
|
||||
is_git_push_request,
|
||||
)
|
||||
from .types import LOG_FULL, Config, Route
|
||||
|
||||
GIT_PUSH_BLOCK_REASON = (
|
||||
"egress: git push over HTTPS is not supported; "
|
||||
"use the bottle.git SSH path (gitleaks-scanned by "
|
||||
"git-gate's pre-receive hook)."
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RoutePolicyResult:
|
||||
"""The flow mutations and outcome produced by general route policy."""
|
||||
|
||||
block_reason: str = ""
|
||||
strip_authorization: bool = False
|
||||
inject_authorization: str | None = None
|
||||
log_request: bool = False
|
||||
|
||||
|
||||
def git_block_reason(
|
||||
routes: Sequence[Route],
|
||||
host: str,
|
||||
request_path: str,
|
||||
query: str,
|
||||
) -> str:
|
||||
"""Return the HTTPS Git policy denial, or ``""`` when allowed."""
|
||||
if is_git_push_request(request_path, query):
|
||||
return GIT_PUSH_BLOCK_REASON
|
||||
if not is_git_fetch_request(request_path, query):
|
||||
return ""
|
||||
decision = decide_git_fetch(routes, host)
|
||||
return decision.reason if decision.action == "block" else ""
|
||||
|
||||
|
||||
def evaluate_route_policy(
|
||||
config: Config,
|
||||
route: Route | None,
|
||||
*,
|
||||
host: str,
|
||||
request_path: str,
|
||||
method: str,
|
||||
headers: Mapping[str, str],
|
||||
env: Mapping[str, str],
|
||||
) -> RoutePolicyResult:
|
||||
"""Evaluate authorization stripping, matching, injection, and logging."""
|
||||
strip_authorization = route is None or not route.preserve_auth
|
||||
effective_headers = {
|
||||
name.lower(): value
|
||||
for name, value in headers.items()
|
||||
if not (strip_authorization and name.lower() == "authorization")
|
||||
}
|
||||
decision = decide(
|
||||
config.routes,
|
||||
host,
|
||||
request_path,
|
||||
env,
|
||||
request_method=method,
|
||||
request_headers=effective_headers,
|
||||
deny_reason=config.deny_reason,
|
||||
)
|
||||
return RoutePolicyResult(
|
||||
block_reason=decision.reason if decision.action == "block" else "",
|
||||
strip_authorization=strip_authorization,
|
||||
inject_authorization=decision.inject_authorization,
|
||||
log_request=config.log >= LOG_FULL,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"GIT_PUSH_BLOCK_REASON",
|
||||
"RoutePolicyResult",
|
||||
"evaluate_route_policy",
|
||||
"git_block_reason",
|
||||
]
|
||||
@@ -0,0 +1,79 @@
|
||||
"""Framework-neutral MCP method and tool dispatch."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Protocol
|
||||
|
||||
from bot_bottle.gateway.egress.context import resolve_client_context
|
||||
from bot_bottle.gateway.egress.schema import route_to_yaml_dict
|
||||
from bot_bottle.gateway.policy_resolver import PolicyResolver
|
||||
from bot_bottle.supervisor import types as _sv
|
||||
|
||||
|
||||
class Request(Protocol):
|
||||
@property
|
||||
def method(self) -> str: ...
|
||||
|
||||
@property
|
||||
def params(self) -> dict[str, object]: ...
|
||||
|
||||
|
||||
class MethodNotFoundError(Exception):
|
||||
"""Raised when a JSON-RPC method has no MCP handler."""
|
||||
|
||||
|
||||
Handler = Callable[[dict[str, object]], object]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Handlers:
|
||||
initialize: Handler
|
||||
tools_list: Handler
|
||||
list_routes: Handler
|
||||
check_proposal: Handler
|
||||
propose: Handler
|
||||
|
||||
|
||||
def dispatch(request: Request, handlers: Handlers) -> object:
|
||||
"""Route one parsed request without depending on the HTTP server."""
|
||||
if request.method == "initialize":
|
||||
return handlers.initialize(request.params)
|
||||
if request.method == "notifications/initialized":
|
||||
return None
|
||||
if request.method == "tools/list":
|
||||
return handlers.tools_list(request.params)
|
||||
if request.method != "tools/call":
|
||||
raise MethodNotFoundError(request.method)
|
||||
|
||||
tool = request.params.get("name")
|
||||
if tool == _sv.TOOL_LIST_EGRESS_ROUTES:
|
||||
return handlers.list_routes(request.params)
|
||||
if tool == _sv.TOOL_CHECK_PROPOSAL:
|
||||
return handlers.check_proposal(request.params)
|
||||
return handlers.propose(request.params)
|
||||
|
||||
|
||||
def resolved_routes_payload(
|
||||
resolver: PolicyResolver,
|
||||
source_ip: str,
|
||||
identity_token: str,
|
||||
) -> dict[str, object]:
|
||||
"""Render the calling bottle's routes, failing closed to an empty list."""
|
||||
config, _slug, _tokens = resolve_client_context(
|
||||
resolver, source_ip, identity_token,
|
||||
)
|
||||
body = json.dumps(
|
||||
{"routes": [route_to_yaml_dict(route) for route in config.routes]},
|
||||
indent=2,
|
||||
)
|
||||
return {"content": [{"type": "text", "text": body}], "isError": False}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Handlers",
|
||||
"MethodNotFoundError",
|
||||
"dispatch",
|
||||
"resolved_routes_payload",
|
||||
]
|
||||
@@ -58,10 +58,15 @@ import typing
|
||||
from dataclasses import dataclass
|
||||
|
||||
from bot_bottle.constants import IDENTITY_HEADER
|
||||
from bot_bottle.gateway.egress.context import resolve_client_context
|
||||
from bot_bottle.gateway.egress.schema import load_config, route_to_yaml_dict
|
||||
from bot_bottle.gateway.egress.schema import load_config
|
||||
from bot_bottle.gateway.egress.types import LOG_OFF
|
||||
from bot_bottle.gateway.policy_resolver import PolicyResolveError, PolicyResolver
|
||||
from bot_bottle.gateway.supervisor.mcp_dispatch import (
|
||||
Handlers as DispatchHandlers,
|
||||
MethodNotFoundError,
|
||||
dispatch,
|
||||
resolved_routes_payload,
|
||||
)
|
||||
from bot_bottle.supervisor import types as _sv
|
||||
|
||||
|
||||
@@ -611,6 +616,11 @@ class MCPHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
try:
|
||||
result = self._dispatch(req, config)
|
||||
except MethodNotFoundError as e:
|
||||
self._write_jsonrpc(
|
||||
jsonrpc_error(req.id, ERR_METHOD_NOT_FOUND, f"method not found: {e}"),
|
||||
)
|
||||
return
|
||||
except _RpcClientError as e:
|
||||
self._write_jsonrpc(jsonrpc_error(req.id, e.code, e.message))
|
||||
return
|
||||
@@ -633,41 +643,37 @@ class MCPHandler(http.server.BaseHTTPRequestHandler):
|
||||
self._write_jsonrpc(jsonrpc_result(req.id, result))
|
||||
|
||||
def _dispatch(self, req: JsonRpcRequest, config: ServerConfig) -> object:
|
||||
method = req.method
|
||||
if method == "initialize":
|
||||
return handle_initialize(req.params)
|
||||
if method == "notifications/initialized":
|
||||
return None # ack-only
|
||||
if method == "tools/list":
|
||||
return handle_tools_list(req.params)
|
||||
if method == "tools/call":
|
||||
# `list-egress-routes` is read-only introspection. The shared gateway
|
||||
# has no static route table (routes are resolved per request by
|
||||
# source IP), so answer it from the calling bottle's resolved policy.
|
||||
# Otherwise the agent sees an empty allowlist and composes an egress
|
||||
# proposal that *replaces* the live routes instead of extending them
|
||||
# — silently dropping base routes like api.anthropic.com on approval.
|
||||
if req.params.get("name") == _sv.TOOL_LIST_EGRESS_ROUTES:
|
||||
return self._resolved_routes_payload()
|
||||
resolver = self._resolver_or_fail()
|
||||
source_ip = self.client_address[0]
|
||||
token = self._identity_token()
|
||||
# `check-proposal` is a non-blocking read of the calling bottle's
|
||||
# own queue — attributed by (source_ip, identity_token) like a
|
||||
# proposal, but it never queues or blocks.
|
||||
if req.params.get("name") == _sv.TOOL_CHECK_PROPOSAL:
|
||||
return handle_check_proposal(
|
||||
req.params, resolver=resolver,
|
||||
source_ip=source_ip, identity_token=token,
|
||||
)
|
||||
# The control plane attributes the proposal to the source-IP + token
|
||||
# resolved bottle, so the one shared queue holds each bottle's
|
||||
# proposal under its own id — no slug is asserted by this daemon.
|
||||
return handle_tools_call(
|
||||
req.params, config, resolver=resolver,
|
||||
source_ip=source_ip, identity_token=token,
|
||||
def check(params: dict[str, object]) -> object:
|
||||
return handle_check_proposal(
|
||||
params,
|
||||
resolver=self._resolver_or_fail(),
|
||||
source_ip=self.client_address[0],
|
||||
identity_token=self._identity_token(),
|
||||
)
|
||||
raise _RpcClientError(ERR_METHOD_NOT_FOUND, f"method not found: {method}")
|
||||
|
||||
def propose(params: dict[str, object]) -> object:
|
||||
return handle_tools_call(
|
||||
params,
|
||||
config,
|
||||
resolver=self._resolver_or_fail(),
|
||||
source_ip=self.client_address[0],
|
||||
identity_token=self._identity_token(),
|
||||
)
|
||||
|
||||
return dispatch(
|
||||
req,
|
||||
DispatchHandlers(
|
||||
initialize=handle_initialize,
|
||||
tools_list=handle_tools_list,
|
||||
list_routes=lambda _params: resolved_routes_payload(
|
||||
self._resolver_or_fail(),
|
||||
self.client_address[0],
|
||||
self._identity_token(),
|
||||
),
|
||||
check_proposal=check,
|
||||
propose=propose,
|
||||
),
|
||||
)
|
||||
|
||||
def _identity_token(self) -> str:
|
||||
"""The agent's per-bottle identity token from the request header (the
|
||||
@@ -686,20 +692,6 @@ class MCPHandler(http.server.BaseHTTPRequestHandler):
|
||||
raise _RpcInternalError("supervise server has no policy resolver")
|
||||
return resolver
|
||||
|
||||
def _resolved_routes_payload(self) -> dict[str, object]:
|
||||
"""The calling bottle's live egress routes as the `list-egress-routes`
|
||||
JSON payload, resolved by (source_ip, identity token). Fail-closed: an
|
||||
unattributed source or an unreachable orchestrator yields an empty route
|
||||
list (never another bottle's), courtesy of `resolve_client_context`."""
|
||||
resolver = self._resolver_or_fail()
|
||||
conf, _slug, _tokens = resolve_client_context(
|
||||
resolver, self.client_address[0], self._identity_token(),
|
||||
)
|
||||
body = json.dumps(
|
||||
{"routes": [route_to_yaml_dict(r) for r in conf.routes]}, indent=2,
|
||||
)
|
||||
return {"content": [{"type": "text", "text": body}], "isError": False}
|
||||
|
||||
def _write_jsonrpc(self, body: bytes) -> None:
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
|
||||
@@ -41,8 +41,8 @@ class TestCmdCleanup(unittest.TestCase):
|
||||
):
|
||||
self.assertEqual(0, cmd.cmd_cleanup([]))
|
||||
|
||||
docker.prepare_cleanup.assert_called_once()
|
||||
fc.prepare_cleanup.assert_called_once()
|
||||
self.assertEqual(2, docker.prepare_cleanup.call_count)
|
||||
self.assertEqual(2, fc.prepare_cleanup.call_count)
|
||||
docker.cleanup.assert_called_once_with(docker_plan)
|
||||
fc.cleanup.assert_called_once_with(fc_plan)
|
||||
|
||||
@@ -68,7 +68,7 @@ class TestCmdCleanup(unittest.TestCase):
|
||||
):
|
||||
self.assertEqual(0, cmd.cmd_cleanup([]))
|
||||
|
||||
docker.prepare_cleanup.assert_called_once()
|
||||
self.assertEqual(2, docker.prepare_cleanup.call_count)
|
||||
docker.cleanup.assert_called_once_with(docker_plan)
|
||||
macos.prepare_cleanup.assert_not_called()
|
||||
|
||||
@@ -135,6 +135,25 @@ class TestCmdCleanup(unittest.TestCase):
|
||||
docker.cleanup.assert_called_once_with(docker_plan)
|
||||
fc.cleanup.assert_not_called()
|
||||
|
||||
def test_executes_refreshed_plan_after_confirmation(self):
|
||||
backend = MagicMock()
|
||||
preview = MagicMock(empty=False)
|
||||
refreshed = MagicMock(empty=False)
|
||||
backend.prepare_cleanup.side_effect = [preview, refreshed]
|
||||
|
||||
with patch.object(
|
||||
cmd, "known_backend_names", return_value=("firecracker",),
|
||||
), patch.object(
|
||||
cmd, "get_bottle_backend", return_value=backend,
|
||||
), patch.object(
|
||||
cmd, "has_backend", return_value=True,
|
||||
), patch.object(
|
||||
cmd, "_prompt_yes", return_value=True,
|
||||
):
|
||||
self.assertEqual(0, cmd.cmd_cleanup([]))
|
||||
|
||||
backend.cleanup.assert_called_once_with(refreshed)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Unit tests for framework-neutral outbound DLP request stages."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from bot_bottle.gateway.egress.outbound_pipeline import (
|
||||
MutableHeaders,
|
||||
redact_request,
|
||||
scan_request,
|
||||
)
|
||||
from bot_bottle.gateway.egress.types import Route
|
||||
|
||||
|
||||
class _Headers(dict[str, str]):
|
||||
pass
|
||||
|
||||
|
||||
class _Request:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
host: str = "api.example.com",
|
||||
path: str = "/v1/messages",
|
||||
headers: dict[str, str] | None = None,
|
||||
body: str = "",
|
||||
) -> None:
|
||||
self.pretty_host = host
|
||||
self.path = path
|
||||
self.headers: MutableHeaders = _Headers(headers or {})
|
||||
self.text = body
|
||||
|
||||
def get_text(self, strict: bool = False) -> str | None:
|
||||
del strict
|
||||
return self.text
|
||||
|
||||
|
||||
class TestOutboundScan(unittest.TestCase):
|
||||
def test_detects_secret_in_body(self) -> None:
|
||||
request = _Request(body="token=sk-" + "a" * 48)
|
||||
|
||||
result = scan_request(request, Route(host="api.example.com"), {})
|
||||
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual("block", result.severity if result else None)
|
||||
|
||||
def test_safe_token_is_ignored(self) -> None:
|
||||
token = "sk-" + "a" * 48
|
||||
request = _Request(body=f"token={token}")
|
||||
|
||||
result = scan_request(
|
||||
request,
|
||||
Route(host="api.example.com"),
|
||||
{},
|
||||
safe_tokens={token},
|
||||
)
|
||||
|
||||
self.assertIsNone(result)
|
||||
|
||||
|
||||
class TestOutboundRedaction(unittest.TestCase):
|
||||
def test_redacts_body_header_and_path_but_preserves_host(self) -> None:
|
||||
token = "sk-" + "a" * 48
|
||||
request = _Request(
|
||||
path=f"/v1/messages?key={token}",
|
||||
headers={"Host": "api.example.com", "X-Token": token + "\r\nInjected: yes"},
|
||||
body=f"token={token}",
|
||||
)
|
||||
|
||||
clean = redact_request(request, Route(host="api.example.com"), {})
|
||||
|
||||
self.assertTrue(clean)
|
||||
self.assertNotIn(token, request.path)
|
||||
self.assertNotIn(token, request.headers["X-Token"])
|
||||
self.assertNotIn("\r", request.headers["X-Token"])
|
||||
self.assertNotIn(token, request.text)
|
||||
self.assertEqual("api.example.com", request.headers["Host"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,89 @@
|
||||
"""Unit tests for framework-neutral egress request policy stages."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from bot_bottle.gateway.egress.request_pipeline import (
|
||||
GIT_PUSH_BLOCK_REASON,
|
||||
evaluate_route_policy,
|
||||
git_block_reason,
|
||||
)
|
||||
from bot_bottle.gateway.egress.types import Config, LOG_FULL, Route
|
||||
|
||||
|
||||
class TestGitPolicy(unittest.TestCase):
|
||||
def test_push_is_always_blocked(self) -> None:
|
||||
reason = git_block_reason(
|
||||
(), "git.example.com", "/repo.git/git-receive-pack", "",
|
||||
)
|
||||
self.assertEqual(GIT_PUSH_BLOCK_REASON, reason)
|
||||
|
||||
def test_fetch_requires_route_opt_in(self) -> None:
|
||||
path = "/repo.git/git-upload-pack"
|
||||
blocked = git_block_reason((), "git.example.com", path, "")
|
||||
allowed = git_block_reason(
|
||||
(Route(host="git.example.com", git_fetch=True),),
|
||||
"git.example.com",
|
||||
path,
|
||||
"",
|
||||
)
|
||||
self.assertTrue(blocked)
|
||||
self.assertEqual("", allowed)
|
||||
|
||||
def test_non_git_request_is_not_decided_here(self) -> None:
|
||||
self.assertEqual(
|
||||
"",
|
||||
git_block_reason((), "api.example.com", "/v1/messages", ""),
|
||||
)
|
||||
|
||||
|
||||
class TestRoutePolicy(unittest.TestCase):
|
||||
def test_strips_agent_auth_and_injects_gateway_auth(self) -> None:
|
||||
route = Route(
|
||||
host="api.example.com",
|
||||
auth_scheme="Bearer",
|
||||
token_env="API_TOKEN",
|
||||
)
|
||||
result = evaluate_route_policy(
|
||||
Config(routes=(route,)),
|
||||
route,
|
||||
host="api.example.com",
|
||||
request_path="/v1/messages",
|
||||
method="POST",
|
||||
headers={"Authorization": "agent-secret"},
|
||||
env={"API_TOKEN": "gateway-secret"},
|
||||
)
|
||||
self.assertTrue(result.strip_authorization)
|
||||
self.assertEqual("Bearer gateway-secret", result.inject_authorization)
|
||||
self.assertFalse(result.block_reason)
|
||||
|
||||
def test_preserved_auth_participates_in_matching(self) -> None:
|
||||
route = Route(host="registry.example.com", preserve_auth=True)
|
||||
result = evaluate_route_policy(
|
||||
Config(routes=(route,), log=LOG_FULL),
|
||||
route,
|
||||
host="registry.example.com",
|
||||
request_path="/v2/",
|
||||
method="GET",
|
||||
headers={"Authorization": "Bearer agent-token"},
|
||||
env={},
|
||||
)
|
||||
self.assertFalse(result.strip_authorization)
|
||||
self.assertTrue(result.log_request)
|
||||
|
||||
def test_missing_route_fails_closed(self) -> None:
|
||||
result = evaluate_route_policy(
|
||||
Config(routes=(), deny_reason="not allowed"),
|
||||
None,
|
||||
host="blocked.example.com",
|
||||
request_path="/",
|
||||
method="GET",
|
||||
headers={},
|
||||
env={},
|
||||
)
|
||||
self.assertEqual("not allowed", result.block_reason)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -132,18 +132,46 @@ class TestCleanupRemoval(unittest.TestCase):
|
||||
vm_pids=(101,),
|
||||
run_dirs=("/run/dev-x",),
|
||||
)
|
||||
with patch.object(fc_cleanup.os, "kill") as kill, \
|
||||
with patch.object(fc_cleanup, "prepare_cleanup", return_value=plan), \
|
||||
patch.object(fc_cleanup, "_run_root", return_value=Path("/run")), \
|
||||
patch.object(fc_cleanup, "_terminate_orphan") as terminate, \
|
||||
patch.object(fc_cleanup.shutil, "rmtree") as rmtree, \
|
||||
patch.object(fc_cleanup, "info"):
|
||||
fc_cleanup.cleanup(plan)
|
||||
kill.assert_called_once()
|
||||
terminate.assert_called_once_with(101, Path("/run"))
|
||||
rmtree.assert_called_once_with("/run/dev-x", ignore_errors=True)
|
||||
|
||||
def test_cleanup_tolerates_dead_pid(self):
|
||||
plan = FirecrackerBottleCleanupPlan(vm_pids=(999,))
|
||||
with patch.object(fc_cleanup.os, "kill", side_effect=ProcessLookupError), \
|
||||
patch.object(fc_cleanup, "info"):
|
||||
fc_cleanup.cleanup(plan) # must not raise
|
||||
def test_cleanup_skips_resources_no_longer_in_refreshed_plan(self):
|
||||
preview = FirecrackerBottleCleanupPlan(
|
||||
vm_pids=(999,), run_dirs=("/run/reused",),
|
||||
)
|
||||
with patch.object(
|
||||
fc_cleanup, "prepare_cleanup",
|
||||
return_value=FirecrackerBottleCleanupPlan(),
|
||||
), patch.object(fc_cleanup, "_terminate_orphan") as terminate, \
|
||||
patch.object(fc_cleanup.shutil, "rmtree") as rmtree:
|
||||
fc_cleanup.cleanup(preview)
|
||||
terminate.assert_not_called()
|
||||
rmtree.assert_not_called()
|
||||
|
||||
def test_pidfd_prevents_pid_reuse_from_signalling_unrelated_process(self):
|
||||
with patch.object(fc_cleanup.os, "pidfd_open", return_value=7), \
|
||||
patch.object(
|
||||
fc_cleanup.Path, "read_bytes",
|
||||
return_value=b"/usr/bin/python\0worker.py\0",
|
||||
), patch.object(fc_cleanup.signal, "pidfd_send_signal") as send, \
|
||||
patch.object(fc_cleanup.os, "close"):
|
||||
fc_cleanup._terminate_orphan(101, Path("/run"))
|
||||
send.assert_not_called()
|
||||
|
||||
def test_pidfd_signals_revalidated_orphan(self):
|
||||
command = b"firecracker\0--config-file\0/run/gone/config.json\0"
|
||||
with patch.object(fc_cleanup.os, "pidfd_open", return_value=7), \
|
||||
patch.object(fc_cleanup.Path, "read_bytes", return_value=command), \
|
||||
patch.object(fc_cleanup.signal, "pidfd_send_signal") as send, \
|
||||
patch.object(fc_cleanup.os, "close"), patch.object(fc_cleanup, "info"):
|
||||
fc_cleanup._terminate_orphan(101, Path("/run"))
|
||||
send.assert_called_once_with(7, fc_cleanup.signal.SIGTERM)
|
||||
|
||||
|
||||
class TestCleanupPlan(unittest.TestCase):
|
||||
|
||||
@@ -42,6 +42,22 @@ class TestMacosContainerCleanup(unittest.TestCase):
|
||||
run.call_args_list[1].args[0],
|
||||
)
|
||||
|
||||
def test_container_enumeration_failure_aborts(self):
|
||||
completed = cleanup.subprocess.CompletedProcess(
|
||||
args=[], returncode=1, stdout="", stderr="service unavailable",
|
||||
)
|
||||
with patch.object(cleanup.subprocess, "run", return_value=completed), \
|
||||
self.assertRaisesRegex(EnumerationError, "service unavailable"):
|
||||
cleanup._list_prefixed_containers()
|
||||
|
||||
def test_network_enumeration_failure_aborts(self):
|
||||
completed = cleanup.subprocess.CompletedProcess(
|
||||
args=[], returncode=1, stdout="", stderr="service unavailable",
|
||||
)
|
||||
with patch.object(cleanup.subprocess, "run", return_value=completed), \
|
||||
self.assertRaisesRegex(EnumerationError, "service unavailable"):
|
||||
cleanup._list_prefixed_networks()
|
||||
|
||||
|
||||
class TestMacosContainerEnumerate(unittest.TestCase):
|
||||
"""The backend launches bottles again (PRD 0070), so enumeration is real
|
||||
|
||||
@@ -13,6 +13,7 @@ import tempfile
|
||||
import threading
|
||||
import time
|
||||
import types
|
||||
import typing
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
@@ -47,6 +48,7 @@ from bot_bottle.gateway.supervisor.server import (
|
||||
jsonrpc_error,
|
||||
jsonrpc_result,
|
||||
parse_jsonrpc,
|
||||
resolved_routes_payload,
|
||||
validate_proposed_file,
|
||||
)
|
||||
|
||||
@@ -701,9 +703,14 @@ class TestResolvedRoutesPayload(unittest.TestCase):
|
||||
" - host: api.anthropic.com\n"
|
||||
" - host: www.google.com\n"
|
||||
)
|
||||
payload = _handler(
|
||||
_FakeSuperviseResolver(bottle_id="b1", policy=policy)
|
||||
)._resolved_routes_payload()
|
||||
payload = resolved_routes_payload(
|
||||
typing.cast(
|
||||
supervise_server.PolicyResolver,
|
||||
_FakeSuperviseResolver(bottle_id="b1", policy=policy),
|
||||
),
|
||||
_SRC,
|
||||
_TOK,
|
||||
)
|
||||
assert payload is not None
|
||||
self.assertFalse(payload["isError"]) # type: ignore[index]
|
||||
data = json.loads(payload["content"][0]["text"]) # type: ignore[index]
|
||||
@@ -713,9 +720,14 @@ class TestResolvedRoutesPayload(unittest.TestCase):
|
||||
def test_orchestrator_error_fails_closed_to_empty(self) -> None:
|
||||
# resolve_client_context swallows resolver errors → deny-all (empty),
|
||||
# never another bottle's routes.
|
||||
payload = _handler(
|
||||
_FakeSuperviseResolver(raises=True)
|
||||
)._resolved_routes_payload()
|
||||
payload = resolved_routes_payload(
|
||||
typing.cast(
|
||||
supervise_server.PolicyResolver,
|
||||
_FakeSuperviseResolver(raises=True),
|
||||
),
|
||||
_SRC,
|
||||
_TOK,
|
||||
)
|
||||
assert payload is not None
|
||||
data = json.loads(payload["content"][0]["text"]) # type: ignore[index]
|
||||
self.assertEqual([], data["routes"])
|
||||
@@ -724,7 +736,13 @@ class TestResolvedRoutesPayload(unittest.TestCase):
|
||||
# A server without a resolver is a misconfig, not a mode: raise rather
|
||||
# than list anything.
|
||||
with self.assertRaises(_RpcInternalError):
|
||||
_handler(None)._resolved_routes_payload()
|
||||
_handler(None)._dispatch(
|
||||
parse_jsonrpc(
|
||||
b'{"jsonrpc":"2.0","id":1,"method":"tools/call",'
|
||||
b'"params":{"name":"list-egress-routes"}}',
|
||||
),
|
||||
ServerConfig(),
|
||||
)
|
||||
|
||||
|
||||
class TestNonBlockingSupervise(unittest.TestCase):
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Unit tests for framework-neutral supervisor MCP dispatch."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from dataclasses import dataclass
|
||||
|
||||
from bot_bottle.gateway.supervisor.mcp_dispatch import (
|
||||
Handlers,
|
||||
MethodNotFoundError,
|
||||
dispatch,
|
||||
)
|
||||
from bot_bottle.supervisor import types as _sv
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Request:
|
||||
method: str
|
||||
params: dict[str, object]
|
||||
|
||||
|
||||
class TestDispatch(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.calls: list[str] = []
|
||||
|
||||
def handler(name: str):
|
||||
def call(_params: dict[str, object]) -> str:
|
||||
self.calls.append(name)
|
||||
return name
|
||||
|
||||
return call
|
||||
|
||||
self.handlers = Handlers(
|
||||
initialize=handler("initialize"),
|
||||
tools_list=handler("tools_list"),
|
||||
list_routes=handler("list_routes"),
|
||||
check_proposal=handler("check_proposal"),
|
||||
propose=handler("propose"),
|
||||
)
|
||||
|
||||
def request(self, method: str, **params: object) -> _Request:
|
||||
return _Request(method=method, params=params)
|
||||
|
||||
def test_routes_protocol_methods(self) -> None:
|
||||
self.assertEqual(
|
||||
"initialize", dispatch(self.request("initialize"), self.handlers),
|
||||
)
|
||||
self.assertEqual(
|
||||
"tools_list", dispatch(self.request("tools/list"), self.handlers),
|
||||
)
|
||||
self.assertIsNone(
|
||||
dispatch(self.request("notifications/initialized"), self.handlers),
|
||||
)
|
||||
|
||||
def test_routes_each_tool_class(self) -> None:
|
||||
self.assertEqual(
|
||||
"list_routes",
|
||||
dispatch(
|
||||
self.request("tools/call", name=_sv.TOOL_LIST_EGRESS_ROUTES),
|
||||
self.handlers,
|
||||
),
|
||||
)
|
||||
self.assertEqual(
|
||||
"check_proposal",
|
||||
dispatch(
|
||||
self.request("tools/call", name=_sv.TOOL_CHECK_PROPOSAL),
|
||||
self.handlers,
|
||||
),
|
||||
)
|
||||
self.assertEqual(
|
||||
"propose",
|
||||
dispatch(self.request("tools/call", name=_sv.TOOL_EGRESS_ALLOW), self.handlers),
|
||||
)
|
||||
|
||||
def test_unknown_method_is_typed(self) -> None:
|
||||
with self.assertRaisesRegex(MethodNotFoundError, "unknown"):
|
||||
dispatch(self.request("unknown"), self.handlers)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user