Prevent signature traversal of raw

This commit is contained in:
xmarre 2026-03-16 09:29:00 +01:00
parent 9f9d37bd9a
commit fce22da313
2 changed files with 107 additions and 4 deletions

View File

@ -465,9 +465,9 @@ class CacheKeySetInputSignature(CacheKeySet):
async def get_immediate_node_signature(self, dynprompt, node_id, ancestor_order_mapping):
"""Build the immediate cache-signature fragment for a node.
Link inputs are reduced to ancestor references here, while non-link
values are appended as-is. Full canonicalization happens later in
`get_node_signature()` via `_signature_to_hashable()`.
Link inputs are reduced to ancestor references here. Non-link values
are canonicalized or failed closed before being appended so the outer
node-signature pass never recurses into live prompt input containers.
"""
if not dynprompt.has_node(node_id):
# This node doesn't exist -- we can't cache it.
@ -485,7 +485,7 @@ class CacheKeySetInputSignature(CacheKeySet):
ancestor_index = ancestor_order_mapping[ancestor_id]
signature.append((key,("ANCESTOR", ancestor_index, ancestor_socket)))
else:
signature.append((key, inputs[key]))
signature.append((key, to_hashable(inputs[key])))
return signature
# This function returns a list of all ancestors of the given node. The order of the list is

View File

@ -0,0 +1,103 @@
import asyncio
from comfy_execution import caching
class _StubDynPrompt:
def __init__(self, nodes):
self._nodes = nodes
def has_node(self, node_id):
return node_id in self._nodes
def get_node(self, node_id):
return self._nodes[node_id]
class _StubIsChangedCache:
async def get(self, node_id):
return None
class _StubNode:
@classmethod
def INPUT_TYPES(cls):
return {"required": {}}
def test_get_immediate_node_signature_canonicalizes_non_link_inputs(monkeypatch):
live_value = [1, {"nested": [2, 3]}]
dynprompt = _StubDynPrompt(
{
"1": {
"class_type": "TestCacheNode",
"inputs": {"value": live_value},
}
}
)
monkeypatch.setitem(caching.nodes.NODE_CLASS_MAPPINGS, "TestCacheNode", _StubNode)
monkeypatch.setattr(caching, "NODE_CLASS_CONTAINS_UNIQUE_ID", {})
keyset = caching.CacheKeySetInputSignature(dynprompt, [], _StubIsChangedCache())
signature = asyncio.run(keyset.get_immediate_node_signature(dynprompt, "1", {}))
assert signature == [
"TestCacheNode",
None,
("value", ("list", (1, ("dict", (("nested", ("list", (2, 3))),))))),
]
def test_get_immediate_node_signature_fails_closed_for_opaque_non_link_input(monkeypatch):
class OpaqueRuntimeValue:
pass
live_value = OpaqueRuntimeValue()
dynprompt = _StubDynPrompt(
{
"1": {
"class_type": "TestCacheNode",
"inputs": {"value": live_value},
}
}
)
monkeypatch.setitem(caching.nodes.NODE_CLASS_MAPPINGS, "TestCacheNode", _StubNode)
monkeypatch.setattr(caching, "NODE_CLASS_CONTAINS_UNIQUE_ID", {})
keyset = caching.CacheKeySetInputSignature(dynprompt, [], _StubIsChangedCache())
signature = asyncio.run(keyset.get_immediate_node_signature(dynprompt, "1", {}))
assert signature[:2] == ["TestCacheNode", None]
assert signature[2][0] == "value"
assert type(signature[2][1]) is caching.Unhashable
def test_get_node_signature_never_visits_raw_non_link_input(monkeypatch):
live_value = [1, 2, 3]
dynprompt = _StubDynPrompt(
{
"1": {
"class_type": "TestCacheNode",
"inputs": {"value": live_value},
}
}
)
monkeypatch.setitem(caching.nodes.NODE_CLASS_MAPPINGS, "TestCacheNode", _StubNode)
monkeypatch.setattr(caching, "NODE_CLASS_CONTAINS_UNIQUE_ID", {})
original_impl = caching._signature_to_hashable_impl
def guarded_impl(obj, *args, **kwargs):
if obj is live_value:
raise AssertionError("raw non-link input reached outer signature canonicalizer")
return original_impl(obj, *args, **kwargs)
monkeypatch.setattr(caching, "_signature_to_hashable_impl", guarded_impl)
keyset = caching.CacheKeySetInputSignature(dynprompt, [], _StubIsChangedCache())
signature = asyncio.run(keyset.get_node_signature(dynprompt, "1"))
assert isinstance(signature, tuple)