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")

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())
 
 

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()