From bf7257448e237325010b70c36c6e5f543f2f6498 Mon Sep 17 00:00:00 2001 From: magictut <131337141+magicwang1111@users.noreply.github.com> Date: Mon, 27 Apr 2026 10:07:30 +0800 Subject: [PATCH] Address asset isolation review feedback --- app/assets/services/asset_management.py | 4 +++- app/assets/services/ingest.py | 5 ++++- execution.py | 3 ++- main.py | 9 ++------- server.py | 14 ++++++++++++-- 5 files changed, 23 insertions(+), 12 deletions(-) diff --git a/app/assets/services/asset_management.py b/app/assets/services/asset_management.py index af0d6db29..7a4631d2c 100644 --- a/app/assets/services/asset_management.py +++ b/app/assets/services/asset_management.py @@ -331,7 +331,9 @@ def is_file_visible_to_owner( owner_id = (owner_id or "").strip() with create_session() as session: ref = get_reference_by_file_path(session, locator) - if not ref or ref.deleted_at is not None: + if not ref: + return os.path.isfile(locator) + if ref.deleted_at is not None: return False return ref.owner_id == "" or ref.owner_id == owner_id diff --git a/app/assets/services/ingest.py b/app/assets/services/ingest.py index 92634b833..ce70c5a13 100644 --- a/app/assets/services/ingest.py +++ b/app/assets/services/ingest.py @@ -193,7 +193,10 @@ def collect_output_absolute_paths(output_data: dict) -> list[str]: abs_path = os.path.abspath( os.path.join(base_dir, item.get("subfolder", ""), filename) ) - if os.path.commonpath((base_dir, abs_path)) != base_dir: + try: + if os.path.commonpath((base_dir, abs_path)) != base_dir: + continue + except ValueError: continue if abs_path not in seen: seen.add(abs_path) diff --git a/execution.py b/execution.py index 5740c6362..aeda10466 100644 --- a/execution.py +++ b/execution.py @@ -551,10 +551,11 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed, if len(output_ui) > 0: register_output_assets = getattr(server, "register_output_assets", None) if register_output_assets is not None: + user_id_key = getattr(server, "INTERNAL_USER_ID_KEY", "_comfy_user_id") register_output_assets( output_ui, prompt_id, - extra_data.get("_comfy_user_id", ""), + extra_data.get(user_id_key, ""), ) ui_outputs[unique_id] = { "meta": { diff --git a/main.py b/main.py index ee1a67915..ab3d3197a 100644 --- a/main.py +++ b/main.py @@ -241,11 +241,6 @@ def cuda_malloc_warning(): logging.warning("\nWARNING: this card most likely does not support cuda-malloc, if you get \"CUDA error\" please run ComfyUI with: --disable-cuda-malloc\n") -def _collect_output_absolute_paths(history_result: dict) -> list[str]: - """Extract absolute file paths for output items from a history result.""" - return collect_output_absolute_paths(history_result) - - def prompt_worker(q, server_instance): current_time: float = 0.0 cache_ram = args.cache_ram @@ -308,8 +303,8 @@ def prompt_worker(q, server_instance): logging.info("Prompt executed in {:.2f} seconds".format(execution_time)) if not asset_seeder.is_disabled(): - paths = _collect_output_absolute_paths(e.history_result) - owner_id = extra_data.get("_comfy_user_id", "") + paths = collect_output_absolute_paths(e.history_result) + owner_id = extra_data.get(server.INTERNAL_USER_ID_KEY, "") register_output_files(paths, job_id=prompt_id, owner_id=owner_id) flags = q.get_flags() diff --git a/server.py b/server.py index 8efb77710..eaacdce93 100644 --- a/server.py +++ b/server.py @@ -69,6 +69,7 @@ def _remove_sensitive_from_queue(queue: list) -> list: def _scrub_prompt_tuple(prompt_tuple): + """Remove internal-only prompt metadata before returning queue data.""" if not isinstance(prompt_tuple, (list, tuple)): return prompt_tuple if len(prompt_tuple) <= 3 or not isinstance(prompt_tuple[3], dict): @@ -81,6 +82,7 @@ def _scrub_prompt_tuple(prompt_tuple): def _scrub_history_for_response(history: dict) -> dict: + """Remove internal-only prompt metadata from history responses.""" out = {} for prompt_id, item in history.items(): if not isinstance(item, dict): @@ -92,25 +94,32 @@ def _scrub_history_for_response(history: dict) -> dict: return out -def _prompt_tuple_owner_id(prompt_tuple) -> str: +def _prompt_tuple_owner_id(prompt_tuple) -> str | None: + """Return the stored owner id, or None for legacy prompts without one.""" try: extra_data = prompt_tuple[3] except Exception: return "default" if not isinstance(extra_data, dict): return "default" + if INTERNAL_USER_ID_KEY not in extra_data: + return None return str(extra_data.get(INTERNAL_USER_ID_KEY) or "default") def _prompt_tuple_visible_to_user(prompt_tuple, owner_id: str) -> bool: - return _prompt_tuple_owner_id(prompt_tuple) == str(owner_id or "default") + """Return whether a prompt tuple is visible to the requesting user.""" + prompt_owner_id = _prompt_tuple_owner_id(prompt_tuple) + return prompt_owner_id is None or prompt_owner_id == str(owner_id or "default") def _filter_queue_for_user(queue: list, owner_id: str) -> list: + """Filter queue entries to those visible to the requesting user.""" return [item for item in queue if _prompt_tuple_visible_to_user(item, owner_id)] def _filter_history_for_user(history: dict, owner_id: str) -> dict: + """Filter history entries to those visible to the requesting user.""" return { prompt_id: item for prompt_id, item in history.items() @@ -120,6 +129,7 @@ def _filter_history_for_user(history: dict, owner_id: str) -> dict: def _slice_history(history: dict, max_items: int | None, offset: int) -> dict: + """Return a stable paginated slice of a history mapping.""" items = list(history.items()) if offset < 0 and max_items is not None: offset = len(items) - max_items