diff --git a/comfy_cli/auth/store.py b/comfy_cli/auth/store.py index d415c426..1ffab6c4 100644 --- a/comfy_cli/auth/store.py +++ b/comfy_cli/auth/store.py @@ -112,6 +112,9 @@ def _read_all(path: Path) -> dict[str, Any]: def _write_all(path: Path, payload: dict[str, Any]) -> None: """Atomic write with 0600 mode from inception. + Deliberately not ``file_utils.atomic_write_text`` — see the write-policy note + in ``comfy_cli/file_utils.py`` (tier 1, secrets). + The tmp file is opened with ``O_CREAT|O_EXCL`` and explicit mode 0o600 so the secrets are never world-readable, even briefly, even on systems with a permissive umask. A unique tmp name (`...tmp`) avoids diff --git a/comfy_cli/command/outdated.py b/comfy_cli/command/outdated.py index 1dd475f6..8a44b752 100644 --- a/comfy_cli/command/outdated.py +++ b/comfy_cli/command/outdated.py @@ -36,6 +36,7 @@ from comfy_cli.command.pack_scan import iter_pack_dirs as _iter_pack_dirs from comfy_cli.command.pack_scan import read_pyproject as _read_pyproject +from comfy_cli.file_utils import atomic_write_text from comfy_cli.registry import RegistryAPI CACHE_TTL_SECONDS = 3600 # 1 hour @@ -77,24 +78,13 @@ def _load_cache() -> dict[str, Any]: def _save_cache(cache: dict[str, Any]) -> None: path = _cache_path() try: - path.parent.mkdir(parents=True, exist_ok=True) - # Write-then-rename: an interrupt mid-write must not leave a truncated - # file that the next `_load_cache` silently resets to `{}`. `os.replace` - # is atomic within a filesystem, and the temp file is a sibling so the - # rename never crosses one. Unique per process — two concurrent writers - # must not share (and truncate) one temp path. - tmp = path.with_name(f"{path.name}.{os.getpid()}.tmp") - try: - # `encoding` pinned rather than left to the platform locale: - # `_load_cache` decodes bytes as JSON, which only accepts - # UTF-8/16/32. Today `json.dumps` defaults to `ensure_ascii=True`, - # so the payload is pure ASCII and any locale would round-trip — - # this keeps that a property of the writer, not a lucky default, if - # a non-ASCII pack id ever reaches the file verbatim. - tmp.write_text(json.dumps(cache), encoding="utf-8") - os.replace(tmp, path) - finally: - tmp.unlink(missing_ok=True) + # Write-then-rename (tier 4, regenerable cache — see the write policy in + # comfy_cli/file_utils.py): an interrupt mid-write must not leave a + # truncated file that the next `_load_cache` silently resets to `{}`. + # The helper creates the parent dir, writes UTF-8 (`_load_cache` decodes + # bytes as JSON, which only accepts UTF-8/16/32), and cleans up its own + # temp file on failure. + atomic_write_text(path, json.dumps(cache), fsync=False) except OSError: # A read-only cache dir must never break a read-only report. pass diff --git a/comfy_cli/command/project.py b/comfy_cli/command/project.py index cb6f19dd..0b42a3e6 100644 --- a/comfy_cli/command/project.py +++ b/comfy_cli/command/project.py @@ -15,7 +15,6 @@ import hashlib import json -import os import urllib.error from datetime import datetime, timezone from pathlib import Path @@ -25,6 +24,7 @@ from comfy_cli import tracking from comfy_cli.command.transfer import _upload_file +from comfy_cli.file_utils import atomic_write_text from comfy_cli.output import get_renderer, rprint from comfy_cli.project import ( ASSETS_LOCK_SCHEMA, @@ -282,18 +282,11 @@ def assets_push_cmd( def _write_assets_lock(path: Path, assets: dict) -> None: - """Atomically rewrite the lock (tmp + fsync + rename, the jobs_state - pattern) so a crash mid-push can't leave a torn JSON file.""" - path.parent.mkdir(parents=True, exist_ok=True) + """Atomically rewrite the lock so a crash mid-push can't leave a torn JSON file.""" doc = {"schema": ASSETS_LOCK_SCHEMA, "assets": assets} - tmp = path.with_suffix(f".{os.getpid()}.tmp") - tmp.write_text(json.dumps(doc, indent=2, sort_keys=True), encoding="utf-8") - fd = os.open(str(tmp), os.O_RDONLY) - try: - os.fsync(fd) - finally: - os.close(fd) - os.replace(tmp, path) + # Tier 3 (cross-process state) per the write policy in comfy_cli/file_utils.py: + # fsync=True, because a lock lost to power failure means re-pushing assets. + atomic_write_text(path, json.dumps(doc, indent=2, sort_keys=True), fsync=True) # --------------------------------------------------------------------------- diff --git a/comfy_cli/command/templates.py b/comfy_cli/command/templates.py index 0f692fff..57c07b20 100644 --- a/comfy_cli/command/templates.py +++ b/comfy_cli/command/templates.py @@ -19,7 +19,6 @@ import json import os -import tempfile import time import urllib.error import urllib.parse @@ -30,6 +29,7 @@ import typer from comfy_cli import tracking +from comfy_cli.file_utils import atomic_write_bytes from comfy_cli.http import plain_urlopen from comfy_cli.output import get_renderer, rprint @@ -138,18 +138,14 @@ def _persist_cache(cache: Path, data: bytes) -> None: propagated. """ try: - cache.parent.mkdir(parents=True, exist_ok=True) - fd, tmp = tempfile.mkstemp(dir=str(cache.parent), prefix=".index-", suffix=".tmp") - try: - with os.fdopen(fd, "wb") as f: - f.write(data) - os.replace(tmp, cache) - except OSError: - try: - os.unlink(tmp) - except OSError: - pass - raise + # Tier 4 (regenerable cache) per the write policy in + # comfy_cli/file_utils.py — fsync=False, wrapped best-effort below. + # The previous inline `tempfile.mkstemp` left the cache at mode 0600 + # (mkstemp hardcodes that and `os.replace` carries it onto the + # destination); the helper restores the umask-derived mode instead, and + # that relaxation is intended: the payload is the public + # template-gallery index, not user data. + atomic_write_bytes(cache, data, fsync=False) except OSError: # Couldn't persist (read-only dir, disk full, …). We still have valid # data in hand, so proceed without caching rather than failing the run. diff --git a/comfy_cli/command/workflow.py b/comfy_cli/command/workflow.py index a0b41422..ca143f0c 100644 --- a/comfy_cli/command/workflow.py +++ b/comfy_cli/command/workflow.py @@ -28,6 +28,7 @@ import typer from comfy_cli import tracking +from comfy_cli.file_utils import atomic_write_text # Aliased at module scope rather than lazy-imported: a class used in ``except`` # clauses at module scope cannot be resolved lazily. ``comfy_cli.http`` is @@ -126,22 +127,6 @@ def _get_graph(input_path: str | None, host: str | None, port: int | None, on_st raise typer.Exit(code=1) from e -def _atomic_write_text(path: Path, content: str) -> None: - """Write via tmp + rename so SIGINT mid-write can't leave a half-written file.""" - import os - - tmp = path.with_suffix(path.suffix + f".{os.getpid()}.tmp") - try: - tmp.write_text(content, encoding="utf-8") - os.replace(tmp, path) - except Exception: - try: - tmp.unlink() - except OSError: - pass - raise - - def _parse_value(raw: str) -> Any: """Parse a CLI-supplied value as JSON; fall back to the literal string.""" try: @@ -306,7 +291,7 @@ def set_slot_cmd( return if not stdout: - _atomic_write_text(p, json.dumps(new_workflow, indent=2)) + atomic_write_text(p, json.dumps(new_workflow, indent=2)) payload: dict[str, Any] = { "workflow": str(p), @@ -421,7 +406,7 @@ def vary_cmd( out.mkdir(parents=True, exist_ok=True) for i, wf in enumerate(workflows): target = out / f"{p.stem}_{i:03d}.json" - _atomic_write_text(target, json.dumps(wf, indent=2)) + atomic_write_text(target, json.dumps(wf, indent=2)) written.append(str(target)) elif renderer.is_pretty(): import sys diff --git a/comfy_cli/cql/loader.py b/comfy_cli/cql/loader.py index 7ed7bac0..d724d0b3 100644 --- a/comfy_cli/cql/loader.py +++ b/comfy_cli/cql/loader.py @@ -30,6 +30,7 @@ from comfy_cli.cql._net import is_loopback_host from comfy_cli.cql.errors import CQLRuntimeError +from comfy_cli.file_utils import atomic_write_text from comfy_cli.http import NoRedirectHandler, build_http_only_opener # Cap raw bytes read from disk or the network. Real `object_info` dumps are a @@ -310,16 +311,11 @@ def write_object_info_cache(host_key: str, data: dict[str, Any]) -> None: """ path = object_info_cache_path(host_key) try: - path.parent.mkdir(parents=True, exist_ok=True) - tmp = path.with_suffix(path.suffix + f".{os.getpid()}.tmp") - tmp.write_text(json.dumps(data), encoding="utf-8") - os.replace(tmp, path) + atomic_write_text(path, json.dumps(data)) except OSError: # A cache we can't write is not worth failing the command over. - try: - tmp.unlink() # type: ignore[possibly-undefined] - except (OSError, NameError, UnboundLocalError): - pass + # atomic_write_text already cleaned up its own tmp file on failure. + pass def read_object_info_cache(host_key: str) -> dict[str, Any] | None: diff --git a/comfy_cli/download_state.py b/comfy_cli/download_state.py index ed7e821f..24ecf4cb 100644 --- a/comfy_cli/download_state.py +++ b/comfy_cli/download_state.py @@ -215,6 +215,9 @@ def write(workspace: Path, state: DownloadState) -> Path: def write_path(path: Path, state: DownloadState) -> Path: """Atomically persist ``state`` at ``path`` (write a tmp file, ``os.replace``). + Deliberately not ``file_utils.atomic_write_text`` — see the write-policy note + in ``comfy_cli/file_utils.py`` (tier 2, secret-adjacent state). + The worker addresses its state file by path rather than by workspace: it is handed ``--state `` and must not have to re-resolve a workspace that the foreground already resolved. diff --git a/comfy_cli/file_utils.py b/comfy_cli/file_utils.py index efd70f83..eb83d497 100644 --- a/comfy_cli/file_utils.py +++ b/comfy_cli/file_utils.py @@ -1,7 +1,9 @@ import json import os import pathlib +import stat import subprocess +import tempfile import time import zipfile from collections.abc import Callable @@ -14,6 +16,150 @@ from comfy_cli import constants, ui from comfy_cli.output.sanitize import sanitize_value +# --------------------------------------------------------------------------- +# Atomic writes — the write policy +# --------------------------------------------------------------------------- +# +# Every hand-rolled tmp+``os.replace`` writer in comfy-cli falls into one of +# four tiers. The tier decides whether the writer uses the shared helpers below +# or stays bespoke; it is decided by *what the file can contain*, never by +# convenience. The helpers deliberately take **no** ``mode=`` parameter — +# parameterizing permissions was considered and rejected, because a writer whose +# payload needs 0600 needs more than a mode (an ``O_EXCL`` open at that mode, a +# 0700 parent, an fsync) and is clearer written out in full at its call site. +# +# Tier 1 — secrets. ``comfy_cli/auth/store.py`` ``_write_all`` stays bespoke: +# the tmp file is opened ``O_CREAT|O_EXCL`` with mode 0600 *at open*, so the +# OAuth tokens are never briefly world-readable under a permissive umask; the +# parent directory is forced to 0700; the write is fsynced. +# +# Tier 2 — secret-adjacent state. ``comfy_cli/download_state.py`` ``write_path`` +# stays bespoke: chmod 0600 + 0700 parent + fsync, because the resolved +# download URL persisted in the state file can embed a presigned/SAS token, +# which is bearer credential material even though the file itself is nominally +# bookkeeping. +# +# Tier 3 — cross-process state. ``atomic_write_text(..., fsync=True)``: +# umask mode, durable. A second process (a watcher, a later CLI invocation, an +# agent shelling out) reads these, and losing one to power failure loses real +# work. Members: ``jobs_state.write``, ``command/project.py`` +# ``_write_assets_lock``. +# +# Tier 4 — regenerable caches and manifests. ``fsync=False``: the content can be +# recomputed or refetched, so paying a sync per write buys nothing. Callers may +# additionally wrap the call best-effort (``except OSError: pass``) when a +# read-only or full cache directory must not break the command. Members: +# ``cql/loader``, ``command/outdated.py`` ``_save_cache``, +# ``command/templates.py`` ``_persist_cache``, the skills manifest writers, +# ``command/workflow.py``. +# +# Two invariants recorded here because they are the *reason* for a tier +# assignment, and a future change to either moves a writer between tiers: +# +# (a) ``JobState`` files are tier 3 (umask mode), not tier 2, because their +# ``outputs`` are plain ``/view?filename=...`` URLs — they carry no embedded +# credential, and the CLI never asks the cloud jobs API for ``?short_link=`` +# responses. If either changes (outputs start carrying signed URLs, or the +# CLI starts requesting short links), jobs_state moves to tier 2 and needs +# 0600 like download_state. +# +# (b) No writer here fsyncs the parent directory *before* ``os.replace`` returns +# durably — ``fsync=True`` syncs the file's contents and then the parent +# directory, but on a power failure between the two an already-``replace``d +# rename can still be lost. This is accepted: the failure mode is losing the +# *newest* write, never observing a torn or half-written file. Contents are +# never torn; renames may be lost. + + +def atomic_write_text(path: pathlib.Path, content: str, *, fsync: bool = False) -> None: + """Atomically write ``content`` to ``path`` via a sibling tmp file + ``os.replace``. + + The write goes to a uniquely-named tmp file in the same directory (so the rename + stays on one filesystem and is atomic), which is then renamed over ``path``. + A SIGINT or crash mid-write therefore never leaves a half-written or empty + file at the destination — readers see either the old contents or the new. + On any failure the tmp file is cleaned up and the exception re-raised. + + The tmp file is created with ``tempfile.mkstemp`` (``O_CREAT | O_EXCL``, no + symlink following) in the destination directory, so concurrent writers never + collide on a shared name and a pre-planted symlink can't redirect the write + (CWE-377). + + Args: + path: destination file. Parent directories are created if missing. + content: text to write (UTF-8). + fsync: if True, flush the tmp file's contents to disk before the rename + and fsync the destination directory afterwards so both the data and + the rename survive power loss, at the cost of a sync. Best-effort: + a failing fsync is ignored, matching the prior per-site behavior. + """ + _atomic_write(path, content.encode("utf-8"), fsync=fsync) + + +def atomic_write_bytes(path: pathlib.Path, data: bytes, *, fsync: bool = False) -> None: + """Atomically write ``data`` to ``path`` — the bytes twin of :func:`atomic_write_text`. + + Identical semantics (same tmp-file creation, same permission handling, same + cleanup-on-failure); it just skips the UTF-8 encode for a caller that already + holds bytes. See :func:`atomic_write_text` for the full contract. + """ + _atomic_write(path, data, fsync=fsync) + + +def _atomic_write(path: pathlib.Path, data: bytes, *, fsync: bool) -> None: + """Shared implementation behind :func:`atomic_write_text` / :func:`atomic_write_bytes`.""" + path.parent.mkdir(parents=True, exist_ok=True) + # mkstemp gives a unique, O_EXCL, non-symlink-following fd opened O_RDWR in the + # destination directory — same filesystem, so the os.replace below is atomic. + fd, tmp_name = tempfile.mkstemp(dir=str(path.parent), prefix=path.name + ".", suffix=".tmp") + try: + with os.fdopen(fd, "wb") as f: + f.write(data) + if fsync: + f.flush() + try: + os.fsync(f.fileno()) # O_RDWR fd, so fsync works on Windows too + except OSError: + pass + # mkstemp hardcodes the tmp file to 0600, and os.replace carries that mode + # onto the destination — so without this a first atomic write would quietly + # strip the group/other read access a shared output is meant to have. Restore + # the intended mode before the rename: reuse the existing destination's bits, + # else fall back to the umask-derived default (0666 & ~umask) for a new file. + try: + dest_mode = stat.S_IMODE(os.stat(path).st_mode) + except OSError: + umask = os.umask(0) + os.umask(umask) + dest_mode = 0o666 & ~umask + try: + os.chmod(tmp_name, dest_mode) + except OSError: + # Windows / filesystems without POSIX perms: best-effort, matching fsync. + pass + os.replace(tmp_name, path) + if fsync: + # Also fsync the parent directory so the rename itself is durable. + try: + dir_fd = os.open(str(path.parent), os.O_RDONLY) + try: + os.fsync(dir_fd) + except OSError: + pass + finally: + os.close(dir_fd) + except OSError: + # e.g. Windows can't open a directory for fsync; best-effort. + pass + except BaseException: + # BaseException (not just Exception) so a KeyboardInterrupt mid-write + # still cleans up the tmp file, per the docstring. + try: + os.unlink(tmp_name) + except OSError: + pass + raise + class DownloadException(Exception): pass diff --git a/comfy_cli/jobs_state.py b/comfy_cli/jobs_state.py index 2758125d..1b83549e 100644 --- a/comfy_cli/jobs_state.py +++ b/comfy_cli/jobs_state.py @@ -39,15 +39,14 @@ from __future__ import annotations import json -import os import re -import secrets as _secrets from dataclasses import asdict, dataclass, field from datetime import datetime, timezone from pathlib import Path from typing import Any from comfy_cli import constants, locking +from comfy_cli.file_utils import atomic_write_text from comfy_cli.utils import get_os TERMINAL_STATUSES = frozenset({"completed", "error", "cancelled"}) @@ -123,18 +122,8 @@ def write(state: JobState) -> Path | None: # Lock per-file so a watcher and a foreground update can't tear each # other's writes. with locking.file_lock(path.with_suffix(".lock")): - tmp = path.with_suffix(f".{os.getpid()}.{_secrets.token_hex(4)}.tmp") - tmp.write_text(json.dumps(state.to_dict(), indent=2, default=str), encoding="utf-8") - # fsync for durability before atomic rename - try: - fd = os.open(str(tmp), os.O_RDONLY) - try: - os.fsync(fd) - finally: - os.close(fd) - except OSError: - pass - os.replace(tmp, path) + # fsync=True: durability against power loss before the atomic rename. + atomic_write_text(path, json.dumps(state.to_dict(), indent=2, default=str), fsync=True) return path diff --git a/comfy_cli/skills/__init__.py b/comfy_cli/skills/__init__.py index 78f22c42..9297e72c 100644 --- a/comfy_cli/skills/__init__.py +++ b/comfy_cli/skills/__init__.py @@ -36,6 +36,8 @@ from pathlib import Path from typing import Literal +from comfy_cli.file_utils import atomic_write_text + # Where the bundled skills live. Each tuple is (skill_name, package_subdir). # ``skill_name`` is the public identifier used in AGENTS.md fences and as the # subdir name in installed targets. ``package_subdir`` is the local resource @@ -199,11 +201,7 @@ def read_manifest() -> dict: def write_manifest(manifest: dict) -> None: """Atomically write the manifest (tmp + rename so a SIGINT can't corrupt it).""" - path = manifest_path() - path.parent.mkdir(parents=True, exist_ok=True) - tmp = path.with_suffix(f".{os.getpid()}.tmp") - tmp.write_text(json.dumps(manifest, indent=2), encoding="utf-8") - os.replace(tmp, path) + atomic_write_text(manifest_path(), json.dumps(manifest, indent=2)) def _sha256(text: str) -> str: @@ -577,24 +575,9 @@ def _backup_if_user_edited(path: Path, expected_content: str) -> Path | None: return bak -def _atomic_write_text(path: Path, content: str) -> None: - """Write via tmp + rename so a SIGINT mid-write can't leave the file empty.""" - path.parent.mkdir(parents=True, exist_ok=True) - tmp = path.with_suffix(path.suffix + f".{os.getpid()}.tmp") - try: - tmp.write_text(content, encoding="utf-8") - os.replace(tmp, path) - except Exception: - try: - tmp.unlink() - except OSError: - pass - raise - - def _write_claude_skill(path: Path, content: str) -> None: _backup_if_user_edited(path, content) - _atomic_write_text(path, content) + atomic_write_text(path, content) def _cursor_description_for(skill_name: str) -> str: @@ -610,14 +593,14 @@ def _write_cursor_rule(path: Path, content: str, *, skill_name: str) -> None: body = _strip_frontmatter(content) rule = f'---\ndescription: {_cursor_description_for(skill_name)}\nglobs: "**/*"\nalwaysApply: false\n---\n\n{body}' _backup_if_user_edited(path, rule) - _atomic_write_text(path, rule) + atomic_write_text(path, rule) def _upsert_agents_md_block(path: Path, content: str, *, skill_name: str) -> None: start, end = _agents_fence(skill_name) block = f"\n{start}\n{content}\n{end}\n" if not path.exists(): - _atomic_write_text(path, block.lstrip("\n")) + atomic_write_text(path, block.lstrip("\n")) return existing = path.read_text(encoding="utf-8") if start in existing and end in existing: @@ -626,7 +609,7 @@ def _upsert_agents_md_block(path: Path, content: str, *, skill_name: str) -> Non new = before.rstrip() + "\n\n" + block.lstrip("\n") + after.lstrip("\n") else: new = existing.rstrip() + "\n" + block - _atomic_write_text(path, new) + atomic_write_text(path, new) def _remove_agents_md_block(path: Path, *, skill_name: str) -> bool: diff --git a/tests/comfy_cli/command/test_node_deps.py b/tests/comfy_cli/command/test_node_deps.py index e85282af..d390ed60 100644 --- a/tests/comfy_cli/command/test_node_deps.py +++ b/tests/comfy_cli/command/test_node_deps.py @@ -807,27 +807,31 @@ def test_save_cache_pins_the_write_encoding(monkeypatch): """The round-trip above cannot catch a dropped `encoding="utf-8"`: the bytes it writes are pure ASCII, which every locale encodes identically, so the pin only matters the day `json.dumps` stops escaping. Nor can the ambient - encoding be faked — CPython resolves `write_text`'s default below the Python - `locale` module. So assert the writer's contract directly: `_save_cache` - names its encoding rather than inheriting the host's, and what lands on disk - is the UTF-8 that `_load_cache`'s `json.loads(bytes)` can decode. + encoding be faked — CPython resolves the text-mode default below the Python + `locale` module. So assert the writer's contract directly, at the seam where + text becomes bytes: `_save_cache` hands `atomic_write_text` a string, and the + bytes that reach the filesystem are that string encoded UTF-8 by us — never a + text-mode write whose codec the host locale gets to pick. What lands on disk + is therefore the UTF-8 that `_load_cache`'s `json.loads(bytes)` can decode. """ import time + from comfy_cli import file_utils from comfy_cli.command import outdated as outdated_cmd seen: dict[str, object] = {} - real_write_text = Path.write_text + real_atomic_write = file_utils._atomic_write - def spy(self, data, encoding=None, **kwargs): - seen["encoding"] = encoding - return real_write_text(self, data, encoding=encoding, **kwargs) + def spy(path, data, **kwargs): + seen["data"] = data + return real_atomic_write(path, data, **kwargs) - monkeypatch.setattr(Path, "write_text", spy) + monkeypatch.setattr(file_utils, "_atomic_write", spy) key = f"{node_deps_cmd.REGISTRY_CACHE_PREFIX}https://api.comfy.org:café-pack" - outdated_cmd._save_cache({key: {"value": "1.0.0", "ts": time.time()}}) + payload = {key: {"value": "1.0.0", "ts": time.time()}} + outdated_cmd._save_cache(payload) - assert seen["encoding"] == "utf-8", "the host locale must not pick the cache encoding" + assert seen["data"] == json.dumps(payload).encode("utf-8"), "the host locale must not pick the cache encoding" assert json.loads(outdated_cmd._cache_path().read_bytes())[key]["value"] == "1.0.0" diff --git a/tests/comfy_cli/command/test_templates.py b/tests/comfy_cli/command/test_templates.py index 86c2f5a6..953bf105 100644 --- a/tests/comfy_cli/command/test_templates.py +++ b/tests/comfy_cli/command/test_templates.py @@ -15,6 +15,7 @@ import pytest from typer.testing import CliRunner +from comfy_cli import file_utils from comfy_cli.caller import Caller from comfy_cli.command import templates as templates_cmd from comfy_cli.output.renderer import ( @@ -542,7 +543,9 @@ def _boom_mkstemp(*args, **kwargs): monkeypatch.setattr(templates_cmd, "_fetch_gallery", _fake_fetch) # Make the real _persist_cache's write fail (read-only dir / disk full); # it must swallow the error and let the command proceed on in-hand data. - monkeypatch.setattr(templates_cmd.tempfile, "mkstemp", _boom_mkstemp) + # Patched at the shared helper's own tmp-file seam so the real + # `_persist_cache` -> `atomic_write_bytes` path is still exercised. + monkeypatch.setattr(file_utils.tempfile, "mkstemp", _boom_mkstemp) _force_json_renderer() runner = CliRunner() diff --git a/tests/comfy_cli/test_file_utils.py b/tests/comfy_cli/test_file_utils.py index dfeed4ea..12eafb20 100644 --- a/tests/comfy_cli/test_file_utils.py +++ b/tests/comfy_cli/test_file_utils.py @@ -1,6 +1,12 @@ +import os +import stat +import sys import zipfile +import pytest + from comfy_cli import file_utils +from comfy_cli.file_utils import atomic_write_bytes, atomic_write_text def test_zip_files_respects_comfyignore(tmp_path, monkeypatch): @@ -81,3 +87,225 @@ def test_zip_files_without_git_falls_back_to_walk(tmp_path, monkeypatch): assert "file.txt" in names assert "node.zip" not in names + + +def test_atomic_write_text_creates_file_and_parents(tmp_path): + target = tmp_path / "sub" / "dir" / "out.json" + atomic_write_text(target, '{"a": 1}') + + assert target.read_text(encoding="utf-8") == '{"a": 1}' + # No stray tmp files left behind. + assert list(target.parent.glob("*.tmp")) == [] + + +def test_atomic_write_text_overwrites_existing(tmp_path): + target = tmp_path / "out.txt" + target.write_text("old", encoding="utf-8") + + atomic_write_text(target, "new") + + assert target.read_text(encoding="utf-8") == "new" + + +def test_atomic_write_text_fsync_true_still_writes(tmp_path, monkeypatch): + target = tmp_path / "durable.txt" + + # Spy on os.fsync so we assert it's actually invoked, not silently no-op'd + # (the class of bug flagged for the Windows O_RDONLY case). Delegate to the + # real fsync so durability behavior is preserved. + real_fsync = file_utils.os.fsync + synced_fds = [] + + def spy_fsync(fd): + synced_fds.append(fd) + return real_fsync(fd) + + monkeypatch.setattr(file_utils.os, "fsync", spy_fsync) + + atomic_write_text(target, "durable", fsync=True) + + assert target.read_text(encoding="utf-8") == "durable" + assert list(target.parent.glob("*.tmp")) == [] + # The tmp file fd is always fsynced (O_RDWR, works cross-platform). On POSIX + # the parent directory fd is fsynced too; Windows can't open a dir for fsync + # and best-effort skips it, so only require the extra sync off Windows. + if sys.platform == "win32": + assert len(synced_fds) >= 1 + else: + assert len(synced_fds) == 2 + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX file-mode semantics") +def test_atomic_write_text_preserves_existing_mode(tmp_path): + # A first atomic write must not clobber the destination's mode down to mkstemp's + # hardcoded 0600. We only need a non-0600 mode to prove restoration happens, so + # use owner-execute (0o700) — a bit mkstemp's 0600 lacks — rather than a + # group/other-readable mode. That keeps this clear of the py/overly-permissive-file + # scanner while still exercising the exact "restore bits beyond 0600" path. + target = tmp_path / "shared.json" + target.write_text("old", encoding="utf-8") + os.chmod(target, 0o700) + + atomic_write_text(target, "new") + + assert stat.S_IMODE(os.stat(target).st_mode) == 0o700 + assert target.read_text(encoding="utf-8") == "new" + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX file-mode semantics") +def test_atomic_write_text_new_file_uses_umask_default(tmp_path): + # A new destination gets the umask-derived default, not mkstemp's hardcoded 0600. + old_umask = os.umask(0o022) + try: + target = tmp_path / "fresh.json" + atomic_write_text(target, "data") + assert stat.S_IMODE(os.stat(target).st_mode) == (0o666 & ~0o022) + finally: + os.umask(old_umask) + + +def test_atomic_write_text_cleans_up_tmp_on_failure(tmp_path, monkeypatch): + target = tmp_path / "out.txt" + target.write_text("original", encoding="utf-8") + + def boom(src, dst): + raise OSError("replace failed") + + # Fail at the rename step, after the tmp file has been written. + monkeypatch.setattr(file_utils.os, "replace", boom) + + with pytest.raises(OSError): + atomic_write_text(target, "new content") + + # The tmp file is cleaned up and the destination is untouched. + assert list(target.parent.glob("*.tmp")) == [] + assert target.read_text(encoding="utf-8") == "original" + + +def test_atomic_write_text_tmp_name_is_unique_per_write(tmp_path, monkeypatch): + # Two writes from the "same" pid must not collide on a shared tmp path. + target = tmp_path / "out.txt" + seen = [] + real_mkstemp = file_utils.tempfile.mkstemp + + def spy(*args, **kwargs): + fd, name = real_mkstemp(*args, **kwargs) + seen.append(name) + return fd, name + + monkeypatch.setattr(file_utils.tempfile, "mkstemp", spy) + monkeypatch.setattr(file_utils.os, "getpid", lambda: 4242) + + atomic_write_text(target, "a") + atomic_write_text(target, "b") + + assert len(seen) == 2 + assert seen[0] != seen[1] + assert target.read_text(encoding="utf-8") == "b" + + +def test_atomic_write_text_does_not_follow_symlinked_tmp(tmp_path): + # A pre-planted symlink at a predictable tmp path must not redirect the write. + target = tmp_path / "out.txt" + victim = tmp_path / "victim.txt" + victim.write_text("do-not-touch", encoding="utf-8") + # The old scheme used "..tmp"; plant a symlink there. + import os as _os + + decoy = tmp_path / f"out.txt.{_os.getpid()}.tmp" + decoy.symlink_to(victim) + + atomic_write_text(target, "new") + + assert target.read_text(encoding="utf-8") == "new" + assert victim.read_text(encoding="utf-8") == "do-not-touch" + + +# --- atomic_write_bytes ----------------------------------------------------- +# The bytes twin shares one private implementation with atomic_write_text, so +# these mirror the text-variant cases over the seam that differs: the payload is +# handed through verbatim, with no encode and no newline translation. + + +def test_atomic_write_bytes_creates_file_and_parents(tmp_path): + target = tmp_path / "sub" / "dir" / "index.json" + atomic_write_bytes(target, b'{"a": 1}') + + assert target.read_bytes() == b'{"a": 1}' + # No stray tmp files left behind. + assert list(target.parent.glob("*.tmp")) == [] + + +def test_atomic_write_bytes_overwrites_existing(tmp_path): + target = tmp_path / "index.json" + target.write_bytes(b"old") + + atomic_write_bytes(target, b"new") + + assert target.read_bytes() == b"new" + + +def test_atomic_write_bytes_writes_payload_verbatim(tmp_path): + # No UTF-8 encode and no platform newline translation: non-UTF-8 bytes and + # bare LFs survive byte-for-byte, which is the reason this variant exists. + target = tmp_path / "raw.bin" + payload = b"\xff\xfe\r\n\x00line\n" + + atomic_write_bytes(target, payload) + + assert target.read_bytes() == payload + + +def test_atomic_write_bytes_fsync_true_still_writes(tmp_path, monkeypatch): + target = tmp_path / "durable.bin" + + real_fsync = file_utils.os.fsync + synced_fds = [] + + def spy_fsync(fd): + synced_fds.append(fd) + return real_fsync(fd) + + monkeypatch.setattr(file_utils.os, "fsync", spy_fsync) + + atomic_write_bytes(target, b"durable", fsync=True) + + assert target.read_bytes() == b"durable" + assert list(target.parent.glob("*.tmp")) == [] + # Same split as the text variant: the tmp fd always syncs; the parent + # directory fd only off Windows, which can't open a dir for fsync. + if sys.platform == "win32": + assert len(synced_fds) >= 1 + else: + assert len(synced_fds) == 2 + + +def test_atomic_write_bytes_cleans_up_tmp_on_failure(tmp_path, monkeypatch): + target = tmp_path / "index.json" + target.write_bytes(b"original") + + def boom(src, dst): + raise OSError("replace failed") + + # Fail at the rename step, after the tmp file has been written. + monkeypatch.setattr(file_utils.os, "replace", boom) + + with pytest.raises(OSError): + atomic_write_bytes(target, b"new content") + + # The tmp file is cleaned up and the destination is untouched. + assert list(target.parent.glob("*.tmp")) == [] + assert target.read_bytes() == b"original" + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX file-mode semantics") +def test_atomic_write_bytes_new_file_uses_umask_default(tmp_path): + # The templates gallery cache relies on this: the old inline mkstemp left it + # at 0600, and the switch to the umask-derived default is intended. + old_umask = os.umask(0o022) + try: + target = tmp_path / "fresh.bin" + atomic_write_bytes(target, b"data") + assert stat.S_IMODE(os.stat(target).st_mode) == (0o666 & ~0o022) + finally: + os.umask(old_umask)