128 lines
5.2 KiB
Python
128 lines
5.2 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: GPL-3.0-or-later
|
|
"""Bytes-only fixed worker result record and generation identity contract."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import hashlib
|
|
import struct
|
|
|
|
|
|
MAGIC = b"CHW10AQ1"
|
|
VERSION = 1
|
|
RECORD_SIZE = 128
|
|
HASHED_SIZE = 96
|
|
MAX_COPY_SIZE = 128 * 1024 * 1024
|
|
STATUS_SUCCESS = 0
|
|
STATUS_COPY_ERROR = 1
|
|
STATUS_RESTORE_ERROR = 2
|
|
STATUS_DEADLINE = 3
|
|
VALID_STATUSES = {STATUS_SUCCESS, STATUS_COPY_ERROR,
|
|
STATUS_RESTORE_ERROR, STATUS_DEADLINE}
|
|
PREFIX = struct.Struct("<8sHHIIIQQQII16s16s8s")
|
|
|
|
|
|
class WorkerRecordError(ValueError):
|
|
"""The fixed worker record is malformed or not precommitted."""
|
|
|
|
|
|
def _token(value: bytes, name: str) -> bytes:
|
|
if type(value) is not bytes or len(value) != 16 or not any(value):
|
|
raise WorkerRecordError(f"{name} must be a nonzero 16-byte token")
|
|
return value
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class WorkerPrecommit:
|
|
attempt_id: bytes
|
|
nonce: bytes
|
|
worker_pid: int
|
|
child_pid: int
|
|
generation: int
|
|
requested: int
|
|
|
|
def __post_init__(self) -> None:
|
|
_token(self.attempt_id, "attempt ID")
|
|
_token(self.nonce, "worker nonce")
|
|
for value in (self.worker_pid, self.child_pid, self.generation,
|
|
self.requested):
|
|
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
|
|
raise WorkerRecordError("precommit integer is invalid")
|
|
if self.worker_pid == self.child_pid or self.worker_pid > 0xffffffff \
|
|
or self.child_pid > 0xffffffff or self.generation > 0xffffffffffffffff \
|
|
or self.requested > MAX_COPY_SIZE:
|
|
raise WorkerRecordError("precommit identity or size is invalid")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class WorkerRecord:
|
|
status: int
|
|
worker_pid: int
|
|
child_pid: int
|
|
generation: int
|
|
requested: int
|
|
copied: int
|
|
restore_failure_bits: int
|
|
flags: int
|
|
attempt_id: bytes
|
|
nonce: bytes
|
|
|
|
@property
|
|
def successful(self) -> bool:
|
|
return (self.status == STATUS_SUCCESS and self.copied == self.requested
|
|
and self.restore_failure_bits == 0 and self.flags == 0)
|
|
|
|
|
|
def encode_record(precommit: WorkerPrecommit, status: int, copied: int,
|
|
restore_failure_bits: int = 0, flags: int = 0) -> bytes:
|
|
"""Encode synthetic bytes; this function has no transport capability."""
|
|
if type(precommit) is not WorkerPrecommit or status not in VALID_STATUSES:
|
|
raise WorkerRecordError("record input is invalid")
|
|
for value in (copied, restore_failure_bits, flags):
|
|
if not isinstance(value, int) or isinstance(value, bool) or value < 0:
|
|
raise WorkerRecordError("record numeric field is invalid")
|
|
if copied > precommit.requested or restore_failure_bits > 3 \
|
|
or flags > 0xffffffff:
|
|
raise WorkerRecordError("record progress, restore bits or flags are invalid")
|
|
if status == STATUS_SUCCESS and (copied != precommit.requested
|
|
or restore_failure_bits or flags):
|
|
raise WorkerRecordError("success record is not exact")
|
|
prefix = PREFIX.pack(MAGIC, VERSION, RECORD_SIZE, status,
|
|
precommit.worker_pid, precommit.child_pid,
|
|
precommit.generation, precommit.requested, copied,
|
|
restore_failure_bits, flags, precommit.attempt_id,
|
|
precommit.nonce, bytes(8))
|
|
if len(prefix) != HASHED_SIZE:
|
|
raise WorkerRecordError("internal record layout differs")
|
|
return prefix + hashlib.sha256(prefix).digest()
|
|
|
|
|
|
def parse_record(raw: bytes, expected: WorkerPrecommit) -> WorkerRecord:
|
|
"""Parse exactly one complete record and enforce the full precommit."""
|
|
if type(raw) is not bytes or len(raw) != RECORD_SIZE \
|
|
or type(expected) is not WorkerPrecommit:
|
|
raise WorkerRecordError("record boundary is invalid")
|
|
prefix, digest = raw[:HASHED_SIZE], raw[HASHED_SIZE:]
|
|
if not hashlib.sha256(prefix).digest() == digest:
|
|
raise WorkerRecordError("record digest differs")
|
|
(magic, version, size, status, worker_pid, child_pid, generation,
|
|
requested, copied, restore_bits, flags, attempt_id, nonce,
|
|
reserved) = PREFIX.unpack(prefix)
|
|
if magic != MAGIC or version != VERSION or size != RECORD_SIZE \
|
|
or status not in VALID_STATUSES or any(reserved):
|
|
raise WorkerRecordError("record header is invalid")
|
|
if (worker_pid, child_pid, generation, requested, attempt_id, nonce) != (
|
|
expected.worker_pid, expected.child_pid, expected.generation,
|
|
expected.requested, expected.attempt_id, expected.nonce):
|
|
raise WorkerRecordError("record identity differs from precommit")
|
|
if copied > requested or restore_bits > 3 or flags != 0:
|
|
raise WorkerRecordError("record result fields are invalid")
|
|
record = WorkerRecord(status, worker_pid, child_pid, generation, requested,
|
|
copied, restore_bits, flags, attempt_id, nonce)
|
|
if status == STATUS_SUCCESS and not record.successful:
|
|
raise WorkerRecordError("success record is incomplete")
|
|
if status == STATUS_RESTORE_ERROR and restore_bits == 0:
|
|
raise WorkerRecordError("restore error has no failure bits")
|
|
return record
|