diff --git a/bot_bottle/egress_addon.py b/bot_bottle/egress_addon.py index c69a635..70abe3d 100644 --- a/bot_bottle/egress_addon.py +++ b/bot_bottle/egress_addon.py @@ -78,6 +78,16 @@ ORCHESTRATOR_URL_ENV = "BOT_BOTTLE_ORCHESTRATOR_URL" # is still stripped defensively (git-http uses that header on its own port). IDENTITY_HEADER = "x-bot-bottle-identity" +# Per-flow key under which `request()` stashes the resolved (Config, supervise +# slug, env) so the later `response()` and `websocket_message()` hooks scan +# against the *calling bottle's* policy. In the consolidated (multi-tenant) +# gateway the static `self.config` is empty — every request's real policy comes +# from the per-request `/resolve` — so a hook that fell back to `self.config` +# would find no route and silently skip its DLP scan (fail-open). Resolving once +# at the request and reusing it also avoids a `/resolve` round-trip per response +# and per WebSocket frame. +_FLOW_CTX_KEY = "bot_bottle_egress_ctx" + def _token_from_proxy_auth(header: str) -> str: """Extract the identity token (the password) from a `Proxy-Authorization: @@ -223,31 +233,39 @@ class EgressAddon: {"Content-Type": "text/plain; charset=utf-8"}, ) - def _log_request(self, flow: http.HTTPFlow) -> None: + def _log_request( + self, flow: http.HTTPFlow, env: "typing.Mapping[str, str]", + ) -> None: + # `env` is the per-flow resolved overlay (process env + this bottle's + # /resolve tokens), so the log redaction scrubs the calling bottle's + # provisioned secrets — not just the process-level ones in os.environ. headers = { - k: redact_tokens(v, env=os.environ) + k: redact_tokens(v, env=env) for k, v in flow.request.headers.items() if k.lower() != "authorization" } - body = redact_tokens(flow.request.get_text(strict=False) or "", env=os.environ) + body = redact_tokens(flow.request.get_text(strict=False) or "", env=env) sys.stderr.write( json.dumps({ "event": "egress_request", - "host": redact_tokens(flow.request.pretty_host, env=os.environ), + "host": redact_tokens(flow.request.pretty_host, env=env), "method": flow.request.method, - "path": redact_tokens(flow.request.path, env=os.environ), + "path": redact_tokens(flow.request.path, env=env), "headers": headers, "body": body, }) + "\n" ) - def _log_response(self, flow: http.HTTPFlow) -> None: + def _log_response( + self, flow: http.HTTPFlow, env: "typing.Mapping[str, str]", + ) -> None: + # Per-flow env overlay (see _log_request): redact this bottle's tokens. headers = { - k: redact_tokens(v, env=os.environ) + k: redact_tokens(v, env=env) for k, v in flow.response.headers.items() } - body = redact_tokens(flow.response.get_text(strict=False) or "", env=os.environ) + body = redact_tokens(flow.response.get_text(strict=False) or "", env=env) sys.stderr.write( json.dumps({ "event": "egress_response", @@ -280,6 +298,36 @@ class EgressAddon: env = {**os.environ, **tokens} if tokens else os.environ return config, slug, env + def _stash_flow_ctx( + self, + flow: http.HTTPFlow, + config: Config, + slug: str, + env: "typing.Mapping[str, str]", + ) -> None: + """Remember the per-flow context `request()` resolved, so the later + `response()` / `websocket_message()` hooks reuse it — scanning against + the same bottle's policy the request was decided on, with one `/resolve` + per flow rather than one per frame.""" + meta = getattr(flow, "metadata", None) + if isinstance(meta, dict): + meta[_FLOW_CTX_KEY] = (config, slug, env) + + def _flow_ctx( + self, flow: http.HTTPFlow, + ) -> "tuple[Config, str, typing.Mapping[str, str]]": + """The `(Config, supervise slug, env)` `request()` resolved for this + flow, so a later hook scans against the calling bottle's policy — not the + empty static config the consolidated gateway carries. Falls back to the + single-tenant static values for a flow that never passed through + `request()` (or a flow object without metadata).""" + meta = getattr(flow, "metadata", None) + if isinstance(meta, dict): + ctx = meta.get(_FLOW_CTX_KEY) + if ctx is not None: + return ctx + return self.config, self._supervise_slug, os.environ + def _request_token(self, flow: http.HTTPFlow) -> str: """The per-bottle identity token for this request, from the proxy credentials — the delivery mechanism (`HTTPS_PROXY=http://id:token@gw`) @@ -319,6 +367,9 @@ class EgressAddon: return config, slug, env = self._resolve_flow(flow) + # Stash for the response / websocket hooks so their DLP scans use this + # bottle's resolved policy, not the empty static config (see _flow_ctx). + self._stash_flow_ctx(flow, config, slug, env) # DLP outbound scan BEFORE stripping auth — catches tokens the # agent tried to smuggle in any header, path, query param, or body. @@ -377,7 +428,7 @@ class EgressAddon: flow.request.headers["authorization"] = decision.inject_authorization if config.log >= LOG_FULL: - self._log_request(flow) + self._log_request(flow, env) def _block_dlp(self, flow: http.HTTPFlow, result: ScanResult) -> None: ctx = self._req_ctx(flow) @@ -586,14 +637,17 @@ class EgressAddon: await asyncio.sleep(TOKEN_ALLOW_POLL_INTERVAL_SECONDS) def response(self, flow: http.HTTPFlow) -> None: - """DLP inbound scan on response headers and body.""" - route = match_route(self.config.routes, flow.request.pretty_host) + """DLP inbound scan on response headers and body, against the calling + bottle's resolved config (multi-tenant) or the static config + (single-tenant) — see `_flow_ctx`.""" + config, _slug, env = self._flow_ctx(flow) + route = match_route(config.routes, flow.request.pretty_host) if route is None: return if flow.response is None: return - if self.config.log >= LOG_FULL: - self._log_response(flow) + if config.log >= LOG_FULL: + self._log_response(flow, env) resp_headers = {k.lower(): v for k, v in flow.response.headers.items()} body = flow.response.get_text(strict=False) or "" scan_text = build_inbound_scan_text(resp_headers, body) @@ -610,7 +664,7 @@ class EgressAddon: resp_ctx = {**resp_ctx, "context": result.context} if result.severity == "block": self._block(flow, f"egress DLP: {result.reason}", ctx=resp_ctx) - elif result.severity == "warn" and self.config.log >= LOG_BLOCKS: + elif result.severity == "warn" and config.log >= LOG_BLOCKS: sys.stderr.write( json.dumps({ "event": "egress_warn", @@ -621,7 +675,10 @@ class EgressAddon: ) def websocket_message(self, flow: http.HTTPFlow) -> None: - """DLP scan on WebSocket frames. + """DLP scan on WebSocket frames, against the calling bottle's resolved + config (see `_flow_ctx`). `request()` resolves and stashes the per-flow + (config, slug, env) at the upgrade, so both the multi-tenant and + single-tenant gateways scan here. Outbound frames (from_client) are scanned for credential leakage; inbound frames are scanned for prompt injection. On a block the @@ -630,10 +687,8 @@ class EgressAddon: """ if flow.websocket is None: # type: ignore[union-attr] return - # WebSocket DLP runs against the static config only (single-tenant); in - # the consolidated gateway self.config has no routes, so this is inert - # until websocket routing is made source-IP-aware (a separate slice). - route = match_route(self.config.routes, flow.request.pretty_host) + config, slug, env = self._flow_ctx(flow) + route = match_route(config.routes, flow.request.pretty_host) if route is None: return message = flow.websocket.messages[-1] # type: ignore[union-attr] @@ -642,8 +697,8 @@ class EgressAddon: # A WebSocket data frame is not an HTTP request line, so CRLF is # not an injection vector here — scan only for credential leakage. result = scan_outbound( - route, content, os.environ, - safe_tokens=self._safe_tokens_for(self._supervise_slug), crlf_text="", + route, content, env, + safe_tokens=self._safe_tokens_for(slug), crlf_text="", ) if result is not None and result.severity == "block": sys.stderr.write(f"egress DLP: {result.reason}\n") diff --git a/tests/unit/test_egress_addon_log_redaction.py b/tests/unit/test_egress_addon_log_redaction.py index 5a648e2..45b60f8 100644 --- a/tests/unit/test_egress_addon_log_redaction.py +++ b/tests/unit/test_egress_addon_log_redaction.py @@ -8,6 +8,7 @@ real mitmproxy package.""" from __future__ import annotations import json +import os import sys import types import unittest @@ -107,14 +108,14 @@ class _Flow: def _log_request(addon: EgressAddon, flow: _Flow) -> dict[str, Any]: buf = StringIO() with patch("sys.stderr", buf): - addon._log_request(flow) # type: ignore[arg-type] + addon._log_request(flow, os.environ) # type: ignore[arg-type] return json.loads(buf.getvalue()) def _log_response(addon: EgressAddon, flow: _Flow) -> dict[str, Any]: buf = StringIO() with patch("sys.stderr", buf): - addon._log_response(flow) # type: ignore[arg-type] + addon._log_response(flow, os.environ) # type: ignore[arg-type] return json.loads(buf.getvalue()) diff --git a/tests/unit/test_egress_addon_request_flow.py b/tests/unit/test_egress_addon_request_flow.py index 71256a0..2deff87 100644 --- a/tests/unit/test_egress_addon_request_flow.py +++ b/tests/unit/test_egress_addon_request_flow.py @@ -141,6 +141,10 @@ class _Flow: self.response = response self.websocket: Any = None self.killed = False + # mitmproxy flows carry a per-flow `metadata` dict for addon use; the + # egress addon stashes the resolved (config, slug, env) there in + # request() so the response/websocket hooks reuse it. + self.metadata: dict[str, Any] = {} def kill(self) -> None: self.killed = True @@ -859,5 +863,91 @@ class TestSuperviseMultiTenant(unittest.TestCase): self.assertNotIn(_OPENAI_KEY, addon._safe_tokens_for("")) +class TestMultiTenantInboundDlp(unittest.TestCase): + """Consolidated gateway: the response + websocket DLP hooks must scan + against the *calling bottle's* config, resolved by source IP in request() + and reused here. The static `self.config` is empty in this mode, so before + the flow-context stash these hooks silently skipped every scan (fail-open). + """ + + def _consolidated_addon(self) -> EgressAddon: + addon = _addon(Config(routes=())) # empty static config, as in prod + addon._resolver = cast(Any, _CtxResolver({"10.0.0.1": "bottle-a"})) + return addon + + def test_response_injection_blocked_after_request_resolves(self) -> None: + addon = self._consolidated_addon() + flow = _with_client_ip(_Flow(_Request(host="api.example.com")), "10.0.0.1") + _run_request(addon, flow) # resolves + stashes bottle-a's allowlist + self.assertIsNone(flow.response) # request forwarded + flow.response = _Response(200, content=_INJECTION_BLOCK) + addon.response(flow) # type: ignore[arg-type] + assert flow.response is not None + # Empty static config would have found no route and left this 200. + self.assertEqual(403, flow.response.status_code) + + def test_websocket_outbound_token_killed_after_request_resolves(self) -> None: + addon = self._consolidated_addon() + flow = _with_client_ip(_Flow(_Request(host="api.example.com")), "10.0.0.1") + _run_request(addon, flow) # the ws upgrade resolves + stashes the config + flow.websocket = _WebSocketData( + [_Message(f"k={_OPENAI_KEY}".encode(), from_client=True)] + ) + addon.websocket_message(flow) # type: ignore[arg-type] + self.assertTrue(flow.killed) # scanned against bottle-a's route now + + def test_websocket_inbound_injection_killed_after_request_resolves(self) -> None: + addon = self._consolidated_addon() + flow = _with_client_ip(_Flow(_Request(host="api.example.com")), "10.0.0.1") + _run_request(addon, flow) + flow.websocket = _WebSocketData( + [_Message(_INJECTION_BLOCK.encode(), from_client=False)] + ) + addon.websocket_message(flow) # type: ignore[arg-type] + self.assertTrue(flow.killed) + + def test_response_log_redacts_per_bottle_resolve_token(self) -> None: + # LOG_FULL response logging now runs in multi-tenant mode, so it must + # scrub the calling bottle's /resolve token — which lives only in the + # resolved env overlay, never in the gateway's os.environ. A + # non-token-shaped secret is caught only via that env, so os.environ + # redaction (the pre-fix behaviour) would leak it into the log. + secret = "bottle-a-provisioned-secret-value" + policy = "log: 2\nroutes:\n - host: api.example.com\n" + + class _TokenResolver: + def resolve_policy_and_bottle_id( + self, ip: str, identity_token: str = "", + ) -> tuple[str | None, str | None, dict[str, str]]: + del ip, identity_token + return policy, "bottle-a", {"EGRESS_TOKEN_0": secret} + + addon = _addon(Config(routes=())) + addon._resolver = cast(Any, _TokenResolver()) + flow = _with_client_ip(_Flow(_Request(host="api.example.com")), "10.0.0.1") + _run_request(addon, flow) + flow.response = _Response(200, content=f"echo {secret} back") + buf = StringIO() + with patch("sys.stderr", buf): + addon.response(flow) # type: ignore[arg-type] + logged = buf.getvalue() + self.assertIn("egress_response", logged) # LOG_FULL logged the response + self.assertNotIn(secret, logged) # redacted via the resolved env overlay + + def test_unattributed_flow_has_no_route_so_frame_passes(self) -> None: + # An unattributed source resolves to a deny-all (empty) config, so the + # request is blocked at the upgrade and any later frame has no route to + # scan against — it passes rather than being attributed to a bottle. + addon = self._consolidated_addon() + flow = _with_client_ip(_Flow(_Request(host="api.example.com")), "10.9.9.9") + _run_request(addon, flow) + self.assertIsNotNone(flow.response) # blocked at the upgrade + flow.websocket = _WebSocketData( + [_Message(f"k={_OPENAI_KEY}".encode(), from_client=True)] + ) + addon.websocket_message(flow) # type: ignore[arg-type] + self.assertFalse(flow.killed) + + if __name__ == "__main__": unittest.main()