feat: taint checking and security by ambrishrawat · Pull Request #197 · generative-computing/mellea

ambrishrawat · 2025-10-14T19:21:17Z

This PR introduces a minimal proof-of-concept for taint and security propagation across CBlock, ModelOutputThunk, and session flows, as discussed in generative-computing/mellea#189
.

mergify · 2025-10-14T19:21:52Z

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Enforce conventional commit

Wonderful, this rule succeeded.

Make sure that we follow https://www.conventionalcommits.org/en/v1.0.0/

title ~= ^(fix|feat|docs|style|refactor|perf|test|build|ci|chore|revert|release)(?:\(.+\))?:

ambrishrawat · 2025-10-17T14:15:05Z

@nrfulton quick clarifications -

What’s the best way for expose taint configuration to devs? e.g. when a description includes a user variable like summarise the following {{email_body}}, should taint be inferred automatically or something they can configure?
Would it make sense to have a global strictness setting to toggle between warnings and exceptions for taint violations? Is blocify the best place for this?

nrfulton · 2025-10-22T19:55:22Z

What’s the best way for expose taint configuration to devs? e.g. when a description includes a user variable like summarise the following {{email_body}}, should taint be inferred automatically or something they can configure?

We should infer automatically where-ever possible. I nthis case, I'm not sure how you would infer taint. I guess you assumption here is that email_boy -- or any user_variable input -- should entail taint?

ambrishrawat · 2025-10-23T12:47:18Z

Yes, that was the thinking. Making it configurable may make more sense for taint. Any thoughts on the best way to expose that? Would love your take on the code too.

davidcox · 2025-10-23T13:24:53Z

If there is a tainted variable in the context, everything downstream should get tainted. As for how variables get tainted in the first place, a common way people do this is to define sources, sinks, and (optionally) washers. These are wrappers around interfaces that produce sensitive data (e.g. HR database api), or where it enters an unsafe place (e.g. sending to a UI).

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

nrfulton · 2025-11-21T01:18:11Z

docs/dev/taint_analysis.md

+component = CBlock("user input")
+component.mark_tainted()  # Sets SecLevel.tainted_by(component)


CBlocks should be immutable.

Naming a variable Ccmponent and assigning it to a CBlock is confusing.

The cyclic reference here is a bit confusing and invites buggy code. Use tainted_by(None) instead of tainted_by(self) for the root node.

c = CBlock("user input", sec_level=SecLevel.tained_by(None))

Updated this; tainted_by(None) for root now

nrfulton · 2025-11-21T01:19:11Z

docs/dev/taint_analysis.md

+component = CBlock("user input")
+component.mark_tainted()  # Sets SecLevel.tainted_by(component)
+
+if component._meta["_security"].is_tainted():


Why not c.sec_level?

Defined it as a property and now this works as c.sec_level.is_tainted()

nrfulton · 2025-11-21T01:20:25Z

docs/examples/security/taint_example.py

+print(f"Original CBlock is tainted: {not tainted_desc.is_safe()}")
+
+# Create session
+session = MelleaSession(OllamaModelBackend("llama3.2"))


Unless the example critically depends on using a particular model, always use session = start_session() instead. This makes the examples easier to maintain.

nrfulton · 2025-11-21T01:21:07Z

docs/examples/security/taint_example.py

+
+# The result should be tainted
+print(f"Result is tainted: {not result.is_safe()}")
+if not result.is_safe():


We should use is_tainted instead of is_safe. The meaning of safe is very ambiguous.

Removed all instances of is_safe

nrfulton · 2025-11-21T01:22:19Z

mellea/security/core.py

