"""One bounded native data-copy child in the existing exclusive restore unit.""" import hashlib import json import os import selectors import signal import stat import subprocess import tempfile import time from collections.abc import Iterator from pathlib import Path from kf_oj.contracts.restore_io import RestoreIOReply from . import backup_manifest, file_copy, preflight, restore_io from .policy import DeliveryError, invalid_constant, unique_object MAX_REQUEST_BYTES = 256 * 1024 MAX_REPLY_BYTES = 64 * 1024 MAX_STDERR_BYTES = 16 * 1024 # Only additional list/newline framing beyond the already bounded source metadata. MAX_SPOOL_BYTES = backup_manifest.MAX_METADATA_BYTES + 4 * backup_manifest.MAX_LOGICAL_ENTRIES class BatchSpool: """Consume validated rows into private, completely verified descriptor-only input.""" def __init__(self, rows: list[dict], directory: Path, deadline: float): self.deadline = deadline self.files = len(rows) self.logical_bytes = sum(row["bytes"] for row in rows) self.started = self.complete = False self.size = 0 self.stream = None restore_io.remaining(deadline) if self.files > backup_manifest.MAX_LOGICAL_ENTRIES: raise DeliveryError("recovery planner inventory exceeds bound") try: self.stream = tempfile.TemporaryFile(mode="w+b", dir=directory) # noqa: SIM115 # Ownership spans construction, consumption and every failure path. info = os.fstat(self.stream.fileno()) if ( not stat.S_ISREG(info.st_mode) or info.st_nlink != 0 or (info.st_uid, info.st_gid) != (os.geteuid(), os.getegid()) or stat.S_IMODE(info.st_mode) != 0o600 or os.get_inheritable(self.stream.fileno()) ): raise DeliveryError("private recovery planner descriptor required") expected = hashlib.sha256() batch = None for batch in restore_io.batches(rows): restore_io.remaining(deadline) raw = json.dumps(batch, separators=(",", ":"), ensure_ascii=True).encode() + b"\n" self.size += len(raw) if len(raw) > MAX_REQUEST_BYTES or self.size > MAX_SPOOL_BYTES: raise DeliveryError("recovery planner input exceeds bound") if self.stream.write(raw) != len(raw): raise DeliveryError("recovery planner write incomplete") expected.update(raw) # No row dictionary or duplicate decoded string remains here while # the native child materializes the full logical data inventory. batch = None rows.clear() self.stream.flush() os.fsync(self.stream.fileno()) self.witness = file_copy._identity(os.fstat(self.stream.fileno())) self.digest = expected.hexdigest() self.stream.seek(0) observed = hashlib.file_digest(self.stream, "sha256").hexdigest() self.unchanged() if observed != self.digest: raise DeliveryError("recovery planner complete input verification failed") self.stream.seek(0) restore_io.remaining(deadline) except BaseException: self.close() raise def unchanged(self) -> None: if ( self.stream is None or self.stream.closed or file_copy._identity(os.fstat(self.stream.fileno())) != self.witness ): raise DeliveryError("recovery planner descriptor changed") def __iter__(self) -> Iterator[list[dict]]: if self.started: raise DeliveryError("recovery planner input already consumed") self.started = True files = logical_bytes = 0 observed = hashlib.sha256() assert self.stream is not None while True: restore_io.remaining(self.deadline) self.unchanged() raw = self.stream.readline(MAX_REQUEST_BYTES + 1) if not raw: break if len(raw) > MAX_REQUEST_BYTES or not raw.endswith(b"\n"): raise DeliveryError("recovery planner input framing differs") observed.update(raw) batch = json.loads( raw, object_pairs_hook=unique_object, parse_constant=invalid_constant ) if not isinstance(batch, list) or not 1 <= len(batch) <= restore_io.BATCH_FILES: raise DeliveryError("recovery planner batch shape differs") files += len(batch) logical_bytes += sum(row["bytes"] for row in batch) if files > self.files or logical_bytes > self.logical_bytes: raise DeliveryError("recovery planner inventory differs") yield batch self.unchanged() if (files, logical_bytes, observed.hexdigest()) != ( self.files, self.logical_bytes, self.digest, ): raise DeliveryError("recovery planner complete input differs") restore_io.remaining(self.deadline) self.complete = True def close(self) -> None: if self.stream is not None: self.stream.close() def __enter__(self): return self def __exit__(self, kind, _value, _traceback): try: if kind is None and not self.complete: raise DeliveryError("recovery planner input incomplete") finally: self.close() class Session: """No fallback after partial writes; all verified witnesses stay in the child.""" def __init__( self, payload: Path, backup_id: str, recovery_id: str, expected_files: int, expected_bytes: int, deadline: float, ): self.deadline = deadline self.batch_id = self.files = self.bytes = 0 self.expected_files, self.expected_bytes = expected_files, expected_bytes self.finished = False self.process = None self.stderr_bytes = 0 helper = payload / "bin/kf-restore-io" preflight.protected(helper) info = helper.lstat() if ( not stat.S_ISREG(info.st_mode) or info.st_nlink != 1 or stat.S_IMODE(info.st_mode) != 0o555 ): raise DeliveryError("selected native recovery helper differs") witness = file_copy._identity(info) self.selector = selectors.DefaultSelector() try: restore_io.remaining(self.deadline) self.process = subprocess.Popen( [str(helper), "--stdio-v1"], cwd="/", env={ "PATH": "/usr/bin:/bin", "LC_ALL": "C.UTF-8", "LD_LIBRARY_PATH": f"{payload}/lib:{payload}/python/lib", }, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, start_new_session=True, close_fds=True, ) if file_copy._identity(helper.lstat()) != witness: raise DeliveryError("selected native recovery helper changed") assert self.process.stdin and self.process.stdout and self.process.stderr for handle in ( self.process.stdin, self.process.stdout, self.process.stderr, ): os.set_blocking(handle.fileno(), False) self.selector.register(self.process.stdout, selectors.EVENT_READ, "stdout") self.selector.register(self.process.stderr, selectors.EVENT_READ, "stderr") self._exchange( { "schema_version": 1, "kind": "restore.io.open", "backup_id": backup_id, "recovery_id": recovery_id, "expected_files": expected_files, "expected_logical_bytes": expected_bytes, "deadline_monotonic_ns": int(deadline * 1_000_000_000), }, "opened", ) except BaseException: self.close() raise def _exchange(self, value: dict, stage: str) -> None: restore_io.remaining(self.deadline) if self.process is None or self.finished: raise DeliveryError("native recovery session is closed") raw = json.dumps(value, separators=(",", ":"), ensure_ascii=True).encode() + b"\n" if len(raw) > MAX_REQUEST_BYTES: raise DeliveryError("native recovery request exceeds bound") assert self.process.stdin self.selector.register(self.process.stdin, selectors.EVENT_WRITE, "stdin") offset = 0 output = bytearray() try: while True: restore_io.remaining(self.deadline) ready = self.selector.select(max(0, self.deadline - time.monotonic())) if not ready: raise DeliveryError("native recovery copy deadline exceeded") for key, _ in ready: if key.data == "stdin": try: count = os.write(key.fd, memoryview(raw)[offset:]) except BlockingIOError: continue if count <= 0: raise DeliveryError("native recovery request pipe closed") offset += count if offset == len(raw): self.selector.unregister(self.process.stdin) continue try: block = os.read(key.fd, 8192) except BlockingIOError: continue if key.data == "stderr": self.stderr_bytes += len(block) if self.stderr_bytes > MAX_STDERR_BYTES: raise DeliveryError("native recovery diagnostic exceeds bound") if not block: self.selector.unregister(key.fileobj) continue if not block: raise DeliveryError("native recovery stopped before verification") output.extend(block) if len(output) > MAX_REPLY_BYTES: raise DeliveryError("native recovery reply exceeds bound") if b"\n" in output: if offset != len(raw) or output[-1:] != b"\n" or output.count(b"\n") != 1: raise DeliveryError("native recovery reply framing differs") reply = RestoreIOReply.model_validate( json.loads( output, object_pairs_hook=unique_object, parse_constant=invalid_constant, ) ) if ( reply.stage != stage or reply.error != "none" or ( reply.batch_id, reply.verified_files, reply.verified_logical_bytes, ) != (self.batch_id, self.files, self.bytes) ): raise DeliveryError("native recovery file verification refused") restore_io.remaining(self.deadline) return except (OSError, ValueError, RecursionError) as error: raise DeliveryError("native recovery transport or reply differs") from error finally: if self.process.stdin in self.selector.get_map(): self.selector.unregister(self.process.stdin) def materialize(self, rows: list[dict]) -> None: self._materialize(restore_io.batches(rows)) def materialize_owned(self, rows: list[dict], directory: Path) -> None: with BatchSpool(rows, directory, self.deadline) as batches: self._materialize(iter(batches)) def _materialize(self, batches: Iterator[list[dict]]) -> None: for batch in batches: self.batch_id += 1 self.files += len(batch) self.bytes += sum(item["bytes"] for item in batch) if self.files > self.expected_files or self.bytes > self.expected_bytes: raise DeliveryError("native recovery inventory exceeds original receipt") # The complete private manifest was already strictly validated. # The native decoder independently validates each row before use. self._exchange( { "schema_version": 1, "kind": "restore.io.batch", "batch_id": self.batch_id, "files": batch, }, "verified", ) def finish(self) -> None: if (self.files, self.bytes) != (self.expected_files, self.expected_bytes): raise DeliveryError("native recovery inventory incomplete") self._exchange({"schema_version": 1, "kind": "restore.io.finish"}, "finished") assert self.process is not None and self.process.stdin self.process.stdin.close() restore_io.remaining(self.deadline) try: code = self.process.wait(timeout=max(0, self.deadline - time.monotonic())) except subprocess.TimeoutExpired as error: raise DeliveryError("native recovery child did not join") from error if code != 0: raise DeliveryError("native recovery child failed after verification") restore_io.remaining(self.deadline) self.finished = True def close(self) -> None: process = self.process try: if process is not None and process.poll() is None: # This child starts its own session and never spawns processes. # Its only descendants are its two native threads. try: if os.getpgid(process.pid) != process.pid: raise DeliveryError("native recovery child ownership differs") os.killpg(process.pid, signal.SIGTERM) except ProcessLookupError: pass try: process.wait(timeout=3) except subprocess.TimeoutExpired: try: os.killpg(process.pid, signal.SIGKILL) except ProcessLookupError: pass process.wait(timeout=2) finally: self.selector.close() if process is not None: for handle in (process.stdin, process.stdout, process.stderr): if handle is not None: handle.close()