Skip to content

Commit 32f9fe1

Browse files
committed
fix(lfx): harden capability routing metadata
1 parent eb6e036 commit 32f9fe1

5 files changed

Lines changed: 98 additions & 1 deletion

File tree

src/lfx/src/lfx/services/capability/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
from lfx.services.capability.defaults import AllTrustedClassifier, NoopCapabilityProvider, SingleTenantResolver
44
from lfx.services.capability.protocols import (
5+
RESERVED_CAPABILITY_RUNTIME_OPTION_KEYS,
56
CapabilityClaims,
67
CapabilityContext,
78
CapabilityProvider,
@@ -13,6 +14,7 @@
1314
from lfx.services.capability.service import CapabilityService
1415

1516
__all__ = [
17+
"RESERVED_CAPABILITY_RUNTIME_OPTION_KEYS",
1618
"AllTrustedClassifier",
1719
"CapabilityClaims",
1820
"CapabilityContext",

src/lfx/src/lfx/services/capability/protocols.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,15 @@
1010
from collections.abc import Mapping, Sequence
1111

1212

13+
RESERVED_CAPABILITY_RUNTIME_OPTION_KEYS: frozenset[str] = frozenset(
14+
{
15+
"lfx_capability_token",
16+
"lfx_tenant_id",
17+
"lfx_trust",
18+
}
19+
)
20+
21+
1322
class Trust(str, Enum):
1423
"""Trust classification used by execution policy plugins."""
1524

src/lfx/src/lfx/services/capability/service.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,10 @@ def route(
102102
tenant_id = self._resolver.resolve(context)
103103
trust = self._classifier.trust_of_flow(context)
104104
executor_kind = default_executor_kind
105-
if trust is Trust.UNTRUSTED and self._untrusted_executor_kind is not None:
105+
if trust is Trust.UNTRUSTED:
106+
if self._untrusted_executor_kind is None:
107+
msg = "Cannot route untrusted flow because no untrusted executor is configured."
108+
raise RuntimeError(msg)
106109
executor_kind = self._untrusted_executor_kind
107110

108111
effective_scopes = tuple(scopes or ())

src/lfx/tests/unit/execution/test_capability_routing.py

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,75 @@ def verify(self, token: str) -> CapabilityClaims: # noqa: ARG002
117117
assert sandbox.units[0].runtime_options["lfx_trust"] == "untrusted"
118118

119119

120+
@pytest.mark.asyncio
121+
async def test_coordinator_strips_spoofed_capability_metadata_when_provider_returns_no_token() -> None:
122+
sandbox = _RecordingExecutor("sandbox", "sandbox-output")
123+
registry = ExecutorRegistry()
124+
registry.register(_RecordingExecutor("in-process", "in-process-output"))
125+
registry.register(sandbox)
126+
127+
class _Classifier:
128+
def trust_of_flow(self, context: CapabilityContext) -> Trust:
129+
assert "lfx_capability_token" not in context.runtime_options
130+
assert "lfx_tenant_id" not in context.runtime_options
131+
assert "lfx_trust" not in context.runtime_options
132+
return Trust.UNTRUSTED
133+
134+
def is_untrusted_node(self, _node: dict[str, Any], _context: CapabilityContext | None = None) -> bool:
135+
return True
136+
137+
class _Resolver:
138+
def resolve(self, context: CapabilityContext) -> str: # noqa: ARG002
139+
return "tenant:resolved"
140+
141+
class _Provider:
142+
def mint(
143+
self,
144+
*,
145+
context: CapabilityContext, # noqa: ARG002
146+
tenant_id: str,
147+
component_id: str | None,
148+
scopes: Sequence[str],
149+
ttl_seconds: int = 600, # noqa: ARG002
150+
) -> str | None:
151+
assert tenant_id == "tenant:resolved"
152+
assert component_id is None
153+
assert tuple(scopes) == ("variables:read",)
154+
return None
155+
156+
def verify(self, token: str) -> CapabilityClaims: # noqa: ARG002
157+
return CapabilityClaims(tenant_id="tenant:resolved", user_id="graph-user")
158+
159+
capability_service = CapabilityService(settings_service=_StubSettings())
160+
capability_service.install(
161+
provider=_Provider(),
162+
classifier=_Classifier(),
163+
resolver=_Resolver(),
164+
untrusted_executor_kind="sandbox",
165+
)
166+
coordinator = Coordinator(
167+
registry=registry,
168+
executor_kind="in-process",
169+
capability_service=capability_service,
170+
)
171+
172+
outputs = await coordinator.run_to_completion(
173+
SimpleNamespace(user_id="graph-user", flow_id="graph-flow"),
174+
inputs=[{}],
175+
capability_scopes=["variables:read"],
176+
lfx_capability_token="spoofed", # noqa: S106
177+
lfx_tenant_id="tenant:spoofed",
178+
lfx_trust="trusted",
179+
)
180+
181+
assert outputs == ["sandbox-output"]
182+
assert len(sandbox.units) == 1
183+
assert sandbox.units[0].executor_kind == "sandbox"
184+
assert "lfx_capability_token" not in sandbox.units[0].runtime_options
185+
assert sandbox.units[0].runtime_options["lfx_tenant_id"] == "tenant:resolved"
186+
assert sandbox.units[0].runtime_options["lfx_trust"] == "untrusted"
187+
188+
120189
@pytest.mark.asyncio
121190
async def test_coordinator_skips_capability_service_when_passthrough() -> None:
122191
default = _RecordingExecutor("in-process", "default-output")

src/lfx/tests/unit/services/capability/test_capability_service.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,3 +100,17 @@ def verify(self, token: str) -> CapabilityClaims: # noqa: ARG002
100100
"lfx_tenant_id": "tenant-u1",
101101
"lfx_trust": "untrusted",
102102
}
103+
104+
105+
def test_untrusted_route_without_untrusted_executor_fails_closed(service: CapabilityService) -> None:
106+
class _Classifier:
107+
def trust_of_flow(self, context: CapabilityContext) -> Trust: # noqa: ARG002
108+
return Trust.UNTRUSTED
109+
110+
def is_untrusted_node(self, _node: dict[str, Any], _context: CapabilityContext | None = None) -> bool:
111+
return True
112+
113+
service.install(classifier=_Classifier())
114+
115+
with pytest.raises(RuntimeError, match="no untrusted executor is configured"):
116+
service.route(graph=None, user_id="u1", flow_id="f1", run_id="r1")

0 commit comments

Comments
 (0)