+        Returns:
+            The CBlock or Component that tainted this content, or None
+        """
+        if self.level_type == "tainted_by":


Especially in a module called security.core, we should avoid use of magic strings.

Created SecLevelType enum

nrfulton · 2025-11-21T01:24:24Z

mellea/security/core.py

+            sources.append(action)
+
+    # For Components, check their constituent parts for taint
+    if hasattr(action, 'parts'):


Instead us something like:

match action: case Component... case CBlock...

(If type(action) :> Component then check is not necessary because the Component protocol has a parts() method. )

Updated it to use match/case

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

ambrishrawat · 2025-11-25T13:27:04Z

Thanks for the review @nrfulton !
I have incorporated your suggestions. Appreciate another pass when you get the chance

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

guicho271828 · 2025-11-26T15:50:04Z

hi ambrish!

mergify · 2026-01-28T09:45:05Z

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Enforce conventional commit

Wonderful, this rule succeeded.

Make sure that we follow https://www.conventionalcommits.org/en/v1.0.0/

title ~= ^(fix|feat|docs|style|refactor|perf|test|build|ci|chore|revert|release)(?:\(.+\))?:

mergify · 2026-02-01T14:47:04Z

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Enforce conventional commit

Wonderful, this rule succeeded.

Make sure that we follow https://www.conventionalcommits.org/en/v1.0.0/

title ~= ^(fix|feat|docs|style|refactor|perf|test|build|ci|chore|revert|release)(?:\(.+\))?:

mergify · 2026-02-08T23:07:56Z

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Enforce conventional commit

Wonderful, this rule succeeded.

Make sure that we follow https://www.conventionalcommits.org/en/v1.0.0/

title ~= ^(fix|feat|docs|style|refactor|perf|test|build|ci|chore|revert|release)(?:\(.+\))?:

mergify · 2026-02-12T14:52:27Z

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Enforce conventional commit

Wonderful, this rule succeeded.

Make sure that we follow https://www.conventionalcommits.org/en/v1.0.0/

title ~= ^(fix|feat|docs|style|refactor|perf|test|build|ci|chore|revert|release)(?:\(.+\))?:

planetf1 · 2026-02-25T14:55:02Z

mellea/security/core.py

+    def wrapper(*args, **kwargs):
+        # Check all arguments for marked content (tainted or classified)
+        for arg in args:
+            if isinstance(arg, TaintChecking):


I'm wondering if there's an issue with here in not looking recursively (ie via taint_sources() )? the args may be ok, but we could have a tainted Component?

Ah I see that, yes need to call taint_sources here and do a recursive check for classified sources as well. I will update this.

mergify · 2026-02-25T14:55:38Z

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Enforce conventional commit

Wonderful, this rule succeeded.

Make sure that we follow https://www.conventionalcommits.org/en/v1.0.0/

title ~= ^(fix|feat|docs|style|refactor|perf|test|build|ci|chore|revert|release)(?:\(.+\))?:

planetf1 · 2026-02-25T15:01:33Z

mellea/security/core.py

+
+
+def declassify(cblock: "CBlock") -> "CBlock":
+    """Create a declassified version of a CBlock (non-mutating).


Only CBlocks or any subclass? If the latter, we return a CBlock which would lose the subclass info?
For example could it be an ImageBlock?
Also note that ImageBlock never calls super().init() which means the value that is being copied does not exist == crash?

Right, this particular declassify is merely illustrative, hence only applicable to CBlocks. One would need bespoke declassification methods for applications to sparsify taint information across the code. That was my initial thinking

mergify · 2026-02-27T14:07:19Z

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Enforce conventional commit

Wonderful, this rule succeeded.

Make sure that we follow https://www.conventionalcommits.org/en/v1.0.0/

title ~= ^(fix|feat|docs|style|refactor|perf|test|build|ci|chore|revert|release)(?:\(.+\))?:

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

mergify · 2026-03-11T12:15:55Z

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Enforce conventional commit

Wonderful, this rule succeeded.

Make sure that we follow https://www.conventionalcommits.org/en/v1.0.0/

title ~= ^(fix|feat|docs|style|refactor|perf|test|build|ci|chore|revert|release)(?:\(.+\))?:

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

ambrishrawat · 2026-03-11T12:29:57Z

@planetf1 Updated the logic for privileged. It recursively checks for both taint and classified.
cc: @nrfulton @jakelorocco

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

mergify · 2026-03-25T10:48:07Z

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Enforce conventional commit

