diff --git a/strix/interface/cloud/billing.py b/strix/interface/cloud/billing.py index a3392d5a..995c0c75 100644 --- a/strix/interface/cloud/billing.py +++ b/strix/interface/cloud/billing.py @@ -284,6 +284,7 @@ def _run_wallet_client( with wallet_payment_bridge( upstream_url=upstream_url, api_token=http.api_token(token), + workspace_id=http.expected_workspace_id(token_override=token is not None), expected_body=body_json.encode(), timeout=getattr(args, "timeout", None), response_observer=upstream_responses.append, diff --git a/strix/interface/cloud/http.py b/strix/interface/cloud/http.py index 7b57f80a..74df0b3f 100644 --- a/strix/interface/cloud/http.py +++ b/strix/interface/cloud/http.py @@ -159,7 +159,7 @@ def request( headers = { "Authorization": f"Bearer {api_token(token)}", } - workspace_id = _expected_workspace_id(token_override=token is not None) + workspace_id = expected_workspace_id(token_override=token is not None) if workspace_id: headers["X-Strix-Workspace"] = workspace_id if idempotency_key is not None: @@ -185,7 +185,7 @@ def request( return response -def _expected_workspace_id(*, token_override: bool) -> str | None: +def expected_workspace_id(*, token_override: bool) -> str | None: """Pin every request in this process to the workspace selected at startup.""" if _workspace_id_override: return _workspace_id_override diff --git a/strix/interface/cloud/payment_proxy.py b/strix/interface/cloud/payment_proxy.py index eada8787..b64ee46c 100644 --- a/strix/interface/cloud/payment_proxy.py +++ b/strix/interface/cloud/payment_proxy.py @@ -45,6 +45,7 @@ _HOP_BY_HOP_HEADERS = frozenset( class _BridgeState: upstream_url: str authorization: str + workspace_id: str | None expected_body: bytes path: str timeout: float @@ -112,6 +113,7 @@ def _forward_request_headers(handler: BaseHTTPRequestHandler) -> dict[str, str]: "x-forwarded-proto", "x-real-ip", "x-strix-authorization", + "x-strix-workspace", "x-vercel-forwarded-for", } return {name: value for name, value in handler.headers.items() if name.lower() not in blocked} @@ -163,6 +165,8 @@ def _make_handler(state: _BridgeState) -> type[BaseHTTPRequestHandler]: headers = _forward_request_headers(self) headers["X-Strix-Authorization"] = state.authorization + if state.workspace_id: + headers["X-Strix-Workspace"] = state.workspace_id try: response = requests.request( "POST", @@ -243,6 +247,7 @@ def wallet_payment_bridge( *, upstream_url: str, api_token: str, + workspace_id: str | None = None, expected_body: bytes, timeout: float | None = None, response_observer: Callable[[WalletUpstreamResponse], None] | None = None, @@ -258,6 +263,7 @@ def wallet_payment_bridge( state = _BridgeState( upstream_url=upstream_url, authorization=f"Bearer {api_token}", + workspace_id=workspace_id, expected_body=expected_body, path=path, timeout=timeout or _DEFAULT_REQUEST_TIMEOUT_S, diff --git a/tests/test_cloud_payment_proxy.py b/tests/test_cloud_payment_proxy.py index c6fb6a79..1fba9e5e 100644 --- a/tests/test_cloud_payment_proxy.py +++ b/tests/test_cloud_payment_proxy.py @@ -84,6 +84,7 @@ def test_bridge_forwards_only_the_approved_request_and_protected_headers( with payment_proxy.wallet_payment_bridge( upstream_url="https://app.example.test/api/v1/billing/topup", api_token="strix-secret", # noqa: S106 + workspace_id="org_trusted", expected_body=b'{"credits":5}', response_observer=observed.append, ) as wallet_url: @@ -94,6 +95,7 @@ def test_bridge_forwards_only_the_approved_request_and_protected_headers( "Authorization": "Payment wallet-proof", "Proxy-Authorization": "Basic drop-me", "X-Strix-Authorization": "Bearer attacker", + "X-Strix-Workspace": "org_attacker", }, ) @@ -101,6 +103,7 @@ def test_bridge_forwards_only_the_approved_request_and_protected_headers( headers = captured[0]["headers"] assert headers["Authorization"] == "Payment wallet-proof" assert headers["X-Strix-Authorization"] == "Bearer strix-secret" + assert headers["X-Strix-Workspace"] == "org_trusted" assert "Proxy-Authorization" not in headers assert not any(name.lower() in {"host", "content-length"} for name in headers) assert observed == [