fix: harden rollback restore checkpoint reads

This commit is contained in:
hinotoi-agent committed 2026-06-18 16:04:39 +08:00
1 parent b6f8fc2a5c
commit f42248d321
2 files changed
+74 -53

No files matched your search

+56 -53
View File
@@ -97,22 +97,28 @@ def _find_git() -> str:
return shutil.which("git") or "git" return shutil.which("git") or "git"
def _checkpoint_entry_mode(git: str, ckpt_dir: Path, rel_path: str) -> int | None: def _checkpoint_entry_modes(git: str, ckpt_dir: Path) -> dict[str, int]:
"""Return the git index mode for a tracked checkpoint path.""" """Return git index modes for tracked checkpoint paths in one pass."""
result = subprocess.run( result = subprocess.run(
[git, "-C", str(ckpt_dir), "ls-files", "-s", "--", rel_path], [git, "-C", str(ckpt_dir), "ls-files", "-s"],
capture_output=True, text=True, timeout=10, capture_output=True, text=True, timeout=10,
) )
if result.returncode != 0 or not result.stdout.strip(): if result.returncode != 0:
return None raise ValueError("Failed to list checkpoint files")
first = result.stdout.splitlines()[0].split(maxsplit=1)[0]
try: modes: dict[str, int] = {}
return int(first, 8) for line in result.stdout.splitlines():
except ValueError: parts = line.split(maxsplit=3)
return None if len(parts) != 4:
continue
try:
modes[parts[3]] = int(parts[0], 8)
except ValueError:
continue
return modes
def _checkpoint_entry_is_regular(git: str, ckpt_dir: Path, rel_path: str) -> bool: def _checkpoint_entry_is_regular(modes: dict[str, int], rel_path: str) -> bool:
"""Return True only for regular tracked checkpoint files. """Return True only for regular tracked checkpoint files.
Checkpoint worktrees can contain tracked symlinks. Pathname reads such as Checkpoint worktrees can contain tracked symlinks. Pathname reads such as
@@ -121,38 +127,44 @@ def _checkpoint_entry_is_regular(git: str, ckpt_dir: Path, rel_path: str) -> boo
mode as the source of truth and refuse symlink/special entries before any mode as the source of truth and refuse symlink/special entries before any
filesystem open. filesystem open.
""" """
mode = _checkpoint_entry_mode(git, ckpt_dir, rel_path) mode = modes.get(rel_path)
return mode is not None and stat.S_ISREG(mode) return mode is not None and stat.S_ISREG(mode)
def _read_checkpoint_text(git: str, ckpt_dir: Path, rel_path: str) -> str | None: def _read_checkpoint_blob(git: str, ckpt_dir: Path, rel_path: str) -> bytes | None:
"""Read a regular tracked checkpoint file without following worktree links.""" """Read a regular tracked checkpoint blob without opening the worktree path."""
if not _checkpoint_entry_is_regular(git, ckpt_dir, rel_path):
return None
result = subprocess.run( result = subprocess.run(
[git, "-C", str(ckpt_dir), "show", f"HEAD:{rel_path}"], [git, "-C", str(ckpt_dir), "show", f"HEAD:{rel_path}"],
capture_output=True, text=True, errors="replace", timeout=10, capture_output=True, timeout=10,
) )
if result.returncode != 0: if result.returncode != 0:
return None return None
return result.stdout return result.stdout
def _read_checkpoint_text(git: str, ckpt_dir: Path, modes: dict[str, int], rel_path: str) -> str | None:
"""Read a regular tracked checkpoint file without following worktree links."""
if not _checkpoint_entry_is_regular(modes, rel_path):
return None
blob = _read_checkpoint_blob(git, ckpt_dir, rel_path)
if blob is None:
return None
return blob.decode(errors="replace")
def _read_workspace_text(workspace_root: Path, rel_path: str) -> str | None: def _read_workspace_text(workspace_root: Path, rel_path: str) -> str | None:
"""Read workspace text through the anchored workspace file API. """Read workspace text through the anchored workspace file API.
The normal workspace helpers reject traversal and symlink escapes. Returning The normal workspace helpers reject traversal and symlink escapes. Treat
an empty string for invalid/unreadable files keeps the diff useful without invalid or unreadable files as absent so the diff does not imply that the
echoing bytes from an escaping symlink target. workspace contains an empty file.
""" """
from api.workspace import read_file_content from api.workspace import read_file_content
try: try:
return read_file_content(workspace_root, rel_path)["content"] return read_file_content(workspace_root, rel_path)["content"]
except FileNotFoundError: except (FileNotFoundError, OSError, ValueError):
return None return None
except (OSError, ValueError):
return ""
# ── Public API functions (called from routes.py) ──────────────────────────── # ── Public API functions (called from routes.py) ────────────────────────────
@@ -257,15 +269,12 @@ def get_checkpoint_diff(workspace: str, checkpoint: str) -> dict[str, Any]:
git = _find_git() git = _find_git()
# Get list of files in the checkpoint try:
ls_result = subprocess.run( entry_modes = _checkpoint_entry_modes(git, ckpt_dir)
[git, "-C", str(ckpt_dir), "ls-files"], except ValueError as e:
capture_output=True, text=True, timeout=10, raise ValueError("Failed to list checkpoint files") from e
)
if ls_result.returncode != 0:
raise ValueError("Failed to list checkpoint files")
ckpt_files = [f for f in ls_result.stdout.strip().split("\n") if f] ckpt_files = list(entry_modes)
files_changed = [] files_changed = []
diff_lines = [] diff_lines = []
@@ -273,7 +282,7 @@ def get_checkpoint_diff(workspace: str, checkpoint: str) -> dict[str, Any]:
# Read checkpoint version from the git object database, not via the # Read checkpoint version from the git object database, not via the
# checkout path. This prevents tracked symlinks in the checkpoint from # checkout path. This prevents tracked symlinks in the checkpoint from
# redirecting diff reads to arbitrary host files. # redirecting diff reads to arbitrary host files.
ckpt_content = _read_checkpoint_text(git, ckpt_dir, rel_path) ckpt_content = _read_checkpoint_text(git, ckpt_dir, entry_modes, rel_path)
if ckpt_content is None: if ckpt_content is None:
continue continue
@@ -310,8 +319,8 @@ def get_checkpoint_diff(workspace: str, checkpoint: str) -> dict[str, Any]:
} }
def _restore_checkpoint_file(workspace_root: Path, ckpt_file: Path, rel_path: str) -> None: def _restore_checkpoint_file(workspace_root: Path, rel_path: str, content: bytes, mode: int) -> None:
"""Restore one checkpoint file without following workspace symlinks outward.""" """Restore one checkpoint blob without following checkpoint or workspace symlinks."""
from api.workspace import open_anchored_create_fd, open_anchored_write_fd, safe_resolve_ws from api.workspace import open_anchored_create_fd, open_anchored_write_fd, safe_resolve_ws
target = safe_resolve_ws(workspace_root, rel_path) target = safe_resolve_ws(workspace_root, rel_path)
@@ -320,19 +329,13 @@ def _restore_checkpoint_file(workspace_root: Path, ckpt_file: Path, rel_path: st
else: else:
fd = open_anchored_create_fd(workspace_root, target) fd = open_anchored_create_fd(workspace_root, target)
src_stat = ckpt_file.stat()
with os.fdopen(fd, "wb") as out: with os.fdopen(fd, "wb") as out:
with ckpt_file.open("rb") as src: out.write(content)
shutil.copyfileobj(src, out)
out.flush() out.flush()
try: try:
os.fchmod(out.fileno(), src_stat.st_mode & 0o777) os.fchmod(out.fileno(), mode & 0o777)
except (AttributeError, OSError): except (AttributeError, OSError):
logger.debug("Failed to apply restored mode to %s", target, exc_info=True) logger.debug("Failed to apply restored mode to %s", target, exc_info=True)
try:
os.utime(out.fileno(), ns=(src_stat.st_atime_ns, src_stat.st_mtime_ns))
except OSError:
logger.debug("Failed to apply restored timestamp to %s", target, exc_info=True)
def restore_checkpoint(workspace: str, checkpoint: str) -> dict[str, Any]: def restore_checkpoint(workspace: str, checkpoint: str) -> dict[str, Any]:
@@ -356,26 +359,26 @@ def restore_checkpoint(workspace: str, checkpoint: str) -> dict[str, Any]:
git = _find_git() git = _find_git()
# Get list of files in the checkpoint try:
ls_result = subprocess.run( entry_modes = _checkpoint_entry_modes(git, ckpt_dir)
[git, "-C", str(ckpt_dir), "ls-files"], except ValueError as e:
capture_output=True, text=True, timeout=10, raise ValueError("Failed to list checkpoint files") from e
)
if ls_result.returncode != 0:
raise ValueError("Failed to list checkpoint files")
ckpt_files = [f for f in ls_result.stdout.strip().split("\n") if f] ckpt_files = list(entry_modes)
restored = [] restored = []
errors = [] errors = []
for rel_path in ckpt_files: for rel_path in ckpt_files:
ckpt_file = ckpt_dir / rel_path mode = entry_modes.get(rel_path)
if mode is None or not stat.S_ISREG(mode):
continue
if not _checkpoint_entry_is_regular(git, ckpt_dir, rel_path): content = _read_checkpoint_blob(git, ckpt_dir, rel_path)
if content is None:
continue continue
try: try:
_restore_checkpoint_file(workspace_root, ckpt_file, rel_path) _restore_checkpoint_file(workspace_root, rel_path, content, mode)
restored.append(rel_path) restored.append(rel_path)
except (OSError, ValueError) as e: except (OSError, ValueError) as e:
errors.append({"file": rel_path, "error": str(e)}) errors.append({"file": rel_path, "error": str(e)})
@@ -82,3 +82,21 @@ def test_restore_checkpoint_skips_checkpoint_symlink_sources(tmp_path, monkeypat
for p in workspace.rglob("*") for p in workspace.rglob("*")
if p.is_file() if p.is_file()
) )
def test_restore_checkpoint_reads_git_blob_after_worktree_symlink_swap(tmp_path, monkeypatch):
workspace, ckpt_dir, checkpoint = _init_checkpoint(tmp_path, monkeypatch)
_commit_checkpoint_file(ckpt_dir, "file.txt", "checkpoint blob content\n")
secret = tmp_path / "outside-secret.txt"
secret.write_text("POST_COMMIT_SECRET_MARKER_SHOULD_NOT_COPY\n", encoding="utf-8")
(ckpt_dir / "file.txt").unlink()
os.symlink(secret, ckpt_dir / "file.txt")
result = rollback.restore_checkpoint(str(workspace), checkpoint)
assert result["files_restored"] == ["file.txt"]
assert (workspace / "file.txt").read_text(encoding="utf-8") == "checkpoint blob content\n"
assert "POST_COMMIT_SECRET_MARKER_SHOULD_NOT_COPY" not in (workspace / "file.txt").read_text(
encoding="utf-8"
)