diff --git a/comfy_execution/asset_enrichment.py b/comfy_execution/asset_enrichment.py index 38e9496a8..54367d428 100644 --- a/comfy_execution/asset_enrichment.py +++ b/comfy_execution/asset_enrichment.py @@ -1,6 +1,7 @@ """Enrich executed-node output entries with asset id.""" import logging -import os + +from comfy_execution.media_enrichment import resolve_output_entry_path def enrich_output_with_assets(output_ui: dict) -> dict: @@ -15,7 +16,6 @@ def enrich_output_with_assets(output_ui: dict) -> dict: if not args.enable_assets: return output_ui - import folder_paths from app.assets.services.ingest import register_file_in_place, DependencyMissingError enriched = {} @@ -29,20 +29,8 @@ def enrich_output_with_assets(output_ui: dict) -> dict: new_entries.append(entry) continue try: - base = folder_paths.get_directory_by_type(entry["type"]) - if base is None: - new_entries.append(entry) - continue - base_abs = os.path.abspath(base) - abs_path = os.path.abspath(os.path.join(base_abs, entry.get("subfolder") or "", entry["filename"])) - try: - if os.path.commonpath([base_abs, abs_path]) != base_abs: - raise ValueError("escapes base") - except ValueError: - logging.warning("Asset enrichment skipped (path escapes base): %s", entry.get("filename")) - new_entries.append(entry) - continue - if not os.path.isfile(abs_path): + abs_path = resolve_output_entry_path(entry) + if abs_path is None: new_entries.append(entry) continue diff --git a/comfy_execution/media_enrichment.py b/comfy_execution/media_enrichment.py new file mode 100644 index 000000000..c83beed62 --- /dev/null +++ b/comfy_execution/media_enrichment.py @@ -0,0 +1,94 @@ +"""Enrich executed-node output entries with media metadata. + +Attaches a ``metadata`` object (``kind``/``width``/``height``, plus +``duration``/``fps``/``frame_count`` for videos) to each file-type output +entry at output-processing time, so consumers of the ``executed`` message +and ``/history`` can read media properties without probing the files +themselves. Unlike asset enrichment this is NOT gated on ``--enable-assets``: +the properties serve history/websocket consumers that don't run the assets +system at all. +""" +import logging +import mimetypes +import os + + +def resolve_output_entry_path(entry) -> str | None: + """Resolve a file-type output entry to an absolute path inside its base dir. + + Returns ``None`` for non-file entries (no ``filename``/``type``, or + non-string fields), unknown folder types, paths that escape the + folder-type base directory (symlinks are resolved before the containment + check), and paths with no file on disk. + """ + import folder_paths + + if not isinstance(entry, dict) or "filename" not in entry or "type" not in entry: + return None + filename = entry["filename"] + subfolder = entry.get("subfolder") or "" + if not isinstance(filename, str) or not isinstance(subfolder, str) or not isinstance(entry["type"], str): + return None + base = folder_paths.get_directory_by_type(entry["type"]) + if base is None: + return None + # realpath (not abspath) so a symlink planted inside the base directory + # can't smuggle an outside target past the containment check; the base is + # resolved too since output/temp dirs are themselves commonly symlinks. + base_real = os.path.realpath(base) + real_path = os.path.realpath(os.path.join(base_real, subfolder, filename)) + try: + if os.path.commonpath([base_real, real_path]) != base_real: + logging.warning("Output entry path escapes base directory: %s", filename) + return None + except ValueError: + logging.warning("Output entry path escapes base directory: %s", filename) + return None + if not os.path.isfile(real_path): + return None + return real_path + + +def enrich_output_with_media_metadata(output_ui: dict) -> dict: + """Attach a ``metadata`` media-properties object to file output entries. + + Runs once per produced output at output-processing time. Returns a new + dict; entries are copied before modification, and an entry that already + carries a ``metadata`` key is left untouched. Best-effort: unreadable + files, non-media types, and per-entry errors leave that entry unchanged + and never block execution. + """ + try: + from app.assets.services.media_metadata import extract_media_metadata + except Exception: + # Broad on purpose: this enrichment is best-effort and is called + # unguarded from the output-processing path, so a dependency failing + # to import for any reason must degrade to a no-op, not block outputs. + logging.warning("Media metadata extraction unavailable; skipping enrichment", exc_info=True) + return output_ui + + enriched = {} + for key, entries in output_ui.items(): + if not isinstance(entries, list): + enriched[key] = entries + continue + new_entries = [] + for entry in entries: + try: + if isinstance(entry, dict) and "metadata" not in entry: + abs_path = resolve_output_entry_path(entry) + if abs_path is not None: + mime_type = mimetypes.guess_type(entry["filename"], strict=False)[0] + media = extract_media_metadata(abs_path, mime_type=mime_type) + if media: + entry = dict(entry) + entry["metadata"] = media + except Exception: + filename = entry.get("filename") if isinstance(entry, dict) else None + logging.warning( + "Failed to enrich output entry with media metadata: %s", + filename, exc_info=True, + ) + new_entries.append(entry) + enriched[key] = new_entries + return enriched diff --git a/execution.py b/execution.py index 7cab4b331..707aca5da 100644 --- a/execution.py +++ b/execution.py @@ -44,6 +44,7 @@ from comfy_execution.validation import validate_node_input from comfy_execution.progress import get_progress_state, reset_progress_state, add_progress_handler, WebUIProgressHandler from comfy_execution.utils import CurrentNodeContext from comfy_execution.asset_enrichment import enrich_output_with_assets +from comfy_execution.media_enrichment import enrich_output_with_media_metadata from comfy_api.internal import _ComfyNodeInternal, _NodeOutputInternal, first_real_override, is_class, make_locked_method_func from comfy_api.latest import io, _io from comfy_execution.cache_provider import _has_cache_providers, _get_cache_providers, _logger as _cache_logger @@ -561,9 +562,11 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed, asyncio.create_task(await_completion()) return (ExecutionResult.PENDING, None, None) if len(output_ui) > 0: - # Enrich at output-processing time (not in the send path) so assets - # are registered even when no client is connected, and the asset id - # flows into ui_outputs and the cache alongside the raw entries. + # Enrich at output-processing time (not in the send path) so the + # added fields flow into ui_outputs and the cache alongside the + # raw entries, even when no client is connected. Media metadata is + # attached unconditionally; asset ids only under --enable-assets. + output_ui = enrich_output_with_media_metadata(output_ui) output_ui = enrich_output_with_assets(output_ui) ui_outputs[unique_id] = { "meta": { diff --git a/tests-unit/execution_test/test_media_enrichment.py b/tests-unit/execution_test/test_media_enrichment.py new file mode 100644 index 000000000..8116d29bb --- /dev/null +++ b/tests-unit/execution_test/test_media_enrichment.py @@ -0,0 +1,220 @@ +"""Tests for enrich_output_with_media_metadata in comfy_execution/media_enrichment.py.""" +import mimetypes +import os +import unittest +from unittest.mock import MagicMock, patch + +# Initialize the mimetypes registry before any test patches os.path.isfile — +# its lazy init consults os.path.isfile to pick candidate files, and a +# patched-True isfile makes it try to open files that don't exist. +mimetypes.init() + +# Platform-appropriate absolute base. tempfile.gettempdir() returns C:\... on +# Windows and /tmp on POSIX, so containment via commonpath behaves naturally. +_DEFAULT_BASE = os.path.join(__import__("tempfile").gettempdir(), "media-enrichment-test-base") + +_VIDEO_META = { + "kind": "video", + "width": 1280, + "height": 720, + "duration": 4.0, + "fps": 24.0, + "frame_count": 96, +} + + +def _mocked_modules(*, extract=None, directory=_DEFAULT_BASE): + return { + "folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=directory)), + "app.assets.services.media_metadata": MagicMock( + extract_media_metadata=extract or MagicMock(return_value=dict(_VIDEO_META)), + ), + } + + +def _call(output_ui, *, extract=None, file_exists=True, directory=_DEFAULT_BASE): + extract_mock = extract or MagicMock(return_value=dict(_VIDEO_META)) + mocked = _mocked_modules(extract=extract_mock, directory=directory) + + # Only os.path.isfile is patched — abspath/join must run natively so the + # containment check sees real platform paths. + with patch.dict("sys.modules", mocked), \ + patch("os.path.isfile", return_value=file_exists): + import comfy_execution.media_enrichment as mod + return mod.enrich_output_with_media_metadata(output_ui), extract_mock + + +class TestEnrichOutputWithMediaMetadata(unittest.TestCase): + + def test_attaches_metadata_object(self): + output = {"images": [{"filename": "clip.mp4", "subfolder": "", "type": "output"}]} + result, _ = _call(output) + self.assertEqual(result["images"][0]["metadata"], _VIDEO_META) + + def test_passes_guessed_mime_type(self): + output = {"images": [{"filename": "clip.mp4", "subfolder": "", "type": "output"}]} + _, extract_mock = _call(output) + _, kwargs = extract_mock.call_args + self.assertEqual(kwargs["mime_type"], "video/mp4") + + def test_original_entry_not_mutated(self): + orig = {"filename": "clip.mp4", "subfolder": "", "type": "output"} + _call({"images": [orig]}) + self.assertNotIn("metadata", orig) + + def test_non_media_extractor_none_leaves_entry_unchanged(self): + extract = MagicMock(return_value=None) + output = {"latent": [{"filename": "a.latent", "subfolder": "", "type": "output"}]} + result, _ = _call(output, extract=extract) + self.assertNotIn("metadata", result["latent"][0]) + + def test_existing_metadata_key_untouched(self): + entry = {"filename": "clip.mp4", "subfolder": "", "type": "output", "metadata": {"kind": "other"}} + result, extract_mock = _call({"images": [entry]}) + self.assertEqual(result["images"][0]["metadata"], {"kind": "other"}) + extract_mock.assert_not_called() + + def test_non_list_value_passed_through(self): + result, _ = _call({"text": "hello"}) + self.assertEqual(result["text"], "hello") + + def test_none_entry_in_list_unchanged(self): + output = {"images": [None, {"filename": "a.mp4", "subfolder": "", "type": "output"}]} + result, _ = _call(output) + self.assertIsNone(result["images"][0]) + self.assertIn("metadata", result["images"][1]) + + def test_entry_without_filename_unchanged(self): + output = {"latent": [{"subfolder": "", "type": "output"}]} + result, extract_mock = _call(output) + self.assertNotIn("metadata", result["latent"][0]) + extract_mock.assert_not_called() + + def test_file_not_on_disk_unchanged(self): + output = {"images": [{"filename": "missing.mp4", "subfolder": "", "type": "output"}]} + result, extract_mock = _call(output, file_exists=False) + self.assertNotIn("metadata", result["images"][0]) + extract_mock.assert_not_called() + + def test_unknown_type_directory_unchanged(self): + output = {"images": [{"filename": "a.mp4", "subfolder": "", "type": "unknown"}]} + result, extract_mock = _call(output, directory=None) + self.assertNotIn("metadata", result["images"][0]) + extract_mock.assert_not_called() + + def test_path_traversal_subfolder_skipped(self): + output = {"images": [{"filename": "passwd", "subfolder": "../../etc", "type": "output"}]} + result, extract_mock = _call(output) + self.assertNotIn("metadata", result["images"][0]) + extract_mock.assert_not_called() + + def test_absolute_filename_skipped(self): + absolute_filename = os.path.abspath(os.sep + "etc" + os.sep + "passwd") + output = {"images": [{"filename": absolute_filename, "subfolder": "", "type": "output"}]} + result, extract_mock = _call(output) + self.assertNotIn("metadata", result["images"][0]) + extract_mock.assert_not_called() + + def test_extractor_unavailable_returns_unchanged(self): + output = {"images": [{"filename": "a.mp4", "subfolder": "", "type": "output"}]} + mocked = { + "folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=_DEFAULT_BASE)), + # A None sys.modules entry makes the lazy import raise ImportError. + "app.assets.services.media_metadata": None, + } + with patch.dict("sys.modules", mocked), \ + patch("os.path.isfile", return_value=True): + import comfy_execution.media_enrichment as mod + result = mod.enrich_output_with_media_metadata(output) + self.assertNotIn("metadata", result["images"][0]) + + def test_extractor_import_failure_beyond_importerror_degrades(self): + class ExplodingModule: + def __getattr__(self, name): + raise RuntimeError("dependency init failed") + + output = {"images": [{"filename": "a.mp4", "subfolder": "", "type": "output"}]} + mocked = { + "folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=_DEFAULT_BASE)), + "app.assets.services.media_metadata": ExplodingModule(), + } + with patch.dict("sys.modules", mocked), \ + patch("os.path.isfile", return_value=True): + import comfy_execution.media_enrichment as mod + result = mod.enrich_output_with_media_metadata(output) + self.assertNotIn("metadata", result["images"][0]) + + def test_non_string_fields_skipped(self): + output = { + "images": [ + {"filename": 42, "subfolder": "", "type": "output"}, + {"filename": "a.mp4", "subfolder": ["nested"], "type": "output"}, + {"filename": "b.mp4", "subfolder": "", "type": 7}, + ] + } + result, extract_mock = _call(output) + for entry in result["images"]: + self.assertNotIn("metadata", entry) + extract_mock.assert_not_called() + + def test_symlink_escape_rejected_real_fs(self): + import shutil + import tempfile + + root = tempfile.mkdtemp(prefix="media-enrichment-symlink-") + try: + base = os.path.join(root, "output") + os.makedirs(base) + secret = os.path.join(root, "secret.txt") + with open(secret, "w") as f: + f.write("outside") + os.symlink(secret, os.path.join(base, "clip.mp4")) + inside = os.path.join(base, "real.mp4") + with open(inside, "w") as f: + f.write("inside") + + mocked = {"folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=base))} + # No isfile patching — this test runs against the real filesystem. + with patch.dict("sys.modules", mocked): + import comfy_execution.media_enrichment as mod + escaped = mod.resolve_output_entry_path( + {"filename": "clip.mp4", "subfolder": "", "type": "output"}) + contained = mod.resolve_output_entry_path( + {"filename": "real.mp4", "subfolder": "", "type": "output"}) + self.assertIsNone(escaped, "symlink pointing outside the base must be rejected") + self.assertEqual(contained, os.path.realpath(inside)) + finally: + shutil.rmtree(root, ignore_errors=True) + + def test_extractor_error_does_not_block_sibling_entries(self): + call_count = [0] + + def extract_side_effect(abs_path, mime_type=None): + call_count[0] += 1 + if call_count[0] == 1: + raise RuntimeError("boom") + return dict(_VIDEO_META) + + extract = MagicMock(side_effect=extract_side_effect) + output = { + "images": [ + {"filename": "bad.mp4", "subfolder": "", "type": "output"}, + {"filename": "good.mp4", "subfolder": "", "type": "output"}, + ] + } + result, _ = _call(output, extract=extract) + self.assertNotIn("metadata", result["images"][0]) + self.assertEqual(result["images"][1]["metadata"], _VIDEO_META) + + def test_multiple_output_keys_all_enriched(self): + output = { + "images": [{"filename": "a.png", "subfolder": "", "type": "output"}], + "videos": [{"filename": "b.mp4", "subfolder": "", "type": "output"}], + } + result, _ = _call(output) + self.assertIn("metadata", result["images"][0]) + self.assertIn("metadata", result["videos"][0]) + + +if __name__ == "__main__": + unittest.main()