rustfs/scripts/scanner_service.py

#!/usr/bin/env python3
"""Node-local service control for isolated scanner benchmarks (Linux/systemd)."""

import argparse
import fcntl
import hashlib
import json
import os
from pathlib import Path
import re
import stat
import subprocess
import sys
import uuid


def require(condition, message):
    if not condition:
        raise ValueError(message)


def positive(value, name):
    require(type(value) is int and value > 0, "invalid " + name)
    return value


def local_path(value):
    require(isinstance(value, str) and re.fullmatch(r"/[A-Za-z0-9_./-]+", value), "invalid local path")
    path = Path(value)
    require(path == path.resolve(), "local path must be canonical")
    return path


def command(*args):
    # Never include argv, environment-file contents or command output in errors.
    try:
        result = subprocess.run(args, capture_output=True, text=True, timeout=45, check=False)
    except (OSError, subprocess.TimeoutExpired):
        raise ValueError("service command failed or timed out; reconcile with status/stop") from None
    require(result.returncode == 0, "service command failed; reconcile with status/stop")
    return result.stdout


class Service:
    def __init__(self, plan):
        require(type(plan) is dict and set(plan) == {
            "run_id", "root", "binary", "binary_sha256", "environment_file", "args",
            "runtime_seconds", "cpu_percent", "memory_bytes",
        }, "invalid service plan fields")
        require(str(uuid.UUID(plan["run_id"])) == plan["run_id"], "invalid run ID")
        require(os.getuid() != 0, "run as the benchmark user, not root")
        self.plan = plan
        self.root = local_path(plan["root"])
        root_stat = self.root.stat()
        require(self.root.is_dir() and root_stat.st_uid == os.getuid()
                and stat.S_IMODE(root_stat.st_mode) == 0o700, "benchmark root must be private and owned")
        self.binary = local_path(plan["binary"])
        self.environment = local_path(plan["environment_file"])
        require(self.binary.is_relative_to(self.root) and self.environment.is_relative_to(self.root),
                "binary and environment file must be inside benchmark root")
        require(isinstance(plan["binary_sha256"], str)
                and re.fullmatch(r"[0-9a-f]{64}", plan["binary_sha256"]), "invalid binary digest")
        require(type(plan["args"]) is list and all(isinstance(arg, str) and arg
                and not any(ord(c) < 32 or c == "%" for c in arg) for arg in plan["args"]), "invalid service arguments")
        for key in ("runtime_seconds", "cpu_percent", "memory_bytes"):
            positive(plan[key], key)
        identity = hashlib.sha256(json.dumps(plan, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
        self.unit = "rustfs-scanner-" + plan["run_id"] + ".service"
        self.description = "RustFS scanner acceptance " + identity

    def state(self):
        raw = command("systemctl", "show", self.unit,
                      "--property=LoadState,ActiveState,MainPID,ControlGroup,Description,User,WorkingDirectory,Transient")
        state = dict(line.split("=", 1) for line in raw.splitlines() if "=" in line)
        if state.get("LoadState") != "not-found":
            require(state.get("Description") == self.description and state.get("User") == str(os.getuid())
                    and state.get("WorkingDirectory") == str(self.root) and state.get("Transient") == "yes",
                    "service ownership mismatch; refusing operation")
        return state

    def observe(self):
        state = self.state()
        pid = int(state["MainPID"])
        result = {"unit": self.unit, "active": state["ActiveState"], "pid": pid,
                  "stopped": state["ActiveState"] in ("inactive", "failed") and pid == 0,
                  "cpu_max": None, "memory_max": None}
        if state["ActiveState"] == "active":
            require(pid > 0 and state["ControlGroup"] == "/system.slice/" + self.unit, "unexpected service cgroup")
            cgroup = Path("/sys/fs/cgroup") / state["ControlGroup"].lstrip("/")
            result["cpu_max"] = (cgroup / "cpu.max").read_text().strip()
            result["memory_max"] = int((cgroup / "memory.max").read_text())
        return result

    def start(self):
        require(self.state()["LoadState"] == "not-found", "service already exists; use status/stop to reconcile")
        for path in (self.binary, self.environment):
            info = path.stat()
            require(stat.S_ISREG(info.st_mode) and info.st_uid == os.getuid()
                    and not info.st_mode & 0o022, "service input must be an owned regular file")
        require(stat.S_IMODE(self.environment.stat().st_mode) == 0o600, "environment file must have mode 0600")
        require(os.access(self.binary, os.X_OK), "binary is not executable")
        with self.binary.open("rb") as source:
            require(hashlib.file_digest(source, "sha256").hexdigest() == self.plan["binary_sha256"], "binary digest mismatch")
        command("sudo", "-n", "systemd-run", "--quiet", "--collect", "--service-type=exec",
                "--expand-environment=no", "--unit=" + self.unit, "--description=" + self.description,
                "--uid=" + str(os.getuid()), "--gid=" + str(os.getgid()),
                "--property=WorkingDirectory=" + str(self.root),
                "--property=EnvironmentFile=" + str(self.environment),
                "--property=KillMode=control-group", "--property=SendSIGKILL=yes",
                "--property=TimeoutStopSec=30", "--property=Restart=no", "--property=NoNewPrivileges=yes",
                "--property=RuntimeMaxSec=" + str(self.plan["runtime_seconds"]),
                "--property=CPUAccounting=yes", "--property=MemoryAccounting=yes", "--property=IOAccounting=yes",
                "--property=CPUQuota=" + str(self.plan["cpu_percent"]) + "%",
                "--property=MemoryMax=" + str(self.plan["memory_bytes"]),
                "--", str(self.binary), *self.plan["args"])
        result = self.observe()
        require(result["active"] == "active", "service exited before observation")
        return result

    def resize(self, cpu_percent, memory_bytes):
        positive(cpu_percent, "cpu_percent")
        positive(memory_bytes, "memory_bytes")
        before = self.observe()
        require(before["active"] == "active", "cannot resize inactive service")
        command("sudo", "-n", "systemctl", "set-property", "--runtime", self.unit,
                "CPUQuota=" + str(cpu_percent) + "%", "MemoryMax=" + str(memory_bytes))
        after = self.observe()
        require(after["active"] == "active" and after["pid"] == before["pid"], "process changed during resize")
        quota, period = map(int, after["cpu_max"].split())
        require(quota * 100 == period * cpu_percent and after["memory_max"] == memory_bytes,
                "kernel resource limits do not match request")
        return after

    def stop(self):
        if self.state()["LoadState"] != "not-found":
            command("sudo", "-n", "systemctl", "stop", self.unit)
        result = self.observe()
        require(result["stopped"], "service has not stopped")
        return result

    def execute(self, action, cpu_percent=None, memory_bytes=None):
        # Separate SSH controllers serialize mutations without relying on a PID file.
        lock_path = self.root / ("." + self.unit + ".lock")
        fd = os.open(lock_path, os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW, 0o600)
        with os.fdopen(fd, "w") as lock:
            fcntl.flock(lock, fcntl.LOCK_EX)
            if action == "start":
                return self.start()
            if action == "stop":
                return self.stop()
            if action == "resize":
                return self.resize(cpu_percent, memory_bytes)
            require(action == "status", "unknown action")
            return self.observe()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("action", choices=("start", "status", "resize", "stop"))
    parser.add_argument("plan", type=Path, help="immutable plan; credentials belong in its private environment file")
    parser.add_argument("--cpu-percent", type=int)
    parser.add_argument("--memory-bytes", type=int)
    args = parser.parse_args()
    try:
        require(args.plan.stat().st_size <= 65536, "service plan is too large")
        result = Service(json.loads(args.plan.read_text())).execute(args.action, args.cpu_percent, args.memory_bytes)
    except (ValueError, TypeError, AttributeError, KeyError, OSError):
        # A malformed plan or a tool error may contain secret-bearing input.
        print("service control failed; retain the plan and reconcile with status/stop", file=sys.stderr)
        return 1
    print(json.dumps(result, sort_keys=True))
    return 0


if __name__ == "__main__":
    sys.exit(main())