ComfyUI/tests/test_invariant_execution.py

"""
Regression tests for the credential-injection trust gate in execution.py.

These tests verify that auth_token_comfy_org and api_key_comfy_org are only
forwarded to node classes that appear in the immutable _TRUSTED_NODE_CLASSES
set, and are withheld from all other nodes (including those that spoof their
__module__ name).
"""
import types
import pytest
import execution as _exec


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_class(name: str, module: str) -> type:
    """Create a minimal node class with a controlled __module__."""
    cls = type(name, (), {})
    cls.__module__ = module
    return cls


def _run_trust_check(class_def: type) -> bool:
    """Thin wrapper around the private helper exposed by execution.py."""
    return _exec._is_trusted_node(class_def)


# ---------------------------------------------------------------------------
# _is_trusted_node
# ---------------------------------------------------------------------------

class TestIsTrustedNode:
    def test_unknown_class_is_not_trusted(self):
        """A freshly created class that was never registered must not be trusted."""
        cls = _make_class("FakeNode", "some_third_party_package")
        assert _run_trust_check(cls) is False

    def test_spoofed_module_name_is_not_trusted(self):
        """
        A class whose __module__ starts with 'comfy_api_nodes' but was NOT
        imported from that package must not be trusted.
        This is the core regression: the old string-prefix check would have
        returned True here.
        """
        spoofed = _make_class("EvilNode", "comfy_api_nodes.evil.payload")
        assert _run_trust_check(spoofed) is False, (
            "Spoofed __module__ name must not grant trust — "
            "only classes in _TRUSTED_NODE_CLASSES are trusted."
        )

    def test_another_spoofed_variant_is_not_trusted(self):
        """Variant: exact prefix match should still be rejected."""
        spoofed = _make_class("EvilNode2", "comfy_api_nodes")
        assert _run_trust_check(spoofed) is False

    def test_trusted_set_is_frozenset(self):
        """_TRUSTED_NODE_CLASSES must be immutable (frozenset)."""
        assert isinstance(_exec._TRUSTED_NODE_CLASSES, frozenset), (
            "_TRUSTED_NODE_CLASSES must be a frozenset so it cannot be mutated at runtime."
        )


# ---------------------------------------------------------------------------
# Credential gating via _is_trusted_node (integration-style)
# ---------------------------------------------------------------------------

class TestCredentialGating:
    """
    Verify that the _is_trusted flag correctly gates credential injection
    in the two code paths inside get_input_data (v3 hidden inputs and legacy
    hidden inputs).
    """

    def _make_extra_data(self) -> dict:
        return {
            "auth_token_comfy_org": "secret-token-abc",
            "api_key_comfy_org": "secret-key-xyz",
        }

    def test_untrusted_node_does_not_receive_credentials(self):
        """
        For any class not in _TRUSTED_NODE_CLASSES, _is_trusted_node returns
        False, so credentials must be None.
        """
        untrusted = _make_class("UntrustedNode", "comfy_api_nodes.spoofed")
        is_trusted = _exec._is_trusted_node(untrusted)
        extra = self._make_extra_data()

        # Simulate the conditional that execution.py applies
        token = extra.get("auth_token_comfy_org", None) if is_trusted else None
        api_key = extra.get("api_key_comfy_org", None) if is_trusted else None

        assert token is None, "Untrusted node must not receive auth_token_comfy_org"
        assert api_key is None, "Untrusted node must not receive api_key_comfy_org"

    def test_trusted_set_membership_gates_credentials(self):
        """
        Only a class that is actually in _TRUSTED_NODE_CLASSES receives
        credentials; identity (not module name) is the gate.
        """
        # Manually inject a sentinel class into a temporary frozenset to
        # simulate what a real comfy_api_nodes class would look like.
        real_trusted = _make_class("RealApiNode", "comfy_api_nodes.real")
        fake_trusted_set = frozenset({real_trusted})

        # Patch _TRUSTED_NODE_CLASSES temporarily
        original = _exec._TRUSTED_NODE_CLASSES
        try:
            _exec._TRUSTED_NODE_CLASSES = fake_trusted_set

            assert _exec._is_trusted_node(real_trusted) is True
            extra = self._make_extra_data()
            token = extra.get("auth_token_comfy_org") if _exec._is_trusted_node(real_trusted) else None
            assert token == "secret-token-abc"

            # A different class with the same module name must still be rejected
            impostor = _make_class("ImpostorNode", "comfy_api_nodes.real")
            assert _exec._is_trusted_node(impostor) is False
            token2 = extra.get("auth_token_comfy_org") if _exec._is_trusted_node(impostor) else None
            assert token2 is None
        finally:
            _exec._TRUSTED_NODE_CLASSES = original