221 lines
9.4 KiB
Python
221 lines
9.4 KiB
Python
"""One bounded native data-copy child in the existing exclusive restore unit."""
|
|
|
|
import json
|
|
import os
|
|
import selectors
|
|
import signal
|
|
import stat
|
|
import subprocess
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from kf_oj.contracts.restore_io import RestoreIOReply
|
|
|
|
from . import 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
|
|
|
|
|
|
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:
|
|
for batch in restore_io.batches(rows):
|
|
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()
|