Wonderful, this rule succeeded.

Make sure that we follow https://www.conventionalcommits.org/en/v1.0.0/

title ~= ^(fix|feat|docs|style|refactor|perf|test|build|ci|chore|revert|release)(?:\(.+\))?:

planetf1

I think there are also some

uses of Any where the types should be stricter
incorrect types
missing type annotations

I know we have omissions elsewhere, but it feels that the focus on security here makes it worth paying close attention.

Need to have a fiddle with mypy/py (the default rules are not strict enough) to be sure

planetf1 · 2026-03-25T13:39:00Z

mellea/core/base.py

+        sec_level: Any = None,
    ):
        """Initialize ModelOutputThunk with an optional pre-computed value and metadata."""
        super().__init__(value, meta)


Suggested change

super().__init__(value, meta)

I think this is left over from conflict resolution as it's a double init() ?

Added sec_level here

planetf1 · 2026-03-25T13:41:27Z

mellea/core/requirement.py

            List of constituent components. Empty by default; subclasses override
            to expose their internal structure.
        """
-        return []


We'll now return none -- inconsistent with annotations, and likely to cause a crash?

planetf1 · 2026-03-25T13:44:19Z

mellea/backends/huggingface.py

@@ -838,6 +844,11 @@ async def _generate_from_context_with_kv_cache(

            output = ModelOutputThunk(None)


left over from conflict since we're assigning a few lines later. And what about assigning _start (not checked if we do that much further down). We'd lose it?

planetf1 · 2026-03-25T13:45:59Z

mellea/backends/huggingface.py

@@ -983,6 +994,11 @@ async def _generate_from_context_standard(

            output = ModelOutputThunk(None)


is this another dup ? similar -- worth checking for this pattern through the code?

planetf1 · 2026-03-25T13:47:50Z

mellea/core/base.py

        meta: dict[str, Any] | None = None,
        parsed_repr: S | None = None,
        tool_calls: dict[str, ModelToolCall] | None = None,
+        sec_level: Any = None,


As per Nathan's comment above there's another example. I think we should be explicit - especially given it's security. May be worth doing a more thorough check on type annotations and use of Any?

planetf1 · 2026-03-25T13:50:31Z

mellea/security/core.py

+                    nested = _collect_sources_by_predicate(item, None, predicate)
+                    sources.extend(nested)
+        except Exception:
+            pass


what else could happen here? Do we need more checks for other exceptions? Or log something at least? I guess we're trying to be non-intrusive (need to understand code better to suggest exact changes!)

planetf1 · 2026-03-25T13:51:33Z

mellea/core/base.py


 @runtime_checkable
-class Component(Protocol, Generic[S]):
+class Component(TaintChecking, Protocol, Generic[S]):


so any contribs/older code outside mellea would be impacted -- should we ensure we capture in release notes/doc as a breaking change

planetf1 · 2026-03-25T13:57:50Z

mellea/stdlib/components/mify.py

+                lambda self: getattr(self, "_sec_level", None),
+                doc="Get the security level for this Component.",
+            )
+            setattr(obj, "sec_level", sec_level_prop)


this works on a class but fails on an instance - other parts of the code have checks for class vs instance to do the right thing, but this is omitted here

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

jakelorocco · 2026-03-25T19:02:19Z

I've started taking a look at this. Will finish my review soon.

ambrishrawat marked this pull request as draft October 14, 2025 19:21

nrfulton self-requested a review October 15, 2025 16:54

ambrishrawat force-pushed the security_poc branch from 4798f64 to d1b09e1 Compare November 11, 2025 19:24

ambrishrawat marked this pull request as ready for review November 12, 2025 11:11

ambrishrawat added 9 commits November 14, 2025 15:34

version 2 of taint tracking

7d4e73e

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

version 2 of taint tracking

4c9ab68

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

taint tracking updates for ollama and litellm

daaf5ee

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

docs taint analysis

2fd7bc0

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

restored formatter to original

0efbee6

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

removed redundant sanitise helper

5e1a266

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

updated taint analysis dev docs

4ea932a

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

updated taint analysis dev docs

b0a23a5

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

updated taint analysis dev docs

5466c31

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

ambrishrawat force-pushed the security_poc branch from d5f8033 to 5466c31 Compare November 14, 2025 15:35

ambrishrawat and others added 3 commits November 18, 2025 16:01

Merge branch 'generative-computing:main' into security_poc

5f85458

Merge branch 'generative-computing:main' into security_poc

0075de9

Merge branch 'main' into security_poc

7876e62

nrfulton requested changes Nov 21, 2025

View reviewed changes

ambrishrawat added 2 commits November 25, 2025 12:17

updates based on PR feedback

c303b87

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

added tests for taint_sources of a Component

f85f953

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

minor doc updates to remove the use of word safe

c1062e5

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

Merge branch 'main' into security_poc

13a853e

nrfulton self-requested a review December 2, 2025 02:15

Merge branch 'generative-computing:main' into security_poc

867e6c8

Merge branch 'main' into security_poc

fdcbc35

Merge branch 'generative-computing:main' into security_poc

24e480f

Merge branch 'main' into security_poc

a55a63b

planetf1 reviewed Feb 25, 2026

View reviewed changes

Merge branch 'main' into security_poc

f305ddd

ambrishrawat requested review from a team and jakelorocco as code owners February 25, 2026 18:12

Merge branch 'main' into security_poc

6601376

added recurssive logic for classified and taint for privlege decorator

470f47b

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

merged with main

33393bd

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

ambrishrawat and others added 4 commits March 11, 2026 13:59

mypy error corrected for sec_level in Component

1df8ded

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

resolved cnflicts and merged with main

7ae5ad9

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

Merge branch 'generative-computing:main' into security_poc

e066a00

resolved cnflicts and merged with main

e0a2b4b

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

ambrishrawat requested a review from planetf1 March 25, 2026 11:40

planetf1 requested changes Mar 25, 2026

View reviewed changes

corrections for sec_level types and hf backend

92c094e

Signed-off-by: Ambrish Rawat <ambrish.rawat@ie.ibm.com>

github-actions bot added the enhancement New feature or request label Mar 25, 2026

		component = CBlock("user input")
		component.mark_tainted() # Sets SecLevel.tainted_by(component)



		def declassify(cblock: "CBlock") -> "CBlock":
		"""Create a declassified version of a CBlock (non-mutating).

		@@ -838,6 +844,11 @@ async def _generate_from_context_with_kv_cache(

		output = ModelOutputThunk(None)

		@@ -983,6 +994,11 @@ async def _generate_from_context_standard(

		output = ModelOutputThunk(None)

Conversation

ambrishrawat commented Oct 14, 2025

Uh oh!

mergify bot commented Oct 14, 2025

Merge Protections

🟢 Enforce conventional commit

Uh oh!

ambrishrawat commented Oct 17, 2025

Uh oh!

nrfulton commented Oct 22, 2025

Uh oh!

ambrishrawat commented Oct 23, 2025

Uh oh!

davidcox commented Oct 23, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ambrishrawat commented Nov 25, 2025

Uh oh!

guicho271828 commented Nov 26, 2025

Uh oh!

mergify bot commented Jan 28, 2026

Merge Protections

🟢 Enforce conventional commit

Uh oh!

mergify bot commented Feb 1, 2026

Merge Protections

🟢 Enforce conventional commit

Uh oh!

mergify bot commented Feb 8, 2026

Merge Protections

🟢 Enforce conventional commit

Uh oh!

mergify bot commented Feb 12, 2026

Merge Protections

🟢 Enforce conventional commit

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ambrishrawat Feb 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Feb 25, 2026

Merge Protections

🟢 Enforce conventional commit

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Feb 27, 2026

Merge Protections

🟢 Enforce conventional commit

Uh oh!

mergify bot commented Mar 11, 2026

Merge Protections

🟢 Enforce conventional commit

Uh oh!

ambrishrawat commented Mar 11, 2026

ambrishrawat Feb 25, 2026 •

edited

Loading