mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-03-20 08:33:44 +08:00
Prevent signature traversal of raw
This commit is contained in:
parent
9f9d37bd9a
commit
fce22da313
@ -465,9 +465,9 @@ class CacheKeySetInputSignature(CacheKeySet):
|
|||||||
async def get_immediate_node_signature(self, dynprompt, node_id, ancestor_order_mapping):
|
async def get_immediate_node_signature(self, dynprompt, node_id, ancestor_order_mapping):
|
||||||
"""Build the immediate cache-signature fragment for a node.
|
"""Build the immediate cache-signature fragment for a node.
|
||||||
|
|
||||||
Link inputs are reduced to ancestor references here, while non-link
|
Link inputs are reduced to ancestor references here. Non-link values
|
||||||
values are appended as-is. Full canonicalization happens later in
|
are canonicalized or failed closed before being appended so the outer
|
||||||
`get_node_signature()` via `_signature_to_hashable()`.
|
node-signature pass never recurses into live prompt input containers.
|
||||||
"""
|
"""
|
||||||
if not dynprompt.has_node(node_id):
|
if not dynprompt.has_node(node_id):
|
||||||
# This node doesn't exist -- we can't cache it.
|
# This node doesn't exist -- we can't cache it.
|
||||||
@ -485,7 +485,7 @@ class CacheKeySetInputSignature(CacheKeySet):
|
|||||||
ancestor_index = ancestor_order_mapping[ancestor_id]
|
ancestor_index = ancestor_order_mapping[ancestor_id]
|
||||||
signature.append((key,("ANCESTOR", ancestor_index, ancestor_socket)))
|
signature.append((key,("ANCESTOR", ancestor_index, ancestor_socket)))
|
||||||
else:
|
else:
|
||||||
signature.append((key, inputs[key]))
|
signature.append((key, to_hashable(inputs[key])))
|
||||||
return signature
|
return signature
|
||||||
|
|
||||||
# This function returns a list of all ancestors of the given node. The order of the list is
|
# This function returns a list of all ancestors of the given node. The order of the list is
|
||||||
|
|||||||
103
tests/execution/test_caching.py
Normal file
103
tests/execution/test_caching.py
Normal 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)
|
||||||
Loading…
Reference in New Issue
Block a user