mptyl/ThothII
1
0
Fork 0
Code Issues 25 Pull Requests Actions Packages Projects Releases 1 Wiki Activity

fix(preprocess): restrict DWH root claims to writers

Browse Source
This commit is contained in:
marcopan
2026-07-12 07:00:51 +02:00
parent c4dd6c8900
commit 22e806a41b
6 changed files with 238 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
harness/tht/jobs/dwh_pipeline.py
+123 -38
View File
@@ -51,7 +51,22 @@ def _binding_digest(binding: dict[str, str]) -> str:
def _read_root_binding(workspace_root: Path) -> dict[str, str]:
marker = workspace_root / ".tht-dwh" / OWNER_MARKER
try:
payload = json.loads(_read_owned(marker, readonly=True).decode("utf-8"))
fd = os.open(marker, os.O_RDONLY | os.O_NOFOLLOW)
try:
info = os.fstat(fd)
if (
not stat.S_ISREG(info.st_mode)
or info.st_uid != os.getuid()
or info.st_nlink != 1
or stat.S_IMODE(info.st_mode) != 0o400
):
raise OSError("unsafe DWH ownership marker")
chunks = []
while chunk := os.read(fd, 1024 * 1024):
chunks.append(chunk)
finally:
os.close(fd)
payload = json.loads(b"".join(chunks).decode("utf-8"))
binding = payload["binding"]
if (
payload.get("schema_version") != 1
@@ -81,6 +96,9 @@ def _validate_root_binding(workspace_root: Path, expected: dict[str, str]) -> No
raise CorruptCheckpointError("DWH generations directory is invalid") from error
else:
try:
info = os.fstat(generations_fd)
if info.st_uid != os.getuid() or stat.S_IMODE(info.st_mode) != 0o700:
raise OSError("unsafe DWH generations directory")
generation_entries = os.listdir(generations_fd)
finally:
os.close(generations_fd)
@@ -94,25 +112,31 @@ def _claim_or_validate_root_binding(workspace_root: Path, binding: dict[str, str
if marker.exists() or marker.is_symlink():
_validate_root_binding(workspace_root, binding)
return
generations = root / "generations"
generations_nonempty = False
root_fd = os.open(root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
try:
generations_fd = os.open(
generations, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
)
except FileNotFoundError:
pass
entries = set(os.listdir(root_fd))
if not entries <= {"generation.lock", "generations"} or "generation.lock" not in entries:
raise CorruptCheckpointError(
"DWH artifacts are unbound; migrate them explicitly or use an empty root"
)
if "generations" in entries:
generations_fd = os.open(
"generations", os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=root_fd
)
try:
info = os.fstat(generations_fd)
if info.st_uid != os.getuid() or stat.S_IMODE(info.st_mode) != 0o700:
raise OSError("unsafe DWH generations directory")
if os.listdir(generations_fd):
raise CorruptCheckpointError(
"DWH artifacts are unbound; migrate them explicitly or use an empty root"
)
finally:
os.close(generations_fd)
except OSError as error:
raise CorruptCheckpointError("DWH generations directory is invalid") from error
else:
try:
generations_nonempty = bool(os.listdir(generations_fd))
finally:
os.close(generations_fd)
if (root / "ACTIVE").exists() or generations_nonempty:
raise CorruptCheckpointError(
"DWH artifacts are unbound; migrate them explicitly or use an empty root"
)
raise CorruptCheckpointError("DWH workspace root is invalid") from error
finally:
os.close(root_fd)
payload = {
"schema_version": 1,
"binding": binding,
@@ -148,13 +172,9 @@ class DwhSnapshotLease:
def __enter__(self) -> DwhArtifactSnapshot:
binding = config_dwh_binding(self.cfg)
claim_fd = _acquire_generation_lock(self.cfg.paths.artifacts.parent, exclusive=True)
try:
_claim_or_validate_root_binding(self.cfg.paths.artifacts.parent, binding)
finally:
fcntl.flock(claim_fd, fcntl.LOCK_UN)
os.close(claim_fd)
self._fd = _acquire_generation_lock(self.cfg.paths.artifacts.parent, exclusive=False)
self._fd = _acquire_existing_generation_lock(
self.cfg.paths.artifacts.parent, exclusive=False
)
try:
self.snapshot = _resolve_dwh_snapshot_locked(self.cfg, binding)
return self.snapshot
@@ -179,7 +199,12 @@ def _acquire_generation_lock(workspace_root: Path, *, exclusive: bool) -> int:
fd = os.open(root / "generation.lock", os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, 0o600)
try:
info = os.fstat(fd)
if not stat.S_ISREG(info.st_mode) or info.st_uid != os.getuid() or info.st_nlink != 1:
if (
not stat.S_ISREG(info.st_mode)
or info.st_uid != os.getuid()
or info.st_nlink != 1
or stat.S_IMODE(info.st_mode) != 0o600
):
raise OSError("unsafe DWH generation lock")
fcntl.flock(fd, fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH)
return fd
@@ -188,6 +213,37 @@ def _acquire_generation_lock(workspace_root: Path, *, exclusive: bool) -> int:
raise
def _acquire_existing_generation_lock(workspace_root: Path, *, exclusive: bool) -> int:
root = workspace_root / ".tht-dwh"
try:
root_fd = os.open(root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
try:
root_info = os.fstat(root_fd)
if root_info.st_uid != os.getuid() or stat.S_IMODE(root_info.st_mode) != 0o700:
raise OSError("unsafe DWH workspace root")
fd = os.open("generation.lock", os.O_RDWR | os.O_NOFOLLOW, dir_fd=root_fd)
finally:
os.close(root_fd)
try:
info = os.fstat(fd)
if (
not stat.S_ISREG(info.st_mode)
or info.st_uid != os.getuid()
or info.st_nlink != 1
or stat.S_IMODE(info.st_mode) != 0o600
):
raise OSError("unsafe DWH generation lock")
fcntl.flock(fd, fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH)
return fd
except BaseException:
os.close(fd)
raise
except OSError as error:
raise CorruptCheckpointError(
"DWH workspace ownership is not initialized; run preprocessing first"
) from error
def _digest(path: Path) -> str:
return hashlib.sha256(path.read_bytes()).hexdigest()
@@ -289,13 +345,9 @@ def validate_generation(
def resolve_dwh_snapshot(cfg) -> DwhArtifactSnapshot:
binding = config_dwh_binding(cfg)
claim_fd = _acquire_generation_lock(cfg.paths.artifacts.parent, exclusive=True)
try:
_claim_or_validate_root_binding(cfg.paths.artifacts.parent, binding)
finally:
fcntl.flock(claim_fd, fcntl.LOCK_UN)
os.close(claim_fd)
lease_fd = _acquire_generation_lock(cfg.paths.artifacts.parent, exclusive=False)
lease_fd = _acquire_existing_generation_lock(
cfg.paths.artifacts.parent, exclusive=False
)
try:
return _resolve_dwh_snapshot_locked(cfg, binding)
finally:
@@ -307,7 +359,7 @@ def _resolve_dwh_snapshot_locked(
cfg, binding: dict[str, str],
) -> DwhArtifactSnapshot:
_validate_root_binding(cfg.paths.artifacts.parent, binding)
target = active_generation_dir(cfg.paths.artifacts.parent)
target = _active_generation_dir_locked(cfg.paths.artifacts.parent, binding)
if target is None:
return DwhArtifactSnapshot(
None, cfg.paths.artifacts / "mschema" / "physical.yaml", cfg.paths.indexes / "lsh"
@@ -316,7 +368,21 @@ def _resolve_dwh_snapshot_locked(
return DwhArtifactSnapshot(target.name, target / "physical.yaml", target)
def active_generation_dir(workspace_root: Path) -> Path | None:
def active_generation_dir(
workspace_root: Path, expected_binding: dict[str, str]
) -> Path | None:
lease_fd = _acquire_existing_generation_lock(workspace_root, exclusive=False)
try:
return _active_generation_dir_locked(workspace_root, expected_binding)
finally:
fcntl.flock(lease_fd, fcntl.LOCK_UN)
os.close(lease_fd)
def _active_generation_dir_locked(
workspace_root: Path, expected_binding: dict[str, str]
) -> Path | None:
_validate_root_binding(workspace_root, expected_binding)
pointer = workspace_root / ".tht-dwh" / "ACTIVE"
try:
generation = _read_owned(pointer, readonly=False).decode("utf-8").strip()
@@ -329,6 +395,7 @@ def active_generation_dir(workspace_root: Path) -> Path | None:
target = pointer.parent / "generations" / generation
if not target.is_dir() or target.is_symlink():
raise CorruptCheckpointError("active DWH generation is missing")
validate_generation(target, expected_binding)
return target
@@ -381,7 +448,7 @@ class DwhPreprocessPipeline:
def _assert_active_binding(self) -> None:
_validate_root_binding(self.workspace_root, self.binding)
active = active_generation_dir(self.workspace_root)
active = _active_generation_dir_locked(self.workspace_root, self.binding)
if active is not None:
validate_generation(active, self.binding)
@@ -391,8 +458,13 @@ class DwhPreprocessPipeline:
self._validate_steps(steps)
lease_fd = _acquire_generation_lock(self.workspace_root, exclusive=True)
try:
self._assert_no_legacy_artifacts()
_claim_or_validate_root_binding(self.workspace_root, self.binding)
self._assert_active_binding()
active = _active_generation_dir_locked(self.workspace_root, self.binding)
if active is not None:
self.current_physical = active / "physical.yaml"
self.current_lsh_dir = active
finally:
fcntl.flock(lease_fd, fcntl.LOCK_UN)
os.close(lease_fd)
@@ -454,7 +526,7 @@ class DwhPreprocessPipeline:
return set()
target = self.workspace_root / ".tht-dwh" / "generations" / source.run_id
validate_generation(target, self.binding)
active = active_generation_dir(self.workspace_root)
active = active_generation_dir(self.workspace_root, self.binding)
if active != target:
raise CorruptCheckpointError("sealed DWH publication is not ACTIVE")
self._validate_published(target, run_dir / "artifacts", running.artifact_files)
@@ -484,6 +556,19 @@ class DwhPreprocessPipeline:
raise CorruptCheckpointError("resume artifact manifest is invalid") from error
self._validate_published(target, run_dir / "artifacts", required)
def _assert_no_legacy_artifacts(self) -> None:
marker = self.workspace_root / ".tht-dwh" / OWNER_MARKER
if marker.exists() or marker.is_symlink():
return
legacy_physical = self.current_physical is not None and self.current_physical.exists()
legacy_lsh = self.current_lsh_dir is not None and any(
(self.current_lsh_dir / name).exists() for name in self.lsh_filenames
)
if legacy_physical or legacy_lsh:
raise CorruptCheckpointError(
"DWH legacy artifacts are unbound; migrate them explicitly or use an empty root"
)
@staticmethod
def _artifacts(context: JobContext) -> Path:
root = context.run_dir / "artifacts"
@@ -630,7 +715,7 @@ class DwhPreprocessPipeline:
if not root.exists():
return
_validate_root_binding(self.workspace_root, self.binding)
active = active_generation_dir(self.workspace_root)
active = _active_generation_dir_locked(self.workspace_root, self.binding)
if active is not None:
validate_generation(active, self.binding)
active_name = active.name if active else None