106 lines
3.7 KiB
Python
106 lines
3.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Private monotonic runtime-generation state shared by runtime hook processes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import fcntl
|
|
import os
|
|
import stat
|
|
from collections.abc import Mapping
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
MAX_GENERATION: Final = (1 << 63) - 1
|
|
MAX_GENERATION_BYTES: Final = 32
|
|
|
|
|
|
def parse_generation(value: object) -> int:
|
|
if not isinstance(value, str) or not value.isascii() or not value.isdigit():
|
|
raise ValueError("invalid runtime generation")
|
|
generation = int(value)
|
|
if generation < 0 or generation > MAX_GENERATION:
|
|
raise ValueError("invalid runtime generation")
|
|
return generation
|
|
|
|
|
|
def _validate_descriptor(descriptor: int) -> None:
|
|
metadata = os.fstat(descriptor)
|
|
if not stat.S_ISREG(metadata.st_mode):
|
|
raise ValueError("runtime generation state is not a regular file")
|
|
if metadata.st_uid != os.geteuid():
|
|
raise ValueError("runtime generation state has the wrong owner")
|
|
if stat.S_IMODE(metadata.st_mode) & 0o077:
|
|
raise ValueError("runtime generation state permissions are not private")
|
|
if metadata.st_size > MAX_GENERATION_BYTES:
|
|
raise ValueError("runtime generation state is oversized")
|
|
|
|
|
|
def _read_descriptor(descriptor: int) -> int:
|
|
os.lseek(descriptor, 0, os.SEEK_SET)
|
|
raw = os.read(descriptor, MAX_GENERATION_BYTES + 1)
|
|
if len(raw) > MAX_GENERATION_BYTES:
|
|
raise ValueError("runtime generation state is oversized")
|
|
try:
|
|
return parse_generation(raw.decode("ascii").strip())
|
|
except UnicodeDecodeError as exc:
|
|
raise ValueError("invalid runtime generation state") from exc
|
|
|
|
|
|
def _write_descriptor(descriptor: int, generation: int) -> None:
|
|
payload = f"{generation}\n".encode("ascii")
|
|
os.lseek(descriptor, 0, os.SEEK_SET)
|
|
os.ftruncate(descriptor, 0)
|
|
remaining = memoryview(payload)
|
|
while remaining:
|
|
written = os.write(descriptor, remaining)
|
|
if written <= 0:
|
|
raise OSError("runtime generation write made no progress")
|
|
remaining = remaining[written:]
|
|
os.fsync(descriptor)
|
|
|
|
|
|
def initialize_runtime_generation(path: Path, generation: int) -> None:
|
|
if generation < 0 or generation > MAX_GENERATION:
|
|
raise ValueError("invalid runtime generation")
|
|
descriptor = os.open(
|
|
path,
|
|
os.O_RDWR | os.O_CREAT | os.O_TRUNC | os.O_CLOEXEC | os.O_NOFOLLOW,
|
|
0o600,
|
|
)
|
|
try:
|
|
os.fchmod(descriptor, 0o600)
|
|
_validate_descriptor(descriptor)
|
|
_write_descriptor(descriptor, generation)
|
|
finally:
|
|
os.close(descriptor)
|
|
|
|
|
|
def read_runtime_generation(environ: Mapping[str, str]) -> int:
|
|
state_path = environ.get("MOSAIC_LEASE_GENERATION_FILE")
|
|
if not state_path:
|
|
return parse_generation(environ["MOSAIC_RUNTIME_GENERATION"])
|
|
descriptor = os.open(state_path, os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW)
|
|
try:
|
|
_validate_descriptor(descriptor)
|
|
return _read_descriptor(descriptor)
|
|
finally:
|
|
os.close(descriptor)
|
|
|
|
|
|
def bump_runtime_generation(environ: Mapping[str, str]) -> int:
|
|
state_path = environ.get("MOSAIC_LEASE_GENERATION_FILE")
|
|
if not state_path:
|
|
raise ValueError("runtime generation file is required for same-PID rollover")
|
|
descriptor = os.open(state_path, os.O_RDWR | os.O_CLOEXEC | os.O_NOFOLLOW)
|
|
try:
|
|
fcntl.flock(descriptor, fcntl.LOCK_EX)
|
|
_validate_descriptor(descriptor)
|
|
current = _read_descriptor(descriptor)
|
|
if current >= MAX_GENERATION:
|
|
raise ValueError("runtime generation exhausted")
|
|
generation = current + 1
|
|
_write_descriptor(descriptor, generation)
|
|
return generation
|
|
finally:
|
|
os.close(descriptor)
|