[security] Mellea Taint Analysis

[nrfulton]

Overall Problem Statement

Some systems co-mingle data of different privilege levels in prompts. This can result in security bugs. The resulting bugs are typically in one of two major classes:

1. data exfiltration: An LLM runs at a higher privilege level than the end-user, which results in either intentional or unintentional leaking of data. Example: Agent has access to an HR database, and employee A is able to extract information about employee B.
2. injection: The output of an LLM is used to take some action, and the action taken is unduly influenced by input from an untrusted source. Example: values from a customer are passed through to an exec call without explicit santitization.

Ambrish's PoC

@ambrishrawat is implementing a taint analysis for Mellea programs at ambrishrawat:security_minimal.

He introduces the following mechanisms:

First, a SecurityLevel: SAFE | TAINTED | PROPRIETARY | SANITIZED is attached to each CBlock | Component.

Next, he adds a @privileged decorator, which rejects unsafe CBlock. he also adds a sanitize : CBlock -> ... function which mutates its input.

Finally, he adds .mark_tainted(level) and .is_safe() to the CBlock class, both of which modify the CBlock.

Here's his example:

def get_summary(email_body: str) -> str:
	return "summary"

@privileged
def open_url(url: str) -> str:
	return "summary"
	
email = CBlock(raw_email).mark_tainted("TAINTED")
summary = get_summary(m, email)            # allowed – propagates taint
try:
    open_url(summary)                      # blocked – raises SecurityError
except SecurityError:
    safe_url = sanitize(summary)           # mark as safe (mock)
    open_url(safe_url)                     # now allowed

He also asks:

1. Implicit str <-> CBlock conversions bypass checks; so unable to get tainted Cblocks via m.instruct or within the args of tool calls. Any idea on how to address this?

Problems with the PoC

Mutation

The PoC mutates CBlocks.

Our goal is for CBlocks and Components to be immutable outside of unsafe and highly audited portions of the code-base (primarily backend/inference engine implementations). If we force CBlocks and Components to be immutable, then you need to know the taint "level" at construction-time.

The downside of immutability: de-classification becomes a bit of a pain. Creating a new CBlock with a lower permission level isn't enough - you have to go through every Component that contains that CBlock.

SecurityLevel is too coarse-grained

When adding this feature, we might as well include integration is capability systems. If we instead make SecurityLevel this type:

SecLevel := None 
                    | Classified of CapabilityType 
                    | TaintedBy of (CBlock | Component)

then we can define various capability schemes, or, more importantly, "plug into" cloud-based IAMs. Using TaintedBy seems like a good idea to me, for reasons I don't have time to elaborate on but are maybe hinted at in the mutation concern.

Problems with Mellea

We need to start tracking parts() again. Ideally this should be automatic. @jakelorocco we should talk about this, because forgetting to add stuff to your parts() is going to be a footgun moving forward. Ideally we should do this with meta-programming or a superclass constructor.

[ambrishrawat]

• Would the intent for SecLevel (or equivalent) be to stay attached to each CBlock/Component as part of their immutable metadata, or managed independently/globally? What CapabilityType do you have in mind?
• Should immutability come from "functional updates" (like a with_meta(...) returning a new object) rather than mutating methods? Or would a subclassed SecureCBlock express this separation better?

[nrfulton]

Would the intent for SecLevel (or equivalent) be to stay attached to each CBlock/Component as part of their immutable metadata

Yes. A security level is assigned at construction. That security level is immutable. We should discuss how that initial assignment happens; likely, it's either an input to the act function or part of the post-processing (ModelOutputThunk.parsed_repr).

To de-classify, you must create a new CBlock.

Existing Components that were constructed from the original higher-permissioned CBlock will still be at the higher permission level, and must be explicitly declassified.

Declassifying a Component that has classified components should result in a critical warning, at least. The actual handling of that situation declassification needs consideration. We could walk the dependency graph (using Component.parts() and de-classify. Or, we could maintain the existing constituent parts and make the permission check shallow. The latter probably makes more sense for explicit downgrades.

What CapabilityType do you have in mind?

For now it can just be a trait. I am not exactly sure what the interface for this should be, but probably something like:

class CapabilityType(Generic[T], abc.ABC): # not a huge fan of the name; feel free to suggest others. Taint?
@AbstractMethod
def has_access(self, entitlement: T | None) -> bool: ...

where can_access takes as input an entitlement and provides as output whether the given entitlement should have access to the tainted information.

You can get back to your original implementation by setting T to any type and defining has_access as a constant function. The reason for future-proofing now: you probably eventually want to use an iam system to manage access; in that case, the user's role or system's iam identifier is the entitlement.

We could discuss whether it makes sense to also add def is_untrusted_llm_output, but I think that's already implied anyways by a thing being a ModelOutputThunk. And sanitiziation should happen at the point in the code where the value is going to be used anyways, right? So I don't think storing that information and marshalling it around has much value add.

[ambrishrawat]

Following up on the discussion, I’ve implemented an updated prototype integrating security propagation into ModelOutputThunk and the broader mellea flow.

New branch ambrishrawat:security_poc

Key changes

1. ModelOutputThunk.from_generation now constructs properly secured instances, ensuring immutability of security critical meta fields.
   • Backends - ollama, litellm, openai for now - have been updated so their *generate* methods call this.
2. Propagation logic is in a new utility taint_source in mellea/security/core.py tracks taint shallowly by context and recursively via parts for actions (which has been updated in instruction.py).
3. Security model now is based on refined SecLevel (as you had described) and an abstract AccessType (suggested name instead of CapabilityType?) for future entitlement-based control
4. declassify is currently a lightweight helper - I agree that the actual sanitization/declassification should be deferred to execution time.
5. Potentially breaking changes are in the signature updates in instruct within mfuncs and session to allow taint propagation through session.instruct; it now accepts both str and Cblock as description

Example test:

tainted_desc = CBlock("Process this sensitive data")
tainted_desc.mark_tainted()
session = MelleaSession(OllamaModelBackend("llama3.2"))

result = session.instruct(description=tainted_desc)

assert not result.is_safe()
print("✅ SUCCESS: Taint preserved!")

[nrfulton]

Cool! Looks like we're starting to converge on something. Would you mind opening up a PR against generative-computing/mellea:main?

You can mark it as a draft. But now that we're down to implementation vagaries it's far easier to have discussions using the code review tool on PRs than by copy/pasting code blocks into comments on a discussion :)

[ambrishrawat]

Done #197