Skip to content
Open
181 changes: 176 additions & 5 deletions src/specify_cli/extensions/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -1343,6 +1343,77 @@ def check_compatibility(

return True

@staticmethod
def _reset_dir(path: Path) -> None:
"""Remove whatever occupies *path* so a fresh directory can take its place.

A symlink is unlinked (never followed); a real directory is rmtree'd; any
other entry is unlinked. Unlike ``rmtree(..., ignore_errors=True)`` this
surfaces failures so a caller never merges into a partially-removed or
still-symlinked path.
"""
if path.is_symlink():
path.unlink()
elif path.is_dir():
shutil.rmtree(path)
elif path.exists():
path.unlink()

def _atomic_restore_config(
self, dest: Path, data: bytes, mode: int, atime: float, mtime: float
) -> None:
"""Restore a preserved config to *dest* atomically and symlink-safely.

Two guards, both required because *dest*'s parent is a fresh copytree
output that a racing process could swap for a symlink:

1. Reject a symlinked ancestor of *dest* (``dest.parent`` and up, under
the project root) via the repo's shared safe-write guard
(``shared_infra._ensure_safe_shared_directory``). Without this,
``mkstemp(dir=dest.parent)`` + ``os.replace`` would create and install
the config (possibly secrets) *inside a symlinked external target*.
2. Write to a sibling temp file, then ``os.replace`` swaps it into the
directory entry — never following a symlink that occupies the leaf
*dest*, and never leaving a partial write.

Mode/timestamps are set on the temp file (independently — either can
succeed alone) so the installed file carries them from the first moment.
"""
from ..shared_infra import _ensure_safe_shared_directory

# (1) Refuse a symlinked ancestor of the destination. dest.parent must
# already exist (a fresh copytree dir, the backup dir, or a rollback dir
# recreated safely just above the call).
_ensure_safe_shared_directory(
self.project_root,
dest.parent,
create=False,
context="extension config directory",
)

# os.replace cannot replace a directory with a file; clear a stray
# (non-symlink) directory at the leaf first. A symlink at the leaf is
# handled by os.replace itself (it swaps the link entry, never follows).
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)
Comment thread
jawwad-ali marked this conversation as resolved.
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,
Expand Down Expand Up @@ -1410,6 +1481,13 @@ def install_from_directory(
f"extension. Install from a copy in a different location instead."
)

# Validate the source ignore file BEFORE any destructive step. It can
# raise on a malformed/invalid-UTF-8 file, and neither the --force
# self.remove() below nor the later rmtree must run if it does — otherwise
# a bad ignore file would uninstall a working extension (and drop its
# registry entry) before validation fails. Reused as-is at the copytree.
ignore_fn = self._load_extensionignore(source_dir)

# Remove existing installation AFTER all validations pass so that a
# validation failure doesn't leave the user with a half-uninstalled
# extension (configs stranded in .backup/).
Expand All @@ -1428,12 +1506,105 @@ def install_from_directory(
backup_config_dir.unlink()
did_remove = self.remove(manifest.id)

# Install extension (dest_dir computed above during self-install guard)
if dest_dir.exists():
shutil.rmtree(dest_dir)
# 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]] = {}
# Require a REAL directory, not a symlink: iterating a symlinked dest_dir
# would follow the link and read arbitrary external files into
# preserved_configs. This path only ever legitimately captures on-disk
# leftover configs inside the real extension directory.
if dest_dir.is_dir() and not dest_dir.is_symlink():
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(),
Comment thread
jawwad-ali marked this conversation as resolved.
st.st_mode,
st.st_atime,
st.st_mtime,
)

ignore_fn = self._load_extensionignore(source_dir)
shutil.copytree(source_dir, dest_dir, ignore=ignore_fn)
# Install extension (dest_dir computed above during self-install guard).
# Preserved configs are the ONLY surviving copy once the old extension dir
# is removed, so guard the whole destructive reinstall against loss:
# 1. The source ignore file is already validated above (before any
# destructive step), so a malformed one never reaches here.
# 2. Stage a durable on-disk backup of the configs.
# 3. Run the destructive rmtree + copytree + restore inside one block;
# on ANY failure, reset dest_dir to a clean, symlink-safe directory
# and roll the configs back, then re-raise. The backup is removed
# only once the configs are provably in place (normal success OR a
# completed rollback) — otherwise it is left on disk so nothing is
# ever lost, even on a rollback failure.
# Each config write is atomic (temp file + os.replace) AND refuses a
# symlinked ancestor of the destination (see _atomic_restore_config).
backup_dir: Path | None = None
if preserved_configs:
backup_dir = Path(
tempfile.mkdtemp(prefix=f".{dest_dir.name}.cfgbak-", dir=dest_dir.parent)
)
for name, (data, mode, atime, mtime) in preserved_configs.items():
self._atomic_restore_config(backup_dir / name, data, mode, atime, mtime)

restored_ok = not preserved_configs
try:
# Reset via _reset_dir (unlink a symlink / rmtree a real dir) rather
# than a bare rmtree, which raises on a symlink or file at dest_dir.
# Handles an unexpected occupant consistently with the rollback path.
self._reset_dir(dest_dir)
shutil.copytree(source_dir, dest_dir, ignore=ignore_fn)
# Restore every preserved config. A failure here is NOT swallowed: a
# config we promised to preserve but could not restore must fail the
# reinstall, not report success with the config silently missing.
for name, (data, mode, atime, mtime) in preserved_configs.items():
self._atomic_restore_config(dest_dir / name, data, mode, atime, mtime)
restored_ok = True
except Exception:
# Roll the configs back into a clean, symlink-safe dest_dir from the
# durable backup so a failed reinstall leaves them intact, then
# re-raise the original error. Reset dest_dir (unlink a symlink /
# rmtree a real dir — surfacing failures, never ignore_errors) and
# recreate it under a verified non-symlink ancestor chain BEFORE
# copying, so recovery can never merge through a symlinked or
# partially-removed directory into an external target.
if backup_dir is not None:
try:
from ..shared_infra import _ensure_safe_shared_directory
self._reset_dir(dest_dir)
_ensure_safe_shared_directory(
self.project_root,
dest_dir,
create=True,
context="extension directory",
)
shutil.copytree(backup_dir, dest_dir, dirs_exist_ok=True)
restored_ok = True
except (OSError, ValueError):
# ValueError covers SymlinkedSharedPathError (a symlinked
# ancestor): leave the durable backup on disk for recovery.
restored_ok = False
raise
finally:
# Drop the backup only when the configs are provably in dest_dir.
if backup_dir is not None and restored_ok and backup_dir.exists():
shutil.rmtree(backup_dir, ignore_errors=True)

# Register commands with AI agents
registered_commands = {}
Expand Down
Loading