From b2f61053ad30c1e5b9a7b2fc1b99cc1c35cb97fd Mon Sep 17 00:00:00 2001 From: didericis Date: Tue, 14 Jul 2026 01:55:49 -0400 Subject: [PATCH] fix(egress+orchestrator): inject per-bottle auth tokens in the shared gateway MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cut-over dropped the per-bottle token flow, so an authed egress route on the shared gateway failed with 'env var EGRESS_TOKEN_0 is unset' — the gateway reads the token from its env, but a shared gateway has no per-bottle env. Now the bottle's egress auth tokens travel to the gateway over /resolve and the addon injects from them, mirroring what the per-bottle sidecar's env did: - launch resolves the token values from the host env and hands them to the orchestrator, which holds them IN MEMORY (keyed by bottle_id, never written to the registry DB) and serves them on /resolve; - PolicyResolver.resolve_policy_and_bottle_id + resolve_client_context now return the token map alongside policy + bottle_id (one round-trip); - the egress addon overlays the process env with the bottle's tokens per request and uses that env for auth injection AND DLP — the agent never sees the credential. Secrets stay off disk (validated: /resolve returns the token, the registry DB does not contain it). SecretProvider (#355) is the future hardening. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01WBMWTEtQdJ4W5UrWuLHCck --- .../backend/docker/consolidated_launch.py | 4 +- bot_bottle/backend/docker/launch.py | 14 +++- bot_bottle/egress_addon.py | 79 +++++++++++-------- bot_bottle/egress_addon_core.py | 23 +++--- bot_bottle/orchestrator/client.py | 8 +- bot_bottle/orchestrator/control_plane.py | 13 ++- bot_bottle/orchestrator/service.py | 23 +++++- bot_bottle/policy_resolver.py | 21 +++-- tests/unit/test_egress_addon_request_flow.py | 49 +++++++++++- tests/unit/test_egress_multitenant.py | 30 ++++--- tests/unit/test_orchestrator_service.py | 12 +++ tests/unit/test_policy_resolver.py | 14 ++-- 12 files changed, 212 insertions(+), 78 deletions(-) diff --git a/bot_bottle/backend/docker/consolidated_launch.py b/bot_bottle/backend/docker/consolidated_launch.py index 50196f3..ab70dcc 100644 --- a/bot_bottle/backend/docker/consolidated_launch.py +++ b/bot_bottle/backend/docker/consolidated_launch.py @@ -98,6 +98,7 @@ def launch_consolidated( git_gate_plan: GitGatePlan, *, image_ref: str = "", + tokens: dict[str, str] | None = None, service: OrchestratorService | None = None, gateway_name: str = GATEWAY_NAME, network: str = GATEWAY_NETWORK, @@ -116,7 +117,8 @@ def launch_consolidated( inputs = registration_inputs(egress_plan) reg = client.register_bottle( - source_ip, image_ref=image_ref, policy=inputs.policy, metadata=inputs.metadata, + source_ip, image_ref=image_ref, policy=inputs.policy, + metadata=inputs.metadata, tokens=tokens, ) try: provision_git_gate(gateway_name, reg.bottle_id, git_gate_plan) diff --git a/bot_bottle/backend/docker/launch.py b/bot_bottle/backend/docker/launch.py index bfcb91c..c1a75c5 100644 --- a/bot_bottle/backend/docker/launch.py +++ b/bot_bottle/backend/docker/launch.py @@ -36,6 +36,7 @@ from contextlib import ExitStack, contextmanager from pathlib import Path from typing import Callable, Generator +from ...egress import egress_resolve_token_values from ...git_gate import ( provision_git_gate_dynamic_keys, revoke_git_gate_provisioned_keys, @@ -120,8 +121,17 @@ def launch( # Step 3: register on the orchestrator + provision this bottle's # git-gate state into the shared gateway; get the agent's attach - # context (pinned source IP, gateway address, shared network). - ctx = launch_consolidated(plan.egress_plan, git_gate_plan, image_ref=plan.image) + # context (pinned source IP, gateway address, shared network). The + # per-bottle egress auth tokens are resolved from the host env now and + # handed to the orchestrator (in memory) for the gateway to inject — + # the agent never sees them. + effective_env = {**os.environ, **plan.agent_provision.provisioned_env} + token_values = egress_resolve_token_values( + plan.egress_plan.token_env_map, effective_env, + ) + ctx = launch_consolidated( + plan.egress_plan, git_gate_plan, image_ref=plan.image, tokens=token_values, + ) stack.callback( teardown_consolidated, ctx.bottle_id, orchestrator_url=ctx.orchestrator_url, ) diff --git a/bot_bottle/egress_addon.py b/bot_bottle/egress_addon.py index d6987c3..4516d0c 100644 --- a/bot_bottle/egress_addon.py +++ b/bot_bottle/egress_addon.py @@ -10,6 +10,7 @@ import json import os import signal import sys +import typing from pathlib import Path from mitmproxy import http # type: ignore[import-not-found] # pylint: disable=import-error @@ -230,21 +231,27 @@ class EgressAddon: + "\n" ) - def _resolve_flow(self, flow: http.HTTPFlow) -> "tuple[Config, str]": - """The `(Config, supervise slug)` to apply to this request. Single-tenant - → the static `self.config` and the env slug. Consolidated → the calling - bottle's Config and its bottle id, resolved by source IP in one - round-trip (fail-closed to deny-all + empty slug if unattributed). The - identity token, if the agent injected one, is read then stripped so it - never leaks upstream. The slug keys both the proposal queue and the - per-bottle safelist, so an approval only ever affects its own bottle.""" + def _resolve_flow( + self, flow: http.HTTPFlow, + ) -> "tuple[Config, str, typing.Mapping[str, str]]": + """The `(Config, supervise slug, env)` to apply to this request. + Single-tenant → the static `self.config`, the env slug, and the process + env. Consolidated → the calling bottle's Config + bottle id + auth + tokens, resolved by source IP in one round-trip (fail-closed to deny-all + + empty slug if unattributed); `env` is the process env overlaid with + the bottle's tokens, so upstream-auth injection (and DLP) use *this* + bottle's credentials — exactly what the per-bottle sidecar's env did. + The identity token, if the agent injected one, is read then stripped so + it never leaks upstream.""" if self._resolver is None: - return self.config, self._supervise_slug + return self.config, self._supervise_slug, os.environ conn = flow.client_conn client_ip = conn.peername[0] if conn and conn.peername else "" token = flow.request.headers.get(IDENTITY_HEADER, "") flow.request.headers.pop(IDENTITY_HEADER, None) - return resolve_client_context(self._resolver, client_ip, token) + config, slug, tokens = resolve_client_context(self._resolver, client_ip, token) + env = {**os.environ, **tokens} if tokens else os.environ + return config, slug, env async def request(self, flow: http.HTTPFlow) -> None: request_path, _, query = flow.request.path.partition("?") @@ -253,14 +260,14 @@ class EgressAddon: self._serve_introspection(flow, request_path) return - config, slug = self._resolve_flow(flow) + config, slug, env = self._resolve_flow(flow) # DLP outbound scan BEFORE stripping auth — catches tokens the # agent tried to smuggle in any header, path, query param, or body. # Hostname is included to catch DNS-tunnelling exfiltration attempts. route = match_route(config.routes, flow.request.pretty_host) if route is not None: - if not await self._handle_outbound_dlp(flow, route, slug): + if not await self._handle_outbound_dlp(flow, route, slug, env): return # The redact policy may have rewritten the request line; recompute # the path/query the git checks below rely on. @@ -299,7 +306,7 @@ class EgressAddon: config.routes, flow.request.pretty_host, request_path, - os.environ, + env, request_method=flow.request.method, request_headers=req_headers, ) @@ -325,10 +332,12 @@ class EgressAddon: flow: http.HTTPFlow, route: Route, slug: str, + env: "typing.Mapping[str, str]", ) -> bool: """Scan the outbound request and apply the route's on-match policy - (PRD 0062). Returns True if the request may be forwarded, False if a - 403 response has been written to `flow`. + (PRD 0062). `env` is the per-bottle env overlay (process env + this + bottle's tokens) used for DLP detection. Returns True if the request may + be forwarded, False if a 403 response has been written to `flow`. Loops so the supervise policy can re-scan after each approval — a second, un-approved token in the same request is still caught.""" @@ -345,7 +354,7 @@ class EgressAddon: flow.request.pretty_host, request_path, query, headers, "", ) result = scan_outbound( - route, scan_text, os.environ, + route, scan_text, env, safe_tokens=self._safe_tokens_for(slug), crlf_text=crlf_text, ) if result is None or result.severity != "block": @@ -356,7 +365,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): + if self._redact_outbound(flow, route, env): if self.config.log >= LOG_BLOCKS: sys.stderr.write(json.dumps({ "event": "egress_redacted", @@ -384,29 +393,31 @@ class EgressAddon: if not self._supervise_available(slug): self._block_dlp(flow, result) return False - approved = await self._supervise_token_block(flow, request_path, result, slug) + approved = await self._supervise_token_block(flow, request_path, result, slug, env) if not approved: 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) -> bool: + 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. 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.""" + 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=os.environ) + 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=os.environ)) + 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=os.environ)) + redacted_path = strip_crlf(redact_tokens(flow.request.path, env=env)) if redacted_path != flow.request.path: flow.request.path = redacted_path @@ -419,7 +430,7 @@ class EgressAddon: crlf_text = build_outbound_scan_text( flow.request.pretty_host, request_path, query, headers, "", ) - result = scan_outbound(route, scan_text, os.environ, crlf_text=crlf_text) + result = scan_outbound(route, scan_text, env, crlf_text=crlf_text) return result is None or result.severity != "block" async def _supervise_token_block( @@ -428,20 +439,22 @@ class EgressAddon: request_path: str, result: ScanResult, slug: str, + env: "typing.Mapping[str, str]", ) -> bool: """Route a token DLP block to the operator's supervisor queue and wait. `slug` attributes the proposal to the calling bottle (its own queue + safelist) — in the shared gateway this is what keeps one bottle's - approval from unblocking another's request. Returns True if the operator - approved (the matched value is added to that bottle's safelist and the - caller re-scans); False if the request must be blocked (a 403 response - has been written to `flow`).""" + approval from unblocking another's request. `env` is the per-bottle env + overlay used to redact secrets from the proposal text. Returns True if + the operator approved (the matched value is added to that bottle's + safelist and the caller re-scans); False if the request must be blocked + (a 403 response has been written to `flow`).""" host = flow.request.pretty_host payload = build_token_allow_payload( - redact_tokens(host, env=os.environ), + redact_tokens(host, env=env), flow.request.method, - redact_tokens(request_path, env=os.environ), + redact_tokens(request_path, env=env), result, ) proposal = _sv.Proposal.new( diff --git a/bot_bottle/egress_addon_core.py b/bot_bottle/egress_addon_core.py index 2e9d152..24f632c 100644 --- a/bot_bottle/egress_addon_core.py +++ b/bot_bottle/egress_addon_core.py @@ -450,28 +450,31 @@ def resolve_client_config( class ContextResolverLike(typing.Protocol): """The bit of `policy_resolver.PolicyResolver` `resolve_client_context` - needs — one round-trip returning both policy and bottle id.""" + needs — one round-trip returning policy, bottle id, and auth tokens.""" def resolve_policy_and_bottle_id( self, source_ip: str, identity_token: str = ..., - ) -> "tuple[str | None, str | None]": + ) -> "tuple[str | None, str | None, dict[str, str]]": ... def resolve_client_context( resolver: ContextResolverLike, client_ip: str, identity_token: str = "", -) -> "tuple[Config, str]": - """The calling client's `(Config, bottle_id)` in one round-trip — +) -> "tuple[Config, str, dict[str, str]]": + """The calling client's `(Config, bottle_id, tokens)` in one round-trip — **fail-closed**. The Config follows `resolve_client_config`'s deny-all rules; the bottle id is `""` whenever unattributed or the orchestrator - errored, which the caller treats as "supervise unavailable for this - bottle" (never another bottle's queue). One `/resolve` keys both the - per-request egress policy and the per-bottle supervise queue + safelist.""" + errored (caller treats as "supervise unavailable", never another bottle's + queue); `tokens` are the per-bottle upstream auth values the addon injects. + One `/resolve` keys the egress policy, the supervise queue + safelist, and + auth injection.""" try: - policy, bottle_id = resolver.resolve_policy_and_bottle_id(client_ip, identity_token) + policy, bottle_id, tokens = resolver.resolve_policy_and_bottle_id( + client_ip, identity_token, + ) except Exception: # noqa: BLE001 # pylint: disable=broad-exception-caught - return Config(routes=()), "" # orchestrator unreachable/errored → deny - return _config_from_policy(policy), (bottle_id or "") + return Config(routes=()), "", {} # orchestrator unreachable/errored → deny + return _config_from_policy(policy), (bottle_id or ""), tokens # --------------------------------------------------------------------------- diff --git a/bot_bottle/orchestrator/client.py b/bot_bottle/orchestrator/client.py index a0c4f93..f6bdc71 100644 --- a/bot_bottle/orchestrator/client.py +++ b/bot_bottle/orchestrator/client.py @@ -90,14 +90,18 @@ class OrchestratorClient: image_ref: str = "", metadata: str = "", policy: str = "", + tokens: dict[str, str] | None = None, ) -> RegisteredBottle: - """Register a bottle and broker its launch (`POST /bottles`). Returns - its minted id + identity token.""" + """Register a bottle and broker its launch (`POST /bottles`). `tokens` + are the per-bottle egress auth values (env_name -> value) the + orchestrator holds in memory for the gateway to inject. Returns the + minted id + identity token.""" payload = self._ok("POST", "/bottles", { "source_ip": source_ip, "image_ref": image_ref, "metadata": metadata, "policy": policy, + "tokens": tokens or {}, }) bottle_id = payload.get("bottle_id") token = payload.get("identity_token") diff --git a/bot_bottle/orchestrator/control_plane.py b/bot_bottle/orchestrator/control_plane.py index a7459c5..2e5dec1 100644 --- a/bot_bottle/orchestrator/control_plane.py +++ b/bot_bottle/orchestrator/control_plane.py @@ -81,11 +81,16 @@ def dispatch( # pylint: disable=too-many-return-statements,too-many-branches image_ref = data.get("image_ref") metadata = data.get("metadata") policy = data.get("policy") + raw_tokens = data.get("tokens") + tokens = { + k: v for k, v in raw_tokens.items() if isinstance(k, str) and isinstance(v, str) + } if isinstance(raw_tokens, dict) else {} rec = orch.launch_bottle( source_ip, image_ref=image_ref if isinstance(image_ref, str) else "", metadata=metadata if isinstance(metadata, str) else "", policy=policy if isinstance(policy, str) else "", + tokens=tokens, ) return 201, {"bottle_id": rec.bottle_id, "identity_token": rec.identity_token} @@ -137,7 +142,13 @@ def dispatch( # pylint: disable=too-many-return-statements,too-many-branches rec = orch.resolve(source_ip, token if isinstance(token, str) else "") if rec is None: return 403, {"error": "unattributed"} - return 200, {"bottle_id": rec.bottle_id, "policy": rec.policy} + # tokens are the in-memory per-bottle egress auth values the gateway + # injects; served here, never persisted. + return 200, { + "bottle_id": rec.bottle_id, + "policy": rec.policy, + "tokens": orch.tokens_for(rec.bottle_id), + } return 404, {"error": "not found"} diff --git a/bot_bottle/orchestrator/service.py b/bot_bottle/orchestrator/service.py index 25bd913..ecd1e86 100644 --- a/bot_bottle/orchestrator/service.py +++ b/bot_bottle/orchestrator/service.py @@ -38,6 +38,12 @@ class Orchestrator: self._broker = broker self._secret = sign_secret self._gateway = gateway + # Per-bottle egress auth tokens (env_name -> value), keyed by bottle_id. + # Held **in memory only** — never written to the registry DB — so the + # gateway can inject each bottle's upstream credential without secrets + # at rest. Lost on restart (re-launch re-registers them); the future + # SecretProvider (#355) replaces this with per-request minting. + self._tokens: dict[str, dict[str, str]] = {} def launch_bottle( self, @@ -47,11 +53,14 @@ class Orchestrator: slot: int | None = None, metadata: str = "", policy: str = "", + tokens: dict[str, str] | None = None, ) -> BottleRecord: - """Register a bottle (with its gateway policy) and broker its launch. - Rolls the registry entry back if the launch doesn't take, so a - failure leaves no orphan.""" + """Register a bottle (with its gateway policy + in-memory egress auth + tokens) and broker its launch. Rolls the registry entry back if the + launch doesn't take, so a failure leaves no orphan.""" rec = self.registry.register(source_ip, metadata=metadata, policy=policy) + if tokens: + self._tokens[rec.bottle_id] = dict(tokens) req = LaunchRequest( op="launch", bottle_id=rec.bottle_id, @@ -66,6 +75,7 @@ class Orchestrator: finally: if not launched: self.registry.deregister(rec.bottle_id) + self._tokens.pop(rec.bottle_id, None) return rec def teardown_bottle(self, bottle_id: str) -> bool: @@ -76,8 +86,15 @@ class Orchestrator: req = LaunchRequest(op="teardown", bottle_id=bottle_id, source_ip=rec.source_ip) self._broker.submit(sign_request(req, self._secret)) self.registry.deregister(bottle_id) + self._tokens.pop(bottle_id, None) return True + def tokens_for(self, bottle_id: str) -> dict[str, str]: + """The bottle's in-memory egress auth tokens (env_name -> value), or + empty. The gateway injects these per request; they are never + persisted.""" + return dict(self._tokens.get(bottle_id, {})) + def attribute(self, source_ip: str, identity_token: str) -> BottleRecord | None: """Fail-closed attribution (delegates to the registry).""" return self.registry.attribute(source_ip, identity_token) diff --git a/bot_bottle/policy_resolver.py b/bot_bottle/policy_resolver.py index cde742f..7449ee5 100644 --- a/bot_bottle/policy_resolver.py +++ b/bot_bottle/policy_resolver.py @@ -102,21 +102,26 @@ class PolicyResolver: def resolve_policy_and_bottle_id( self, source_ip: str, identity_token: str = "", - ) -> tuple[str | None, str | None]: - """Both the policy blob and the bottle id in a single `/resolve` — so - a caller that needs each (the egress addon: policy for routing, bottle - id to key the calling bottle's supervise queue + safelist) makes one - round-trip, not two. Returns `(None, None)` when unattributed (a clean - `403`). Raises `PolicyResolveError` if the orchestrator can't be - reached, so the caller still fails closed.""" + ) -> tuple[str | None, str | None, dict[str, str]]: + """The policy blob, bottle id, **and per-bottle egress auth tokens** in + a single `/resolve` — so the egress addon gets everything it needs + (policy for routing, bottle id for the supervise queue/safelist, tokens + to inject upstream auth) in one round-trip. Returns `(None, None, {})` + when unattributed (a clean `403`). Raises `PolicyResolveError` if the + orchestrator can't be reached, so the caller still fails closed.""" payload = self._post_resolve(source_ip, identity_token) if payload is None: - return None, None + return None, None, {} policy = payload.get("policy") bottle_id = payload.get("bottle_id") + raw = payload.get("tokens") + tokens = { + k: v for k, v in raw.items() if isinstance(k, str) and isinstance(v, str) + } if isinstance(raw, dict) else {} return ( policy if isinstance(policy, str) else "", bottle_id if isinstance(bottle_id, str) and bottle_id else None, + tokens, ) diff --git a/tests/unit/test_egress_addon_request_flow.py b/tests/unit/test_egress_addon_request_flow.py index 779816d..aa46628 100644 --- a/tests/unit/test_egress_addon_request_flow.py +++ b/tests/unit/test_egress_addon_request_flow.py @@ -233,16 +233,19 @@ class _CtxResolver: """Fake orchestrator resolver: maps source IP -> bottle id, and grants the same allow-list to any attributed bottle (unattributed -> deny).""" - def __init__(self, ip_to_bottle: dict[str, str]) -> None: + def __init__( + self, ip_to_bottle: dict[str, str], tokens: dict[str, str] | None = None, + ) -> None: self._map = ip_to_bottle + self._tokens = tokens or {} def resolve_policy_and_bottle_id( self, source_ip: str, identity_token: str = "", - ) -> tuple[str | None, str | None]: + ) -> tuple[str | None, str | None, dict[str, str]]: del identity_token bottle_id = self._map.get(source_ip) policy = "routes:\n - host: api.example.com\n" if bottle_id else None - return policy, bottle_id + return policy, bottle_id, (self._tokens if bottle_id else {}) # --------------------------------------------------------------------------- @@ -803,6 +806,46 @@ class TestSuperviseMultiTenant(unittest.TestCase): _run_request(addon, flow) self.assertEqual(["bottle-b"], seen) # proposal keyed by the resolved bottle + def test_auth_token_injected_from_resolved_tokens(self) -> None: + # The bottle's upstream token comes from /resolve (in-memory on the + # orchestrator), NOT the gateway's env — the request is forwarded with + # Authorization injected. This is the token-injection cut-over fix. + policy = ( + 'routes:\n - host: api.example.com\n' + ' auth_scheme: "Bearer"\n token_env: "EGRESS_TOKEN_0"\n' + ) + + class _AuthResolver: + def resolve_policy_and_bottle_id(self, ip: str, identity_token: str = ""): + del ip, identity_token + return policy, "bottle-a", {"EGRESS_TOKEN_0": "sk-secret"} + + addon = _addon(Config(routes=())) + addon._resolver = cast(Any, _AuthResolver()) + flow = _with_client_ip(_Flow(_Request(host="api.example.com", method="GET")), "10.0.0.1") + _run_request(addon, flow) + self.assertIsNone(flow.response) # forwarded, not blocked + self.assertEqual("Bearer sk-secret", flow.request.headers.get("authorization")) + + def test_authed_route_without_token_is_blocked(self) -> None: + # No token resolved for the bottle → the authed route can't inject, so + # the request is blocked (the error the cut-over originally produced). + policy = ( + 'routes:\n - host: api.example.com\n' + ' auth_scheme: "Bearer"\n token_env: "EGRESS_TOKEN_0"\n' + ) + + class _NoTokenResolver: + def resolve_policy_and_bottle_id(self, ip: str, identity_token: str = ""): + del ip, identity_token + return policy, "bottle-a", {} # no tokens + + addon = _addon(Config(routes=())) + addon._resolver = cast(Any, _NoTokenResolver()) + flow = _with_client_ip(_Flow(_Request(host="api.example.com", method="GET")), "10.0.0.1") + _run_request(addon, flow) + self.assertIsNotNone(flow.response) # blocked — token unset + def test_unattributed_source_ip_cannot_supervise(self) -> None: addon = self._consolidated_addon() # 10.9.9.9 is not in the resolver map -> deny-all config, empty slug. diff --git a/tests/unit/test_egress_multitenant.py b/tests/unit/test_egress_multitenant.py index cff26a7..e35ec95 100644 --- a/tests/unit/test_egress_multitenant.py +++ b/tests/unit/test_egress_multitenant.py @@ -51,45 +51,55 @@ class TestResolveClientConfig(unittest.TestCase): class _FakeContextResolver: def __init__( - self, policy: str | None = None, bottle_id: str | None = None, raises: bool = False, + self, policy: str | None = None, bottle_id: str | None = None, + raises: bool = False, tokens: dict[str, str] | None = None, ) -> None: self._policy = policy self._bottle_id = bottle_id self._raises = raises + self._tokens = tokens or {} def resolve_policy_and_bottle_id( self, source_ip: str, identity_token: str = "", - ) -> tuple[str | None, str | None]: + ) -> tuple[str | None, str | None, dict[str, str]]: if self._raises: raise PolicyResolveError("orchestrator down") - return self._policy, self._bottle_id + return self._policy, self._bottle_id, self._tokens class TestResolveClientContext(unittest.TestCase): - def test_returns_config_and_bottle_id(self) -> None: - cfg, slug = resolve_client_context( - _FakeContextResolver(policy="routes:\n - host: example.com\n", bottle_id="b1"), + def test_returns_config_bottle_id_and_tokens(self) -> None: + cfg, slug, tokens = resolve_client_context( + _FakeContextResolver( + policy="routes:\n - host: example.com\n", bottle_id="b1", + tokens={"EGRESS_TOKEN_0": "sekret"}, + ), "10.243.0.1", ) self.assertEqual(("example.com",), tuple(r.host for r in cfg.routes)) self.assertEqual("b1", slug) + self.assertEqual({"EGRESS_TOKEN_0": "sekret"}, tokens) # for auth injection def test_unattributed_denies_and_empty_slug(self) -> None: - cfg, slug = resolve_client_context( + cfg, slug, tokens = resolve_client_context( _FakeContextResolver(policy=None, bottle_id=None), "10.243.0.9", ) self.assertEqual((), cfg.routes) self.assertEqual("", slug) # no bottle → supervise unavailable + self.assertEqual({}, tokens) - def test_resolver_error_denies_and_empty_slug(self) -> None: - cfg, slug = resolve_client_context(_FakeContextResolver(raises=True), "10.243.0.1") + def test_resolver_error_denies_empty_slug_no_tokens(self) -> None: + cfg, slug, tokens = resolve_client_context( + _FakeContextResolver(raises=True), "10.243.0.1", + ) self.assertEqual((), cfg.routes) self.assertEqual("", slug) + self.assertEqual({}, tokens) def test_unparseable_policy_denies_but_keeps_slug(self) -> None: # A bad policy denies egress, but the bottle is still attributed (its # supervise queue is keyed by the id, independent of route parsing). - cfg, slug = resolve_client_context( + cfg, slug, _tokens = resolve_client_context( _FakeContextResolver(policy="routes: notalist\n", bottle_id="b2"), "10.243.0.1", ) self.assertEqual((), cfg.routes) diff --git a/tests/unit/test_orchestrator_service.py b/tests/unit/test_orchestrator_service.py index 5b36436..5dcbbd1 100644 --- a/tests/unit/test_orchestrator_service.py +++ b/tests/unit/test_orchestrator_service.py @@ -92,6 +92,18 @@ class TestOrchestrator(unittest.TestCase): assert got is not None self.assertEqual('{"routes":[]}', got.policy) + def test_tokens_held_in_memory_and_cleared_on_teardown(self) -> None: + rec = self.orch.launch_bottle("10.243.0.5", tokens={"EGRESS_TOKEN_0": "sk"}) + self.assertEqual({"EGRESS_TOKEN_0": "sk"}, self.orch.tokens_for(rec.bottle_id)) + # not persisted in the registry (redacted record has no tokens field) + self.assertNotIn("tokens", rec.redacted()) + self.orch.teardown_bottle(rec.bottle_id) + self.assertEqual({}, self.orch.tokens_for(rec.bottle_id)) + + def test_tokens_default_empty(self) -> None: + rec = self.orch.launch_bottle("10.243.0.6") + self.assertEqual({}, self.orch.tokens_for(rec.bottle_id)) + def test_set_policy_live_reload(self) -> None: rec = self.orch.launch_bottle("10.243.0.3") self.assertTrue(self.orch.set_policy(rec.bottle_id, '{"x":1}')) diff --git a/tests/unit/test_policy_resolver.py b/tests/unit/test_policy_resolver.py index 8110ddc..c06b7a4 100644 --- a/tests/unit/test_policy_resolver.py +++ b/tests/unit/test_policy_resolver.py @@ -89,13 +89,17 @@ class TestPolicyResolver(unittest.TestCase): self.r.resolve_bottle_id("10.243.0.1", "tok") def test_resolve_policy_and_bottle_id_one_call(self) -> None: - with patch(_URLOPEN, return_value=_resp({"bottle_id": "b1", "policy": "P"})) as m: - self.assertEqual(("P", "b1"), self.r.resolve_policy_and_bottle_id("10.243.0.1", "t")) - self.assertEqual(1, m.call_count) # both from a single /resolve + payload = {"bottle_id": "b1", "policy": "P", "tokens": {"EGRESS_TOKEN_0": "s"}} + with patch(_URLOPEN, return_value=_resp(payload)) as m: + self.assertEqual( + ("P", "b1", {"EGRESS_TOKEN_0": "s"}), + self.r.resolve_policy_and_bottle_id("10.243.0.1", "t"), + ) + self.assertEqual(1, m.call_count) # policy + id + tokens from a single /resolve - def test_resolve_policy_and_bottle_id_403_is_none_none(self) -> None: + def test_resolve_policy_and_bottle_id_403_is_none_none_empty(self) -> None: with patch(_URLOPEN, side_effect=_http_error(403)): - self.assertEqual((None, None), self.r.resolve_policy_and_bottle_id("10.243.0.9")) + self.assertEqual((None, None, {}), self.r.resolve_policy_and_bottle_id("10.243.0.9")) def test_resolve_policy_and_bottle_id_error_raises(self) -> None: with patch(_URLOPEN, side_effect=urllib.error.URLError("refused")):