diff --git a/src/specify_cli/extensions/__init__.py b/src/specify_cli/extensions/__init__.py index b84b836545..aa37c3c9e0 100644 --- a/src/specify_cli/extensions/__init__.py +++ b/src/specify_cli/extensions/__init__.py @@ -1312,6 +1312,44 @@ def check_compatibility( return True + @staticmethod + def _atomic_restore_config( + dest: Path, data: bytes, mode: int, atime: float, mtime: float + ) -> None: + """Restore a preserved config to *dest* atomically (temp file + os.replace). + + Writes to a sibling temp file, then ``os.replace()`` swaps it into place. + ``os.replace`` operates on the directory entry, so it (a) never follows a + symlink that may occupy *dest* — closing the check-then-use race where a + racing process swaps in a symlink to overwrite an external target — and + (b) leaves either the old file or the fully-written new one, never a + partial write. Mirrors ``shared_infra._write_shared_bytes``. Mode and + timestamps are set on the temp file (independently — either can succeed + on its own) so the installed file carries them from the first moment. + """ + # os.replace cannot replace a directory with a file; clear a stray + # (non-symlink) directory at the path first. A symlink is handled by + # os.replace itself (it swaps the link entry without following it). + if dest.is_dir() and not dest.is_symlink(): + shutil.rmtree(dest, ignore_errors=True) + fd, temp_name = tempfile.mkstemp(prefix=f".{dest.name}.", dir=dest.parent) + temp_path = Path(temp_name) + try: + with os.fdopen(fd, "wb") as fh: + fh.write(data) + try: + temp_path.chmod(mode & 0o7777) # permission bits only + except OSError: + pass + try: + os.utime(temp_path, (atime, mtime)) + except OSError: + pass + os.replace(temp_path, dest) + finally: + if temp_path.exists(): + temp_path.unlink() + def install_from_directory( self, source_dir: Path, @@ -1397,12 +1435,74 @@ def install_from_directory( backup_config_dir.unlink() did_remove = self.remove(manifest.id) - # Install extension (dest_dir computed above during self-install guard) + # A prior `remove(keep_config=True)` leaves the top-level + # *-config.yml / *-config.local.yml in place while dropping the + # registry entry. A later reinstall is not a --force removal + # (did_remove is False), so the rmtree below would wipe those + # preserved configs and the backup-restore path never runs. Capture + # them here and write them back after the fresh copytree, mirroring + # the *-config filter the --force restore path uses. + # Capture (bytes, mode, atime, mtime) so restore preserves both + # permissions and timestamps — truly mirroring the --force restore's + # shutil.copy2 (which preserves mode/mtime). Config files may hold + # secrets (API keys); recreating them with default perms could widen + # access. + preserved_configs: dict[str, tuple[bytes, int, float, float]] = {} + if dest_dir.exists(): + for cfg_file in dest_dir.iterdir(): + if ( + cfg_file.is_file() + and not cfg_file.is_symlink() + and ( + cfg_file.name.endswith("-config.yml") + or cfg_file.name.endswith("-config.local.yml") + ) + ): + st = cfg_file.stat() + preserved_configs[cfg_file.name] = ( + cfg_file.read_bytes(), + st.st_mode, + st.st_atime, + st.st_mtime, + ) + + # Install extension (dest_dir computed above during self-install guard). + # rmtree removes the only on-disk copy of the preserved configs; the sole + # remaining copy is in `preserved_configs` (in memory). If copytree then + # fails (permissions, disk full, a bad source entry), write the configs + # back into a clean dest_dir before re-raising so a failed reinstall never + # loses the user's config. if dest_dir.exists(): shutil.rmtree(dest_dir) ignore_fn = self._load_extensionignore(source_dir) - shutil.copytree(source_dir, dest_dir, ignore=ignore_fn) + try: + shutil.copytree(source_dir, dest_dir, ignore=ignore_fn) + except Exception: + if preserved_configs: + if dest_dir.exists(): + shutil.rmtree(dest_dir, ignore_errors=True) + dest_dir.mkdir(parents=True, exist_ok=True) + for name, (data, mode, atime, mtime) in preserved_configs.items(): + try: + self._atomic_restore_config( + dest_dir / name, data, mode, atime, mtime + ) + except OSError: + pass + raise + + # Restore configs preserved from a keep-config leftover (see above), + # preserving the original file mode and timestamps. The --force + # backup-restore below (did_remove) still takes precedence for that + # path, which uses .backup/ rather than an in-place leftover. Each write + # is atomic (temp file + os.replace) so it never follows a symlink that + # may occupy the path (see _atomic_restore_config). + for name, (data, mode, atime, mtime) in preserved_configs.items(): + try: + self._atomic_restore_config(dest_dir / name, data, mode, atime, mtime) + except OSError: + continue # can't safely restore this one — skip rather than clobber # Register commands with AI agents registered_commands = {} diff --git a/tests/test_extensions.py b/tests/test_extensions.py index 2885180360..6d6be42312 100644 --- a/tests/test_extensions.py +++ b/tests/test_extensions.py @@ -1557,6 +1557,164 @@ def test_config_backup_on_remove(self, extension_dir, project_dir): assert backup_file.exists() assert backup_file.read_text() == "test: config" + def test_reinstall_preserves_keep_config_leftover(self, extension_dir, project_dir): + """A reinstall after `remove(keep_config=True)` must not wipe the + preserved config. remove(keep_config=True) leaves the *-config.yml in + place and drops the registry entry; the subsequent reinstall is not a + --force removal, so the unconditional rmtree used to destroy it.""" + manager = ExtensionManager(project_dir) + manager.install_from_directory(extension_dir, "0.1.0", register_commands=False) + + ext_dir = project_dir / ".specify" / "extensions" / "test-ext" + config_file = ext_dir / "test-ext-config.yml" + config_file.write_text("api_key: MY-CUSTOMIZED-VALUE") + # Restrictive perms — a config may hold secrets; the restore must not + # widen them (POSIX only; Windows chmod ignores group/other bits). + config_file.chmod(0o600) + # Distinctive mtime to prove timestamps survive too (mirrors copy2). + old_mtime = 1_600_000_000 + os.utime(config_file, (old_mtime, old_mtime)) + + # keep-config removal: registry entry dropped, config left in place. + assert manager.remove("test-ext", keep_config=True) is True + assert config_file.exists() + + # Reinstall (no --force; registry now reports not-installed). + manager.install_from_directory(extension_dir, "0.1.0", register_commands=False) + + # The user's config must survive the reinstall (was silently wiped). + assert config_file.exists() + assert "MY-CUSTOMIZED-VALUE" in config_file.read_text() + # ...and keep its original mtime (copy2-style timestamp preservation). + assert config_file.stat().st_mtime == old_mtime + # ...and its original (restrictive) mode, not default perms. + if platform.system() != "Windows": + assert config_file.stat().st_mode & 0o777 == 0o600 + + def test_reinstall_config_restore_does_not_follow_symlink( + self, extension_dir, project_dir, monkeypatch + ): + """The keep-config restore must never write *through* a symlink. If the + config path is a symlink at restore time (e.g. swapped in between the + copytree and the restore), the write must not follow it and clobber the + external target — it must replace the link with a fresh regular file.""" + # Probe symlink capability (Windows usually needs privilege). + probe_target = project_dir / "_probe_target" + probe_link = project_dir / "_probe_link" + try: + probe_target.write_text("x") + probe_link.symlink_to(probe_target) + except (OSError, NotImplementedError): + pytest.skip("symlinks not supported in this environment") + finally: + probe_link.unlink(missing_ok=True) + probe_target.unlink(missing_ok=True) + + manager = ExtensionManager(project_dir) + manager.install_from_directory(extension_dir, "0.1.0", register_commands=False) + + ext_dir = project_dir / ".specify" / "extensions" / "test-ext" + config_file = ext_dir / "test-ext-config.yml" + config_file.write_text("api_key: MY-CUSTOMIZED-VALUE") + + # keep-config removal: registry entry dropped, config left in place. + assert manager.remove("test-ext", keep_config=True) is True + + # A victim file outside the extension dir that the symlink points at. + victim = project_dir / "victim.txt" + victim.write_text("DO-NOT-OVERWRITE") + + # Simulate the config path being a symlink to the victim right after the + # fresh copytree (the TOCTOU / symlinks=True case the guard defends). + real_copytree = shutil.copytree + + def copytree_then_symlink(src, dst, *args, **kwargs): + result = real_copytree(src, dst, *args, **kwargs) + link = Path(dst) / "test-ext-config.yml" + if link.exists() or link.is_symlink(): + link.unlink() + link.symlink_to(victim) + return result + + monkeypatch.setattr(shutil, "copytree", copytree_then_symlink) + + # Reinstall: restore must drop the symlink, not follow it. + manager.install_from_directory(extension_dir, "0.1.0", register_commands=False) + + # The victim outside the extension dir is untouched. + assert victim.read_text() == "DO-NOT-OVERWRITE" + # The config path is now a regular file (not a symlink) with the + # preserved content written in place. + assert not config_file.is_symlink() + assert config_file.is_file() + assert "MY-CUSTOMIZED-VALUE" in config_file.read_text() + + def test_reinstall_config_restore_replaces_directory_at_path( + self, extension_dir, project_dir, monkeypatch + ): + """Cross-platform guard check: if a directory occupies the config path at + restore time, write_bytes() would raise (uncaught, crashing the whole + reinstall). The guard must clear the directory and write a fresh regular + file. Exercises the same non-regular-file guard as the symlink case, but + needs no symlink privilege so it runs everywhere (incl. Windows).""" + manager = ExtensionManager(project_dir) + manager.install_from_directory(extension_dir, "0.1.0", register_commands=False) + + ext_dir = project_dir / ".specify" / "extensions" / "test-ext" + config_file = ext_dir / "test-ext-config.yml" + config_file.write_text("api_key: MY-CUSTOMIZED-VALUE") + assert manager.remove("test-ext", keep_config=True) is True + + # Simulate a directory occupying the config path right after copytree. + real_copytree = shutil.copytree + + def copytree_then_dir(src, dst, *args, **kwargs): + result = real_copytree(src, dst, *args, **kwargs) + clash = Path(dst) / "test-ext-config.yml" + clash.mkdir() + (clash / "sentinel").write_text("stray") + return result + + monkeypatch.setattr(shutil, "copytree", copytree_then_dir) + + # Must not raise, and must restore a regular file with preserved content. + manager.install_from_directory(extension_dir, "0.1.0", register_commands=False) + + assert config_file.is_file() + assert "MY-CUSTOMIZED-VALUE" in config_file.read_text() + + def test_reinstall_config_survives_copytree_failure( + self, extension_dir, project_dir, monkeypatch + ): + """The preserved config is captured in memory and the old dir is removed + before copytree runs. If copytree fails, the config must be written back + (not lost) so a failed reinstall leaves the user's config intact.""" + manager = ExtensionManager(project_dir) + manager.install_from_directory(extension_dir, "0.1.0", register_commands=False) + + ext_dir = project_dir / ".specify" / "extensions" / "test-ext" + config_file = ext_dir / "test-ext-config.yml" + config_file.write_text("api_key: MY-CUSTOMIZED-VALUE") + config_file.chmod(0o600) + assert manager.remove("test-ext", keep_config=True) is True + + # Simulate a mid-reinstall failure after the old dir was removed. + def boom(src, dst, *args, **kwargs): + raise OSError("simulated disk-full during copytree") + + monkeypatch.setattr(shutil, "copytree", boom) + + with pytest.raises(OSError): + manager.install_from_directory( + extension_dir, "0.1.0", register_commands=False + ) + + # Config recovered despite the failure (was permanently lost before). + assert config_file.is_file() + assert "MY-CUSTOMIZED-VALUE" in config_file.read_text() + if platform.system() != "Windows": + assert config_file.stat().st_mode & 0o777 == 0o600 + # ===== CommandRegistrar Tests =====