diff --git a/src/specify_cli/_init_options.py b/src/specify_cli/_init_options.py index 9f509da256..a53730d852 100644 --- a/src/specify_cli/_init_options.py +++ b/src/specify_cli/_init_options.py @@ -1,6 +1,8 @@ """Helpers for interpreting persisted init options.""" import json +import os +import tempfile from collections.abc import Mapping from pathlib import Path from typing import Any, Union @@ -23,10 +25,22 @@ def save_init_options(project_path: Path, options: dict[str, Any]) -> None: """Persist the CLI options used during ``specify init``.""" dest = project_path / INIT_OPTIONS_FILE dest.parent.mkdir(parents=True, exist_ok=True) - dest.write_text( - json.dumps(options, indent=2, sort_keys=True, ensure_ascii=False) + "\n", - encoding="utf-8", + fd, tmp = tempfile.mkstemp( + dir=str(dest.parent), + prefix=f".{dest.name}.", + suffix=".tmp", ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + json.dump(options, f, indent=2, sort_keys=True, ensure_ascii=False) + f.write("\n") + os.replace(tmp, dest) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise def load_init_options(project_path: Path) -> dict[str, Any]: diff --git a/src/specify_cli/events.py b/src/specify_cli/events.py index 17a9d7ffbe..9b7a506400 100644 --- a/src/specify_cli/events.py +++ b/src/specify_cli/events.py @@ -17,6 +17,7 @@ import sys import subprocess import platform +import tempfile from pathlib import Path, PurePosixPath, PureWindowsPath from typing import TYPE_CHECKING, Any @@ -2151,7 +2152,21 @@ def _merge_toml_fragment(dst: Path, fragment: str) -> bool: flags=re.DOTALL, ) dst.parent.mkdir(parents=True, exist_ok=True) - dst.write_text(existing.rstrip() + "\n\n" + fragment + "\n", encoding="utf-8") + fd, tmp = tempfile.mkstemp( + dir=str(dst.parent), prefix=f".{dst.name}.", suffix=".tmp" + ) + try: + if dst.exists() and hasattr(os, "fchmod"): + os.fchmod(fd, dst.stat(follow_symlinks=False).st_mode & 0o7777) + with os.fdopen(fd, "w", encoding="utf-8") as f: + f.write(existing.rstrip() + "\n\n" + fragment + "\n") + os.replace(tmp, dst) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise return True @@ -2230,7 +2245,21 @@ def _remove_toml_entries(dst: Path) -> bool: if not stripped: dst.unlink(missing_ok=True) return True - dst.write_text(cleaned, encoding="utf-8") + fd, tmp = tempfile.mkstemp( + dir=str(dst.parent), prefix=f".{dst.name}.", suffix=".tmp" + ) + try: + if dst.exists() and hasattr(os, "fchmod"): + os.fchmod(fd, dst.stat(follow_symlinks=False).st_mode & 0o7777) + with os.fdopen(fd, "w", encoding="utf-8") as f: + f.write(cleaned) + os.replace(tmp, dst) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise return False @@ -2520,10 +2549,25 @@ def _load_user_json(path: Path) -> dict | None: def _safe_write_json(dst: Path, data: dict) -> None: - """Write *data* as JSON to *dst* after validating the destination (#12).""" + """Write *data* as JSON to *dst* atomically after validating the destination (#12).""" _ensure_safe_destination(dst) dst.parent.mkdir(parents=True, exist_ok=True) - dst.write_text(json.dumps(data, indent=2) + "\n", encoding="utf-8") + fd, tmp = tempfile.mkstemp( + dir=str(dst.parent), prefix=f".{dst.name}.", suffix=".tmp" + ) + try: + if dst.exists() and hasattr(os, "fchmod"): + os.fchmod(fd, dst.stat(follow_symlinks=False).st_mode & 0o7777) + with os.fdopen(fd, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2) + f.write("\n") + os.replace(tmp, dst) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise def _ensure_safe_destination(dst: Path) -> None: diff --git a/src/specify_cli/presets/__init__.py b/src/specify_cli/presets/__init__.py index b22a440661..083f939f67 100644 --- a/src/specify_cli/presets/__init__.py +++ b/src/specify_cli/presets/__init__.py @@ -63,6 +63,25 @@ def _content_sha256(content: bytes) -> str: return hashlib.sha256(content).hexdigest() +def _atomic_write_text(path: Path, content: str) -> None: + """Write *content* to *path* atomically via mkstemp + os.replace.""" + fd, tmp = tempfile.mkstemp( + dir=str(path.parent), prefix=f".{path.name}.", suffix=".tmp" + ) + try: + if path.exists() and hasattr(os, "fchmod"): + os.fchmod(fd, path.stat(follow_symlinks=False).st_mode & 0o7777) + with os.fdopen(fd, "w", encoding="utf-8") as f: + f.write(content) + os.replace(tmp, path) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise + + def _is_comparable_version(value: str) -> bool: """Return whether a recorded version can be evaluated against a specifier. @@ -1178,7 +1197,7 @@ def _register_commands( composed_dir = preset_dir / ".composed" composed_dir.mkdir(parents=True, exist_ok=True) composed_file = composed_dir / f"{cmd['name']}.md" - composed_file.write_text(composed, encoding="utf-8") + _atomic_write_text(composed_file, composed) commands_to_register.append({ **cmd, "file": f".composed/{cmd['name']}.md", @@ -2170,7 +2189,7 @@ def record_written(written: Dict[str, List[str]]) -> None: composed_dir = pack_dir / ".composed" composed_dir.mkdir(parents=True, exist_ok=True) composed_file = composed_dir / f"{cmd_name}.md" - composed_file.write_text(composed, encoding="utf-8") + _atomic_write_text(composed_file, composed) written = self._register_for_non_skill_agents( registrar, [{**tmpl, "file": f".composed/{cmd_name}.md"}], @@ -2190,7 +2209,7 @@ def record_written(written: Dict[str, List[str]]) -> None: shared_composed = self.presets_dir / ".composed" shared_composed.mkdir(parents=True, exist_ok=True) composed_file = shared_composed / f"{cmd_name}.md" - composed_file.write_text(composed, encoding="utf-8") + _atomic_write_text(composed_file, composed) source = layers[0]["source"] if source.startswith("extension:"): source_id = source.split(":", 1)[1].split(" ", 1)[0] diff --git a/tests/test_atomic_writes.py b/tests/test_atomic_writes.py new file mode 100644 index 0000000000..18ce109f08 --- /dev/null +++ b/tests/test_atomic_writes.py @@ -0,0 +1,216 @@ +"""Crash-injection tests for atomic write helpers. + +Verifies that mkstemp + os.replace write patterns across the codebase +leave the original file intact when a mid-write failure occurs, and that +temp files are cleaned up. +""" + +from __future__ import annotations + +import json +import os +import tempfile +from pathlib import Path +from unittest.mock import patch + +import pytest + + +# -- _safe_write_json (events.py) ------------------------------------------- + +class TestSafeWriteJsonAtomic: + """Atomic writes for _safe_write_json in events.py.""" + + def test_successful_write(self, tmp_path: Path) -> None: + from specify_cli.events import _safe_write_json + + dst = tmp_path / "config.json" + data = {"key": "value", "nested": {"a": 1}} + + _safe_write_json(dst, data) + + assert dst.exists() + loaded = json.loads(dst.read_text(encoding="utf-8")) + assert loaded == data + + def test_original_intact_on_write_failure(self, tmp_path: Path) -> None: + """Mid-write failure must leave the original file untouched.""" + from specify_cli.events import _safe_write_json + + dst = tmp_path / "config.json" + original = {"original": True} + dst.write_text(json.dumps(original) + "\n", encoding="utf-8") + + with patch("specify_cli.events.os.replace", side_effect=OSError("disk full")): + with pytest.raises(OSError, match="disk full"): + _safe_write_json(dst, {"corrupted": True}) + + loaded = json.loads(dst.read_text(encoding="utf-8")) + assert loaded == original + + def test_temp_file_cleaned_up_after_failure(self, tmp_path: Path) -> None: + """Temp files must not be left behind after a failure.""" + from specify_cli.events import _safe_write_json + + dst = tmp_path / "config.json" + before = set(tmp_path.iterdir()) + + with patch("specify_cli.events.os.replace", side_effect=OSError("fail")): + with pytest.raises(OSError): + _safe_write_json(dst, {"data": 1}) + + after = set(tmp_path.iterdir()) + new_files = after - before + # Only the original file should exist (if created), no temp files + for f in new_files: + assert not f.name.startswith(".config.json.") + + def test_preserves_file_permissions(self, tmp_path: Path) -> None: + """Atomic write should preserve original file permissions.""" + from specify_cli.events import _safe_write_json + + dst = tmp_path / "config.json" + dst.write_text("{}\n", encoding="utf-8") + original_mode = dst.stat().st_mode & 0o7777 + + _safe_write_json(dst, {"updated": True}) + + new_mode = dst.stat().st_mode & 0o7777 + assert new_mode == original_mode + + +# -- _merge_toml_fragment (events.py) --------------------------------------- + +class TestMergeTomlFragmentAtomic: + """Atomic writes for _merge_toml_fragment in events.py.""" + + def test_successful_write(self, tmp_path: Path) -> None: + from specify_cli.events import _merge_toml_fragment + + dst = tmp_path / "hooks.toml" + fragment = "[[hooks.session_start]]\ncommand = 'echo hello'\nspeckit_marker = true" + + result = _merge_toml_fragment(dst, fragment) + + assert result is True + assert dst.exists() + content = dst.read_text(encoding="utf-8") + assert "echo hello" in content + + def test_original_intact_on_write_failure(self, tmp_path: Path) -> None: + from specify_cli.events import _merge_toml_fragment + + dst = tmp_path / "hooks.toml" + original = "[[hooks.custom]]\ncommand = 'original'\n" + dst.write_text(original, encoding="utf-8") + + with patch("specify_cli.events.os.replace", side_effect=OSError("fail")): + with pytest.raises(OSError, match="fail"): + _merge_toml_fragment(dst, "fragment") + + assert dst.read_text(encoding="utf-8") == original + + +# -- _remove_toml_entries (events.py) ---------------------------------------- + +class TestRemoveTomlEntriesAtomic: + """Atomic writes for _remove_toml_entries in events.py.""" + + def test_successful_removal(self, tmp_path: Path) -> None: + from specify_cli.events import _remove_toml_entries + + dst = tmp_path / "hooks.toml" + content = "[[hooks.custom]]\ncommand = 'keep'\n\n[[hooks.session_start]]\ncommand = 'remove'\nspeckit_marker = true\n" + dst.write_text(content, encoding="utf-8") + + result = _remove_toml_entries(dst) + + assert result is False # file still has user content + remaining = dst.read_text(encoding="utf-8") + assert "custom" in remaining + assert "speckit_marker" not in remaining + + def test_original_intact_on_write_failure(self, tmp_path: Path) -> None: + from specify_cli.events import _remove_toml_entries + + dst = tmp_path / "hooks.toml" + # Must have BOTH speckit AND user content so the function tries to write + original = "[[hooks.custom]]\ncommand = 'keep'\n\n[[hooks.session_start]]\ncommand = 'test'\nspeckit_marker = true\n" + dst.write_text(original, encoding="utf-8") + + with patch("specify_cli.events.os.replace", side_effect=OSError("fail")): + with pytest.raises(OSError, match="fail"): + _remove_toml_entries(dst) + + assert dst.read_text(encoding="utf-8") == original + + +# -- save_init_options (_init_options.py) ------------------------------------ + +class TestSaveInitOptionsAtomic: + """Atomic writes for save_init_options in _init_options.py.""" + + def test_successful_write(self, tmp_path: Path) -> None: + from specify_cli._init_options import save_init_options, load_init_options + + options = {"integration": "copilot", "script_type": "sh"} + save_init_options(tmp_path, options) + + loaded = load_init_options(tmp_path) + assert loaded == options + + def test_original_intact_on_write_failure(self, tmp_path: Path) -> None: + from specify_cli._init_options import save_init_options, load_init_options + + original = {"integration": "original"} + save_init_options(tmp_path, original) + + with patch("specify_cli._init_options.os.replace", side_effect=OSError("fail")): + with pytest.raises(OSError, match="fail"): + save_init_options(tmp_path, {"corrupted": True}) + + loaded = load_init_options(tmp_path) + assert loaded == original + + +# -- _atomic_write_text (presets/__init__.py) -------------------------------- + +class TestAtomicWriteTextPresets: + """Atomic writes for _atomic_write_text in presets.""" + + def test_successful_write(self, tmp_path: Path) -> None: + from specify_cli.presets import _atomic_write_text + + path = tmp_path / "command.md" + content = "# My Command\n\nRun this." + + _atomic_write_text(path, content) + + assert path.read_text(encoding="utf-8") == content + + def test_original_intact_on_write_failure(self, tmp_path: Path) -> None: + from specify_cli.presets import _atomic_write_text + + path = tmp_path / "command.md" + path.write_text("original content", encoding="utf-8") + + with patch("specify_cli.presets.os.replace", side_effect=OSError("fail")): + with pytest.raises(OSError, match="fail"): + _atomic_write_text(path, "new content") + + assert path.read_text(encoding="utf-8") == "original content" + + def test_temp_file_cleaned_up_after_failure(self, tmp_path: Path) -> None: + from specify_cli.presets import _atomic_write_text + + path = tmp_path / "command.md" + before = set(tmp_path.iterdir()) + + with patch("specify_cli.presets.os.replace", side_effect=OSError("fail")): + with pytest.raises(OSError): + _atomic_write_text(path, "content") + + after = set(tmp_path.iterdir()) + new_files = after - before + for f in new_files: + assert not f.name.startswith(".command.md.")