fix(egress): scan response + websocket DLP against the resolved per-flow config
In the consolidated (multi-tenant) gateway the addon's static `self.config` is empty — each request's real policy comes from the per-request `/resolve`. `response()` and `websocket_message()` still matched routes against that empty config, so inbound prompt-injection DLP and WebSocket credential/injection DLP silently skipped every scan (fail-open) whenever the gateway ran multi-tenant. This is backend-agnostic: the gateway image (and this addon) is shared by the Firecracker, macOS, and docker consolidated backends. Resolve the per-flow (config, slug, env) once in `request()`, stash it on `flow.metadata`, and have both hooks read it back — falling back to the static single-tenant values for a flow that never passed through `request()`. Reusing the request's one `/resolve` avoids a round-trip per response and per WebSocket frame. Adds multi-tenant regression tests for both hooks that fail against the old fall-open behaviour. Refs: audit issue #400 (finding #2) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UoEZHDjv84ChoZbozQERhJ
This commit is contained in:
@@ -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,63 @@ 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_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()
|
||||
|
||||
Reference in New Issue
Block a user