feat: supervise stelloauth and cloakbrowser

2026-09-24 17:24:14 +02:00
parent 8a7e9895c8
commit 9410d83053
2 changed files with 1074 additions and 0 deletions
@@ -0,0 +1,386 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import dataclasses
+import json
+import logging
+import os
+import re
+import shutil
+import signal
+import subprocess
+import sys
+import time
+import urllib.parse
+import urllib.request
+from pathlib import Path
+from typing import Callable, NamedTuple
+
+
+OPTIONS_PATH = Path("/data/options.json")
+PROFILE_PATH = Path("/tmp/cloakserve")
+CLOAK_ROOT = "http://127.0.0.1:9222/"
+CLOAK_VERSION = "http://127.0.0.1:9222/json/version?fingerprint=addon-readiness"
+CLOAK_CLOSE = "http://127.0.0.1:9222/fingerprint/addon-readiness/close"
+STELLOAUTH_ROOT = "http://127.0.0.1:8080/"
+DURATION = re.compile(r"^[1-9][0-9]*(?:ms|s|m|h)$")
+CLOAK_COMMAND = [
+    "/usr/local/bin/cloakserve",
+    "--headless=true",
+    "--idle-timeout=30",
+    "--data-dir=/tmp/cloakserve",
+]
+STELLOAUTH_COMMAND = ["/usr/local/bin/stelloauth"]
+
+
+class ConfigError(RuntimeError):
+    pass
+
+
+@dataclasses.dataclass(frozen=True)
+class Options:
+    queue_timeout: str
+    rate_limit_count: int
+    rate_limit_duration: str
+
+
+class HttpResponse(NamedTuple):
+    status: int
+    body: bytes
+
+
+class ReadinessError(RuntimeError):
+    pass
+
+
+def load_options(path: Path) -> Options:
+    try:
+        value = json.loads(path.read_text(encoding="utf-8"))
+        if not isinstance(value, dict) or set(value) != {
+            "queue_timeout",
+            "rate_limit_count",
+            "rate_limit_duration",
+        }:
+            raise ValueError
+
+        queue_timeout = value["queue_timeout"]
+        rate_limit_count = value["rate_limit_count"]
+        rate_limit_duration = value["rate_limit_duration"]
+        if not isinstance(queue_timeout, str) or DURATION.fullmatch(queue_timeout) is None:
+            raise ValueError
+        if (
+            isinstance(rate_limit_count, bool)
+            or not isinstance(rate_limit_count, int)
+            or not 1 <= rate_limit_count <= 20
+        ):
+            raise ValueError
+        if (
+            not isinstance(rate_limit_duration, str)
+            or DURATION.fullmatch(rate_limit_duration) is None
+        ):
+            raise ValueError
+        return Options(queue_timeout, rate_limit_count, rate_limit_duration)
+    except (OSError, UnicodeError, json.JSONDecodeError, KeyError, TypeError, ValueError):
+        raise ConfigError("Invalid add-on configuration") from None
+
+
+def build_environment(options: Options) -> dict[str, str]:
+    environment = os.environ.copy()
+    environment.update(
+        {
+            "CLOAK_CDP_URL": "http://127.0.0.1:9222",
+            "CLOAK_MAX_SESSIONS": "1",
+            "CLOAK_QUEUE_TIMEOUT": options.queue_timeout,
+            "RATE_LIMIT_COUNT": str(options.rate_limit_count),
+            "RATE_LIMIT_DURATION": options.rate_limit_duration,
+            "HTTP_ADDRESS": "0.0.0.0",
+            "PORT": "8080",
+            "METRICS_ADDRESS": "127.0.0.1",
+            "METRICS_PORT": "9090",
+        }
+    )
+    return environment
+
+
+def http_request(method: str, url: str, timeout: float) -> HttpResponse:
+    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
+    data = b"" if method == "POST" else None
+    request = urllib.request.Request(url, data=data, method=method)
+    with opener.open(request, timeout=timeout) as response:
+        return HttpResponse(response.status, response.read())
+
+
+def probe_cloak(
+    request: Callable[[str, str, float], HttpResponse] = http_request,
+) -> None:
+    if request("GET", CLOAK_ROOT, 2.0).status != 200:
+        raise ReadinessError("CloakBrowser root probe failed")
+
+    version = request("GET", CLOAK_VERSION, 2.0)
+    if version.status != 200:
+        raise ReadinessError("CloakBrowser CDP probe failed")
+    try:
+        document = json.loads(version.body)
+        if not isinstance(document, dict):
+            raise ValueError
+        websocket_url = document["webSocketDebuggerUrl"]
+        if (
+            not isinstance(websocket_url, str)
+            or not websocket_url.startswith("ws://127.0.0.1:")
+        ):
+            raise ValueError
+        parsed = urllib.parse.urlsplit(websocket_url)
+        if (
+            parsed.scheme != "ws"
+            or parsed.hostname != "127.0.0.1"
+            or parsed.username is not None
+            or parsed.password is not None
+            or parsed.port is None
+        ):
+            raise ValueError
+    except (json.JSONDecodeError, KeyError, TypeError, UnicodeError, ValueError):
+        raise ReadinessError("CloakBrowser CDP response invalid") from None
+
+    if request("POST", CLOAK_CLOSE, 2.0).status != 200:
+        raise ReadinessError("CloakBrowser readiness profile close failed")
+
+
+def probe_stelloauth(
+    request: Callable[[str, str, float], HttpResponse] = http_request,
+) -> None:
+    if request("GET", STELLOAUTH_ROOT, 2.0).status != 200:
+        raise ReadinessError("Stelloauth root probe failed")
+
+
+def wait_until_ready(
+    name: str,
+    probe: Callable[[], None],
+    timeout: float,
+    stopping: Callable[[], bool],
+    monotonic: Callable[[], float] = time.monotonic,
+    sleep: Callable[[float], None] = time.sleep,
+) -> None:
+    deadline = monotonic() + timeout
+    delay = 0.25
+    while monotonic() < deadline and not stopping():
+        try:
+            probe()
+            return
+        except (OSError, ValueError, ReadinessError):
+            sleep(delay)
+            delay = min(delay * 2, 2.0)
+    raise ReadinessError(f"{name} did not become ready")
+
+
+class ProcessManager:
+    def __init__(
+        self,
+        environment: dict[str, str],
+        *,
+        popen: Callable[..., subprocess.Popen] = subprocess.Popen,
+        killpg: Callable[[int, int], None] = os.killpg,
+        cloak_probe: Callable[[], None] = probe_cloak,
+        stelloauth_probe: Callable[[], None] = probe_stelloauth,
+        monotonic: Callable[[], float] = time.monotonic,
+        sleep: Callable[[float], None] = time.sleep,
+        profile_path: Path = PROFILE_PATH,
+    ) -> None:
+        self.environment = environment
+        self._popen = popen
+        self._killpg = killpg
+        self._cloak_probe = cloak_probe
+        self._stelloauth_probe = stelloauth_probe
+        self._monotonic = monotonic
+        self._sleep = sleep
+        self._profile_path = profile_path
+        self._children: list[tuple[str, subprocess.Popen]] = []
+        self._stopping = False
+        self._term_sent = False
+        self._shutdown_deadline: float | None = None
+
+    def _handle_signal(self, _signum: int, _frame: object) -> None:
+        if self._stopping:
+            return
+        logging.info("Shutdown requested")
+        self._stopping = True
+        self._begin_shutdown()
+
+    def _clean_profiles(self) -> None:
+        logging.info("Cleaning CloakBrowser profiles")
+        if self._profile_path.is_symlink():
+            self._profile_path.unlink()
+        elif self._profile_path.exists():
+            shutil.rmtree(self._profile_path)
+        self._profile_path.mkdir(parents=True, mode=0o700)
+        self._profile_path.chmod(0o700)
+
+    def _spawn(self, name: str, command: list[str]) -> subprocess.Popen:
+        process = self._popen(
+            command,
+            env=self.environment,
+            start_new_session=True,
+        )
+        self._children.append((name, process))
+        if self._stopping:
+            try:
+                self._killpg(process.pid, signal.SIGTERM)
+            except OSError:
+                pass
+        return process
+
+    def _first_exited_child(self) -> tuple[str, int] | None:
+        for name, process in self._children:
+            returncode = process.poll()
+            if returncode is not None:
+                return name, returncode
+        return None
+
+    def _startup_stopping(self) -> bool:
+        return self._stopping or self._first_exited_child() is not None
+
+    def _signal_running(self, sent_signal: int) -> None:
+        for _name, process in self._children:
+            if process.poll() is not None:
+                continue
+            try:
+                self._killpg(process.pid, sent_signal)
+            except OSError:
+                pass
+
+    def _begin_shutdown(self) -> None:
+        if self._shutdown_deadline is None:
+            self._shutdown_deadline = self._monotonic() + 10.0
+        if not self._term_sent:
+            self._signal_running(signal.SIGTERM)
+            self._term_sent = True
+
+    def _reap_all(self) -> None:
+        for _name, process in self._children:
+            try:
+                process.wait(timeout=None)
+            except OSError:
+                logging.error("Child process reap failed")
+
+    def _shutdown(self) -> None:
+        if not self._children:
+            return
+        logging.info("Stopping child processes")
+        self._begin_shutdown()
+        assert self._shutdown_deadline is not None
+        while (
+            any(process.poll() is None for _name, process in self._children)
+            and self._monotonic() < self._shutdown_deadline
+        ):
+            remaining = self._shutdown_deadline - self._monotonic()
+            self._sleep(min(0.25, max(0.0, remaining)))
+
+        if any(process.poll() is None for _name, process in self._children):
+            logging.info("Forcing child processes to stop")
+            self._signal_running(signal.SIGKILL)
+        self._reap_all()
+        logging.info("Child processes stopped")
+
+    @staticmethod
+    def _failure_status(returncode: int) -> int:
+        return returncode if returncode != 0 else 1
+
+    def _startup_failure(self, service: str) -> int:
+        exited = self._first_exited_child()
+        if exited is not None:
+            name, returncode = exited
+            logging.error(f"{name} exited unexpectedly")
+            status = self._failure_status(returncode)
+        else:
+            logging.error(f"{service} readiness failed")
+            status = 1
+        self._shutdown()
+        return status
+
+    def _run(self) -> int:
+        try:
+            self._clean_profiles()
+            if self._stopping:
+                self._shutdown()
+                return 0
+
+            logging.info("Starting CloakBrowser")
+            self._spawn("CloakBrowser", CLOAK_COMMAND)
+            wait_until_ready(
+                "CloakBrowser",
+                self._cloak_probe,
+                60.0,
+                self._startup_stopping,
+                self._monotonic,
+                self._sleep,
+            )
+            if self._stopping:
+                self._shutdown()
+                return 0
+            if self._first_exited_child() is not None:
+                return self._startup_failure("CloakBrowser")
+            logging.info("CloakBrowser ready")
+
+            logging.info("Starting Stelloauth")
+            self._spawn("Stelloauth", STELLOAUTH_COMMAND)
+            wait_until_ready(
+                "Stelloauth",
+                self._stelloauth_probe,
+                30.0,
+                self._startup_stopping,
+                self._monotonic,
+                self._sleep,
+            )
+            if self._stopping:
+                self._shutdown()
+                return 0
+            if self._first_exited_child() is not None:
+                return self._startup_failure("Stelloauth")
+            logging.info("Stelloauth listening on 0.0.0.0:8080")
+        except ReadinessError:
+            if self._stopping:
+                self._shutdown()
+                return 0
+            service = "Stelloauth" if len(self._children) > 1 else "CloakBrowser"
+            return self._startup_failure(service)
+        except OSError:
+            logging.error("Process startup failed")
+            self._shutdown()
+            return 1
+
+        while not self._stopping:
+            exited = self._first_exited_child()
+            if exited is not None:
+                name, returncode = exited
+                logging.error(f"{name} exited unexpectedly")
+                status = self._failure_status(returncode)
+                self._shutdown()
+                return status
+            self._sleep(0.25)
+
+        self._shutdown()
+        return 0
+
+    def run(self) -> int:
+        previous_handlers = {
+            signal.SIGTERM: signal.signal(signal.SIGTERM, self._handle_signal),
+            signal.SIGINT: signal.signal(signal.SIGINT, self._handle_signal),
+        }
+        try:
+            return self._run()
+        finally:
+            for handled_signal, previous_handler in previous_handlers.items():
+                signal.signal(handled_signal, previous_handler)
+
+
+def main() -> int:
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+    try:
+        options = load_options(OPTIONS_PATH)
+    except ConfigError:
+        logging.error("Invalid add-on configuration")
+        return 2
+    return ProcessManager(build_environment(options)).run()
+
+
+if __name__ == "__main__":
+    sys.exit(main())