fix(extensions): preserve keep-config leftover on reinstall (data loss)

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 = {}

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 =====