atlas-iac/services/hermes/scripts/execution_pool_scm.py
2026-08-17 09:55:55 +00:00

451 lines
19 KiB
Python

#!/usr/bin/env python3
"""Ordinal-scoped Atlas SCM boundary; the model container never gets its token."""
from __future__ import annotations
import configparser
import json
import os
import re
import shutil
import stat
import subprocess
import threading
import urllib.error
import urllib.parse
import urllib.request
import uuid
from http.server import BaseHTTPRequestHandler
from pathlib import Path
from typing import Any
from execution_pool_protocol import (
MAX_WIRE_BYTES,
BoundedHTTPServer,
ProtocolError,
parse_wire,
read_key,
verify_envelope,
)
WORKSPACE_ROOT = Path(os.environ.get("HERMES_WORKER_ROOT", "/workspace"))
SCM_ROOT = Path(os.environ.get("HERMES_SCM_STATE_ROOT", "/scm-state"))
KEY_PATH = Path(os.environ.get("HERMES_EXECUTION_POOL_KEY_FILE", "/pool-access/execution-pool-key"))
TOKEN_PATH = Path(os.environ.get("HERMES_GITEA_TOKEN_FILE", "/vault/secrets/gitea-token"))
USERNAME_PATH = Path(os.environ.get("HERMES_GITEA_USERNAME_FILE", "/vault/secrets/gitea-username"))
ASKPASS = os.environ.get(
"HERMES_GITEA_ASKPASS", "/opt/coordinator/execution_pool_askpass.sh"
)
ORDINAL = int(os.environ.get("HERMES_WORKER_ORDINAL", "-1"))
PORT = int(os.environ.get("HERMES_SCM_BOUNDARY_PORT", "9008"))
REPO = re.compile(r"^https://scm\.bstein\.dev/atlas/([A-Za-z0-9_.-]+)\.git$")
BRANCH = re.compile(r"^(?:feature|fix|chore|docs|test|refactor)/[A-Za-z0-9][A-Za-z0-9._/-]{0,119}$")
MAX_BUNDLE_BYTES = 128 * 1024 * 1024
def _private_text(path: Path) -> str:
descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0))
try:
value = os.read(descriptor, 65537).decode("utf-8").strip()
finally:
os.close(descriptor)
if not value or len(value) > 65536:
raise ProtocolError(f"SCM credential is missing or invalid: {path.name}")
return value
def _git_env(authenticated: bool) -> dict[str, str]:
"""Expose credential paths only to explicit boundary-owned network calls."""
sensitive = {
"GIT_ASKPASS", "SSH_ASKPASS", "HERMES_SCM_PASSWORD_FILE",
"HERMES_SCM_USERNAME_FILE",
}
environment = {
name: value for name, value in os.environ.items()
if name not in sensitive and not name.startswith("GIT_CONFIG_")
}
environment["GIT_TERMINAL_PROMPT"] = "0"
if authenticated:
environment.update(
{
"GIT_ASKPASS": ASKPASS,
"HERMES_SCM_PASSWORD_FILE": str(TOKEN_PATH),
"HERMES_SCM_USERNAME_FILE": str(USERNAME_PATH),
"GIT_CONFIG_COUNT": "1",
"GIT_CONFIG_KEY_0": (
"url.http://gitea.gitea.svc.cluster.local:3000/.insteadOf"
),
"GIT_CONFIG_VALUE_0": "https://scm.bstein.dev/",
}
)
return environment
def _run(
*arguments: str,
cwd: Path | None = None,
timeout: int = 300,
authenticated: bool = False,
) -> str:
completed = subprocess.run(
list(arguments), cwd=cwd, env=_git_env(authenticated), stdin=subprocess.DEVNULL,
text=True, capture_output=True, timeout=timeout, check=False,
)
if completed.returncode:
message = (completed.stderr or completed.stdout or "SCM operation failed")[-2000:]
raise RuntimeError(message.strip())
return completed.stdout.strip()
def _regular_text(path: Path, limit: int, encoding: str = "utf-8") -> str:
"""Read bounded model-controlled metadata without following or blocking."""
flags = os.O_RDONLY | os.O_NONBLOCK | getattr(os, "O_NOFOLLOW", 0)
try:
descriptor = os.open(path, flags)
except OSError as error:
raise ProtocolError(f"workspace metadata is unavailable: {path.name}") from error
try:
info = os.fstat(descriptor)
if not stat.S_ISREG(info.st_mode) or info.st_size > limit:
raise ProtocolError(f"workspace metadata is invalid: {path.name}")
raw = os.read(descriptor, limit + 1)
if len(raw) > limit:
raise ProtocolError(f"workspace metadata is oversized: {path.name}")
return raw.decode(encoding)
except UnicodeError as error:
raise ProtocolError(f"workspace metadata is malformed: {path.name}") from error
finally:
os.close(descriptor)
def _binding(envelope: dict[str, Any]) -> tuple[dict[str, Any], str, str, str]:
if envelope["kind"] != "assignment" or int(envelope["worker_ordinal"]) != ORDINAL:
raise ProtocolError("assignment does not belong to this worker ordinal")
payload = envelope["payload"]
if not isinstance(payload, dict):
raise ProtocolError("assignment payload must be an object")
repo = str(payload.get("repo_url") or "")
branch = str(payload.get("branch") or "")
base_branch = str(payload.get("base_branch") or "main")
match = REPO.fullmatch(repo)
if (
not match
or not BRANCH.fullmatch(branch)
or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._/-]{0,119}", base_branch)
or ".." in branch
or "//" in branch
or ".." in base_branch
or "//" in base_branch
):
raise ProtocolError("assignment is outside the Atlas SCM policy")
run_name = f"{envelope['task_id']}-{envelope['run_id']}"
return payload, repo, branch, match.group(1) + ":" + run_name
def workspace_path(envelope: dict[str, Any]) -> Path:
"""Derive a contained path; no caller-provided filesystem path is accepted."""
parts = (str(envelope["board"]), str(envelope["task_id"]), str(envelope["run_id"]))
if any(not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,127}", part) for part in parts):
raise ProtocolError("assignment path binding is invalid")
workspace_root = WORKSPACE_ROOT.resolve()
root = WORKSPACE_ROOT / "runs"
if root.is_symlink():
raise ProtocolError("workspace run root must not be a symlink")
root.mkdir(mode=0o700, parents=True, exist_ok=True)
root = root.resolve()
root.relative_to(workspace_root)
candidate = root.joinpath(*parts)
if candidate.is_symlink():
raise ProtocolError("workspace must not be a symlink")
current = root
for part in parts[:-1]:
current = current / part
if current.is_symlink():
raise ProtocolError("workspace parent must not be a symlink")
current.mkdir(mode=0o700, exist_ok=True)
if current.is_symlink():
raise ProtocolError("workspace parent must not be a symlink")
try:
candidate.resolve(strict=False).relative_to(root)
except ValueError as error:
raise ProtocolError("workspace escaped its ordinal root") from error
return candidate
def _workspace_identity(destination: Path, repo: str, branch: str) -> str:
"""Read identity as data; never execute Git in a model-controlled checkout."""
git_dir = destination / ".git"
if git_dir.is_symlink() or not git_dir.is_dir():
raise ProtocolError("workspace Git metadata must be a private directory")
raw_config = _regular_text(git_dir / "config", 64 * 1024)
parser = configparser.ConfigParser(interpolation=None, strict=True)
try:
parser.read_string(raw_config)
origin = parser.get('remote "origin"', "url")
except (configparser.Error, KeyError, UnicodeError) as error:
raise ProtocolError("workspace Git config cannot prove its origin") from error
if origin != repo:
raise ProtocolError("durable workspace origin does not match assignment")
head = _regular_text(git_dir / "HEAD", 4096).strip()
if head != f"ref: refs/heads/{branch}":
raise ProtocolError("durable workspace branch does not match assignment")
ref_path = git_dir / "refs/heads" / Path(branch)
if ref_path.is_symlink():
raise ProtocolError("workspace branch ref must not be a symlink")
if ref_path.exists():
commit = _regular_text(ref_path, 128, "ascii").strip()
else:
packed = _regular_text(git_dir / "packed-refs", 1024 * 1024, "ascii")
matches = [
line.split(" ", 1)[0]
for line in packed.splitlines()
if line.endswith(f" refs/heads/{branch}")
]
if len(matches) != 1:
raise ProtocolError("workspace branch ref is unavailable")
commit = matches[0]
if not re.fullmatch(r"[0-9a-f]{40,64}", commit):
raise ProtocolError("workspace branch ref is invalid")
return commit
def _private_repo(envelope: dict[str, Any]) -> Path:
root = SCM_ROOT.resolve()
root.mkdir(mode=0o700, parents=True, exist_ok=True)
name = f"{envelope['board']}-{envelope['task_id']}-{envelope['run_id']}.git"
path = root / name
path.resolve(strict=False).relative_to(root)
if path.is_symlink():
raise ProtocolError("private SCM state must not be a symlink")
return path
def _copy_bundle(source: Path, destination: Path) -> None:
"""Copy a bounded regular bundle into boundary-private storage."""
source_fd = os.open(
source, os.O_RDONLY | os.O_NONBLOCK | getattr(os, "O_NOFOLLOW", 0)
)
temporary = destination.with_name(f".{destination.name}.{uuid.uuid4().hex}.tmp")
destination.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
target_fd = -1
try:
info = os.fstat(source_fd)
if not stat.S_ISREG(info.st_mode) or not 0 < info.st_size <= MAX_BUNDLE_BYTES:
raise ProtocolError("submission bundle is empty, oversized, or not regular")
target_fd = os.open(
temporary,
os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0),
0o600,
)
remaining = info.st_size
while remaining:
chunk = os.read(source_fd, min(1024 * 1024, remaining))
if not chunk:
raise ProtocolError("submission bundle ended early")
view = memoryview(chunk)
while view:
view = view[os.write(target_fd, view) :]
remaining -= len(chunk)
os.fsync(target_fd)
except Exception:
temporary.unlink(missing_ok=True)
raise
finally:
os.close(source_fd)
if target_fd >= 0:
os.close(target_fd)
try:
os.replace(temporary, destination)
finally:
temporary.unlink(missing_ok=True)
class Boundary:
def __init__(self, key: bytes):
self.key = key
self.lock = threading.RLock()
def verify(self, raw: Any) -> dict[str, Any]:
return verify_envelope(self.key, raw, expected_kind="assignment")
def checkout(self, envelope: dict[str, Any]) -> dict[str, Any]:
"""Clone/fetch exactly one assigned branch into its private durable path."""
payload, repo, branch, _ = _binding(envelope)
base_branch = str(payload.get("base_branch") or "main")
destination = workspace_path(envelope)
with self.lock:
if (destination / ".git").exists():
return {
"workspace": str(destination), "preserved_dirty_state": True,
"baseline_sha": _workspace_identity(destination, repo, branch),
}
if destination.exists() and any(destination.iterdir()):
raise ProtocolError("workspace is non-empty and unmanaged")
destination.parent.mkdir(parents=True, exist_ok=True)
try:
_run(
"git", "clone", "--single-branch", "--branch", branch,
"--no-tags", repo, str(destination), timeout=900,
authenticated=True,
)
except RuntimeError as error:
if destination.exists() and any(destination.iterdir()):
raise RuntimeError(
"assigned branch clone failed and left state for review"
) from error
_run(
"git", "clone", "--single-branch", "--branch", base_branch,
"--no-tags", repo, str(destination), timeout=900,
authenticated=True,
)
_run("git", "checkout", "-b", branch, cwd=destination)
_run("git", "config", "user.name", "Hermes Execution Worker", cwd=destination)
_run("git", "config", "user.email", "hermes@bstein.dev", cwd=destination)
return {
"workspace": str(destination), "preserved_dirty_state": False,
"baseline_sha": _run("git", "rev-parse", "HEAD", cwd=destination),
}
def submit(self, envelope: dict[str, Any], request: dict[str, Any]) -> dict[str, Any]:
"""Push only the assignment branch and create/reuse its draft pull request."""
_payload, repo, branch, repo_binding = _binding(envelope)
repo_name, _ = repo_binding.split(":", 1)
destination = workspace_path(envelope)
title = str(request.get("title") or f"Hermes task {envelope['task_id']}").strip()[:240]
body = str(request.get("body") or "Automated draft from Hermes execution pool.")[:12000]
if not title or not (destination / ".git").exists():
raise ProtocolError("submission workspace or title is invalid")
with self.lock:
_workspace_identity(destination, repo, branch)
private = _private_repo(envelope)
bundle = private.with_suffix(".bundle")
_copy_bundle(destination / ".git/hermes-submit.bundle", bundle)
if not private.exists():
_run("git", "init", "--bare", str(private))
_run("git", "--git-dir", str(private), "remote", "add", "origin", repo)
else:
_run("git", "--git-dir", str(private), "remote", "set-url", "origin", repo)
try:
_run(
"git", "--git-dir", str(private), "fetch", "origin",
f"+refs/heads/{branch}:refs/remotes/origin/{branch}", timeout=900,
authenticated=True,
)
except RuntimeError as error:
if "couldn't find remote ref" not in str(error).lower():
raise
_run(
"git", "--git-dir", str(private), "fetch", str(bundle),
"HEAD:refs/pool/candidate", timeout=900,
)
_run(
"git", "--git-dir", str(private), "push", "origin",
f"refs/pool/candidate:refs/heads/{branch}", timeout=900,
authenticated=True,
)
pull = self._draft_pull(
repo_name, branch, str(_payload.get("base_branch") or "main"), title, body
)
bundle.unlink(missing_ok=True)
shutil.rmtree(private)
return {"workspace": str(destination), "branch": branch, "pull_request": pull}
@staticmethod
def _api(path: str, data: dict[str, Any] | None = None) -> Any:
token = _private_text(TOKEN_PATH)
request = urllib.request.Request(
"http://gitea.gitea.svc.cluster.local:3000" + path,
data=json.dumps(data, separators=(",", ":")).encode() if data else None,
method="POST" if data else "GET",
headers={
"Authorization": f"token {token}", "Accept": "application/json",
"Content-Type": "application/json", "User-Agent": "hermes-scm-boundary/1",
},
)
with urllib.request.urlopen(request, timeout=30) as response:
body = response.read(1024 * 1024 + 1)
if len(body) > 1024 * 1024:
raise ProtocolError("Gitea response exceeds the SCM boundary limit")
value = json.loads(body)
if not isinstance(value, (dict, list)):
raise ProtocolError("Gitea response has an invalid shape")
return value
def _draft_pull(
self, repo: str, branch: str, base_branch: str, title: str, body: str
) -> str:
query = urllib.parse.urlencode({"state": "open", "head": f"atlas:{branch}", "limit": 10})
existing = self._api(f"/api/v1/repos/atlas/{repo}/pulls?{query}")
if isinstance(existing, list) and existing:
return str(existing[0].get("html_url") or "")
created = self._api(
f"/api/v1/repos/atlas/{repo}/pulls",
{
"base": base_branch,
"head": branch, "title": title, "body": body, "draft": True,
},
)
return str(created.get("html_url") or "")
def garbage_collect(self, envelope: dict[str, Any]) -> dict[str, Any]:
"""Online requests never delete model-controlled or private SCM state."""
return {"removed": False, "eligible": False, "workspace": str(workspace_path(envelope))}
def handler_factory(boundary: Boundary) -> type[BaseHTTPRequestHandler]:
class Handler(BaseHTTPRequestHandler):
server_version = "hermes-scm-boundary/1"
def _reply(self, status: int, value: dict[str, Any]) -> None:
body = json.dumps(value, separators=(",", ":"), sort_keys=True).encode()
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.send_header("Cache-Control", "no-store")
self.end_headers()
self.wfile.write(body)
def do_GET(self) -> None: # noqa: N802
self._reply(200, {"ready": True}) if self.path == "/ready" else self._reply(404, {"error": "not found"})
def do_POST(self) -> None: # noqa: N802
try:
length = int(self.headers.get("Content-Length", "0"))
request = parse_wire(self.rfile.read(length)) if 0 < length <= MAX_WIRE_BYTES else None
if not request or set(request) - {"operation", "assignment", "title", "body"}:
raise ProtocolError("invalid SCM request")
envelope = boundary.verify(request.get("assignment"))
operations = {
"checkout": lambda: boundary.checkout(envelope),
"submit": lambda: boundary.submit(envelope, request),
"gc-check": lambda: boundary.garbage_collect(envelope),
}
operation = str(request.get("operation") or "")
if operation not in operations:
raise ProtocolError("unsupported SCM operation")
self._reply(200, operations[operation]())
except (ProtocolError, RuntimeError, OSError, urllib.error.URLError) as error:
self._reply(409, {"error": str(error)[:2000]})
def log_message(self, _format: str, *_arguments: Any) -> None:
return
return Handler
def main() -> int:
if ORDINAL not in range(3):
raise SystemExit("HERMES_WORKER_ORDINAL must be 0, 1, or 2")
key = read_key(KEY_PATH)
_private_text(TOKEN_PATH)
server = BoundedHTTPServer(
("0.0.0.0", PORT), handler_factory(Boundary(key)), max_workers=4
)
server.serve_forever()
return 0
if __name__ == "__main__":
raise SystemExit(main())