refactor(egress): extract outbound DLP request stage

bot_bottle/gateway/egress/addon.py

@@ -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,12 +25,11 @@ 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,
     match_route,
@@ -496,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
@@ -519,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",
@@ -552,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,

bot_bottle/gateway/egress/outbound_pipeline.py

@@ -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"]

tests/unit/test_egress_outbound_pipeline.py

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