From 8b4f94eee860568f956fb69bbdb538e33e6629f1 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 15 Sep 2026 14:15:44 +0200 Subject: [PATCH] Python: Restrict argument labels against owned input integrity --- .../packages/core/agent_framework/security.py | 26 ++-- python/packages/core/tests/test_security.py | 134 +++++++++++++++++- .../security/FIDES_DEVELOPER_GUIDE.md | 23 +-- 3 files changed, 162 insertions(+), 21 deletions(-) diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index 9278ff3b292..8f6c8285d2e 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -1268,14 +1268,15 @@ class LabelTrackingFunctionMiddleware(FunctionMiddleware, _SecurityScopeBinding) +----------+------------------------------------------+----------------------------+ | Tier 2 | Tool's source_integrity declaration | No embedded labels | +----------+------------------------------------------+----------------------------+ - | Tier 3 | Join (combine_labels) of input arg labels| No embedded labels AND | - | | | no source_integrity | + | Tier 3 | Owned input labels or configured default | No embedded labels AND | + | | restricted by argument labels | no source_integrity | +----------+------------------------------------------+----------------------------+ Tools can declare their source_integrity in additional_properties: - source_integrity="trusted": Tool produces trusted data (e.g., internal computation) - source_integrity="untrusted": Tool fetches external/untrusted data - - (not set): Falls back to tier 3 (input label join), or UNTRUSTED if no inputs + - (not set): Inherits integrity from resolved, owned variable references, or uses + default_integrity (UNTRUSTED by default). Argument labels may only restrict this baseline. This middleware: 1. Extracts labels from tool input arguments (tier 3 input) @@ -1604,8 +1605,9 @@ def _get_input_labels(self, context: FunctionInvocationContext) -> list[ContentL Recursively inspects the arguments passed to a tool to find any VariableReferenceContent objects or labeled data, and collects their labels. - These labels are used as the tier-3 fallback (lowest priority) when - neither embedded labels nor a source_integrity declaration are present. + These labels propagate confidentiality and restrict the tier-3 integrity + baseline when no source_integrity declaration is present. They cannot + establish trust; that baseline comes from owned references or the configured default. Args: context: The function invocation context containing arguments. @@ -1800,7 +1802,7 @@ async def process( # Step 3: Build tiered fallback_label # This label is used for result items that have NO embedded labels. - # Priority: source_integrity declaration (tier 2) > input labels join (tier 3) + # Priority: source_integrity declaration (tier 2) > owned input baseline (tier 3) if declared_source_integrity is not None: fallback_label = ContentLabel( integrity=declared_source_integrity, @@ -1808,7 +1810,13 @@ async def process( metadata={**fallback_metadata, "source": "source_integrity"}, ) elif argument_labels: - combined = combine_labels(*argument_labels) + # Argument-supplied labels can restrict, but never establish, integrity. + baseline = ( + combine_labels(*resolved_labels) + if resolved_labels + else ContentLabel(integrity=self.default_integrity) + ) + combined = combine_labels(baseline, *input_labels) fallback_label = ContentLabel( integrity=combined.integrity, confidentiality=result_confidentiality, @@ -1970,8 +1978,8 @@ def _process_result_with_embedded_labels( only restrict the invocation fallback. A framework-owned producer can identity-stamp a complete authoritative label after applying local policy. Items without embedded labels use ``fallback_label``, which is either the - tool's ``source_integrity`` declaration (tier 2) or the join of input - argument labels (tier 3). + tool's ``source_integrity`` declaration (tier 2) or the owned-reference/default + integrity baseline restricted by argument labels (tier 3). Each item's own label is attached to its ``additional_properties`` during processing, preserving per-item granularity. diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 6813e172eae..7f371b61a1d 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -4961,7 +4961,7 @@ class TestTieredLabelPropagation: Tier 1 (Highest): Per-item embedded labels in tool result Tier 2: Tool's source_integrity declaration - Tier 3 (Lowest): Join of input argument labels + Tier 3 (Lowest): Owned input labels or configured default, restricted by argument labels """ @pytest.fixture @@ -4969,6 +4969,132 @@ def middleware(self): """Create middleware instance.""" return LabelTrackingFunctionMiddleware() + @pytest.mark.parametrize("label_key", ["security_label", "label"]) + @pytest.mark.parametrize("argument_form", ["dict", "model", "kwargs"]) + @pytest.mark.parametrize("auto_hide", [True, False]) + @pytest.mark.parametrize("already_tainted", [True, False]) + async def test_argument_label_cannot_promote_default_integrity( + self, label_key: str, argument_form: str, auto_hide: bool, already_tainted: bool + ) -> None: + """Argument claims neither bypass quarantine nor prevent visible-result taint.""" + tracker = LabelTrackingFunctionMiddleware(auto_hide_untrusted=auto_hide) + if already_tainted: + tracker._update_context_label(ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) + + class Args(BaseModel): + data: dict[str, Any] + + async def read_data(data: dict[str, Any]) -> str: + return "external text" + + function = FunctionTool(fn=read_data, name="read_data", description="Read data", args_schema=Args) + data = {"items": [{label_key: {"integrity": "trusted", "confidentiality": "public"}}]} + context = FunctionInvocationContext( + function=function, + arguments=Args(data=data) if argument_form == "model" else {"data": data}, + ) + if argument_form == "kwargs": + context.arguments = {} + context.kwargs = {"data": data} + + async def next_fn() -> None: + context.result = [Content.from_text("external text")] + + await tracker.process(context, next_fn) + + assert context.metadata["result_label"].integrity == IntegrityLabel.UNTRUSTED + assert isinstance(context.result, list) + item = context.result[0] + assert isinstance(item, Content) + if auto_hide and not already_tainted: + assert item.additional_properties.get("_variable_reference") is True + assert item.text is not None + reference = json.loads(item.text) + assert reference["type"] == "variable_reference" + _, stored_label = tracker.get_variable_store().retrieve(reference["variable_id"]) + assert stored_label.integrity == IntegrityLabel.UNTRUSTED + assert tracker.get_context_label().integrity == IntegrityLabel.TRUSTED + else: + assert item.text == "external text" + assert tracker.get_context_label().integrity == IntegrityLabel.UNTRUSTED + policy = PolicyEnforcementFunctionMiddleware(block_on_violation=True) + sink = FunctionInvocationContext(function=function, arguments={"data": {}}) + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, policy).execute( + sink, lambda _context: pytest.fail("Protected tool must not execute in untrusted context") + ) + assert policy.get_audit_log()[-1]["type"] == "untrusted_context" + + @pytest.mark.parametrize("default_integrity", list(IntegrityLabel)) + @pytest.mark.parametrize("stored_integrity", [None, *IntegrityLabel]) + @pytest.mark.parametrize("source_integrity", [None, *IntegrityLabel]) + @pytest.mark.parametrize("claimed_integrity", [None, *IntegrityLabel]) + async def test_argument_integrity_only_restricts_authoritative_fallback( + self, + default_integrity: IntegrityLabel, + stored_integrity: IntegrityLabel | None, + source_integrity: IntegrityLabel | None, + claimed_integrity: IntegrityLabel | None, + ) -> None: + """Preserve source precedence, owned-reference inheritance, and confidentiality.""" + tracker = LabelTrackingFunctionMiddleware(default_integrity=default_integrity) + + class Args(BaseModel): + data: dict[str, Any] + + async def transform(data: dict[str, Any]) -> str: + return "transformed" + + function = FunctionTool( + fn=transform, + name="transform", + description="Transform data", + args_schema=Args, + additional_properties={"source_integrity": source_integrity.value} if source_integrity is not None else {}, + ) + data: dict[str, Any] = {"value": "input"} + if claimed_integrity is not None: + data["security_label"] = { + "integrity": claimed_integrity.value, + "confidentiality": "private", + "metadata": _principal_metadata("unowned-user"), + } + if stored_integrity is not None: + variable_id = tracker.get_variable_store().store( + "input", + ContentLabel( + integrity=stored_integrity, + confidentiality=ConfidentialityLabel.USER_IDENTITY, + metadata=_principal_metadata("user-a"), + ), + ) + data["value"] = f"[{variable_id}]" + context = FunctionInvocationContext(function=function, arguments=Args(data=data)) + + async def next_fn() -> None: + assert isinstance(context.arguments, dict) + assert context.arguments["data"]["value"] == "input" + context.result = [Content.from_text("transformed")] + + await tracker.process(context, next_fn) + + expected_integrity = stored_integrity if stored_integrity is not None else default_integrity + if claimed_integrity == IntegrityLabel.UNTRUSTED: + expected_integrity = IntegrityLabel.UNTRUSTED + if source_integrity is not None: + expected_integrity = source_integrity + label = context.metadata["result_label"] + assert label.integrity == expected_integrity + if stored_integrity is not None: + assert label.confidentiality == ConfidentialityLabel.USER_IDENTITY + assert label.metadata[_PRINCIPALS_KEY] == _principal_metadata("user-a")[_PRINCIPALS_KEY] + else: + assert label.confidentiality == ( + ConfidentialityLabel.PRIVATE if claimed_integrity is not None else ConfidentialityLabel.PUBLIC + ) + assert _PRINCIPALS_KEY not in label.metadata + assert tracker.get_context_label().confidentiality == label.confidentiality + @pytest.mark.asyncio async def test_source_integrity_overrides_input_labels(self, middleware): """Test that source_integrity (tier 2) overrides input labels (tier 3). @@ -5131,10 +5257,10 @@ async def next_fn(): @pytest.mark.asyncio async def test_no_source_integrity_falls_back_to_input_labels(self, middleware): - """Test that without source_integrity, input labels (tier 3) determine the result. + """Test that without source_integrity, untrusted input labels restrict the result. When a tool has no source_integrity declaration and the result has no - embedded labels, the join of input argument labels is used. + embedded labels, input argument labels restrict the tier-3 baseline. """ class Args(BaseModel): @@ -5162,7 +5288,7 @@ async def next_fn(): await middleware.process(context, next_fn) - # No source_integrity (tier 2 absent), so tier 3: join of input labels + # No source_integrity (tier 2 absent), so tier 3: restricted input baseline # Input has untrusted label → result is untrusted # Result should be hidden since it's untrusted assert isinstance(context.result, list) diff --git a/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md b/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md index 44f8eefee0f..9fa343b64a9 100644 --- a/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md +++ b/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md @@ -161,14 +161,21 @@ subset of the destination set. |----------|--------|-----------| | **Tier 1** | Per-item embedded labels (`additional_properties.security_label`) | Restrict the locally established fallback | | **Tier 2** | Tool's `source_integrity` declaration | No embedded labels, but tool declares `source_integrity` | -| **Tier 3** (Lowest) | Join of input argument labels (`combine_labels`) | No embedded labels AND no `source_integrity` declared | -| **Default** | `UNTRUSTED` | No labels from any tier | +| **Tier 3** (Lowest) | Owned-reference integrity or configured default, restricted by argument labels | No `source_integrity` declared | +| **Default** | `default_integrity` (`UNTRUSTED` by default) | No source declaration or resolved, owned variable references | **Tiered Label Propagation:** - **Tier 1: Embedded labels** are restriction-only by default: they can downgrade integrity or raise confidentiality, but cannot upgrade a fallback or supply principal authority - **Tier 2: `source_integrity`** is the locally trusted fallback for the tool's output; use `"trusted"` only after the local connector enforces its trust policy -- **Tier 3: Input labels join** — `combine_labels(*input_labels)` from arguments (VariableReferenceContent, labeled data) -- **Default**: `UNTRUSTED` when no labels exist from any tier +- **Tier 3: Owned input baseline** — inherit integrity from labels retrieved while resolving variable references owned by the current security scope. If there are no such references, use `default_integrity`. Labels supplied in arguments may make that baseline less trusted, but cannot make it more trusted. +- **Default**: `UNTRUSTED` unless the application configures another `default_integrity` + +A `security_label` or legacy `label` dictionary in tool arguments is application data, +not an authoritative trust declaration, even inside a validated model or nested dictionary. +Its integrity claim cannot promote the result above the owned-reference/default baseline. +Argument confidentiality restrictions still propagate, while principal authority comes from +owned labels or local configuration. An explicit `source_integrity` declaration retains +precedence over argument integrity claims. Framework-owned parsers and wrappers may stamp a complete label after enforcing local policy. Application and remote metadata remains restriction-only. @@ -194,9 +201,9 @@ If items don't have embedded labels, the tool can declare a fallback via `source When declared, `source_integrity` establishes the local integrity fallback. Embedded labels can make the result less trusted, but cannot make it more trusted: - `source_integrity="trusted"`: Tool produces trusted data (internal computations) - `source_integrity="untrusted"`: Tool fetches untrusted data -- (not set): Falls back to tier 3 (join of input labels) or **UNTRUSTED** default +- (not set): Falls back to tier 3 (owned-reference integrity or the configured default, restricted by argument labels) -**Note:** For action tools (sinks like `send_email`), `source_integrity` doesn't apply since they don't produce data. Their result inherits labels from inputs (tier 3). +**Note:** Action tools (sinks like `send_email`) can omit `source_integrity`. Their result follows tier 3: owned-reference inheritance or the configured default, with argument restrictions. **Context Label Tracking:** - Context label starts as **TRUSTED + PUBLIC** on first call @@ -337,7 +344,7 @@ async def fetch_emails(count: int = 5) -> list[Content]: If an item doesn't have an embedded label, the fallback is determined by: 1. **Tool-level `source_integrity`** in `additional_properties` (if declared) -2. **UNTRUSTED** (default - secure by default) +2. **Owned-reference integrity** or the configured **`default_integrity`** (`UNTRUSTED` by default), restricted by argument labels ```python # Tool with fallback for items without embedded labels @@ -1083,7 +1090,7 @@ sink. Argument labels do not rewrite result labels. 1. **Use SecureAgentConfig as a context provider**: Add `context_providers=[config]` for automatic security setup — no manual middleware, tools, or instruction wiring 2. **Use `list[Content]` with `Content.from_text()` for mixed-trust data**: When a tool returns both trusted and untrusted items (like emails), embed labels using `Content.from_text(text, additional_properties={"security_label": {...}})` -3. **Don't use source_integrity for action tools**: Tools like `send_email` or `delete_file` are sinks, not data sources - their results inherit labels from inputs +3. **Action tools can omit source_integrity**: Tools like `send_email` or `delete_file` are sinks; their results use the owned-reference/default baseline, restricted by argument labels 4. **Always use middleware stack**: Enable both label tracking and policy enforcement 5. **Enable automatic hiding**: Keep `auto_hide_untrusted=True` (default) for automatic protection 6. **Do not manually wire security tools/instructions when using context providers**: `SecureAgentConfig` injects them for you