434 lines
17 KiB
Python
434 lines
17 KiB
Python
"""Cross-turn broker for ManagedLockSet issue/repository FLOCKs."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import errno
|
|
import json
|
|
import os
|
|
import secrets
|
|
import signal
|
|
import socket
|
|
import stat
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
if __package__ in (None, ""):
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
|
|
from harness.lib.managed_lock import LockTimeoutError, ManagedLockSet, probe_free
|
|
|
|
ROLES = {"gitea-impl-agent", "gitea-spec-agent", "gitea-review-agent"}
|
|
MAX_TTL = 3600
|
|
|
|
|
|
def _abs(value: str, name: str) -> Path:
|
|
p = Path(value)
|
|
if not p.is_absolute() or "\x00" in value:
|
|
raise ValueError(f"{name} must be absolute")
|
|
return p
|
|
|
|
|
|
def _safe_dir(path: Path, name: str) -> None:
|
|
try:
|
|
fd = _open_directory(path)
|
|
except OSError as exc:
|
|
raise ValueError(f"unsafe {name} parent") from exc
|
|
try:
|
|
st = os.fstat(fd)
|
|
if (not stat.S_ISDIR(st.st_mode) or st.st_uid != os.getuid()
|
|
or stat.S_IMODE(st.st_mode) != 0o700):
|
|
raise ValueError(f"unsafe {name} parent")
|
|
finally:
|
|
os.close(fd)
|
|
|
|
|
|
def _open_directory(path: Path) -> int:
|
|
parts = path.parts
|
|
fd = os.open(parts[0], os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW | os.O_CLOEXEC)
|
|
try:
|
|
for part in parts[1:]:
|
|
nxt = os.open(part, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW | os.O_CLOEXEC,
|
|
dir_fd=fd)
|
|
os.close(fd)
|
|
fd = nxt
|
|
return fd
|
|
except BaseException:
|
|
os.close(fd)
|
|
raise
|
|
|
|
|
|
def _open_parent(path: Path) -> tuple[int, str]:
|
|
return _open_directory(path.parent), path.name
|
|
|
|
|
|
def _validate_handle_fd(fd: int, path: Path) -> os.stat_result:
|
|
fd_st = os.fstat(fd)
|
|
if (not stat.S_ISREG(fd_st.st_mode)
|
|
or fd_st.st_uid != os.getuid()
|
|
or stat.S_IMODE(fd_st.st_mode) != 0o600
|
|
or fd_st.st_nlink != 1):
|
|
raise ValueError("unsafe handle file")
|
|
return fd_st
|
|
|
|
|
|
def _open_handle(path: Path, flags: int) -> int:
|
|
parent_fd, name = _open_parent(path)
|
|
try:
|
|
fd = os.open(name, flags | os.O_NOFOLLOW | os.O_CLOEXEC, dir_fd=parent_fd)
|
|
except OSError as exc:
|
|
if exc.errno in (errno.ELOOP, errno.ENOTDIR):
|
|
raise ValueError("unsafe handle file") from exc
|
|
raise
|
|
try:
|
|
_validate_handle_fd(fd, path)
|
|
return fd
|
|
except BaseException:
|
|
os.close(fd)
|
|
raise
|
|
finally:
|
|
os.close(parent_fd)
|
|
|
|
|
|
def _safe_existing_regular(path: Path) -> None:
|
|
fd = _open_handle(path, os.O_RDONLY)
|
|
os.close(fd)
|
|
|
|
|
|
def _write_all(fd: int, data: bytes) -> None:
|
|
written = 0
|
|
while written < len(data):
|
|
written += os.write(fd, data[written:])
|
|
|
|
|
|
def _write_handle(path: Path, state: dict, *, exclusive: bool = False, on_created=None) -> None:
|
|
flags = os.O_WRONLY | os.O_CREAT | (os.O_EXCL if exclusive else 0)
|
|
if not exclusive:
|
|
fd = _open_handle(path, flags)
|
|
parent_fd = None
|
|
else:
|
|
parent_fd, name = _open_parent(path)
|
|
fd = None
|
|
try:
|
|
fd = os.open(name, flags | os.O_NOFOLLOW | os.O_CLOEXEC, 0o600, dir_fd=parent_fd)
|
|
created = os.fstat(fd)
|
|
if on_created is not None:
|
|
on_created((created.st_dev, created.st_ino))
|
|
except BaseException:
|
|
if fd is not None:
|
|
os.close(fd)
|
|
raise
|
|
finally:
|
|
os.close(parent_fd)
|
|
try:
|
|
if exclusive:
|
|
os.fchmod(fd, 0o600)
|
|
_validate_handle_fd(fd, path)
|
|
os.lseek(fd, 0, os.SEEK_SET)
|
|
os.ftruncate(fd, 0)
|
|
_write_all(fd, (json.dumps(state, sort_keys=True, separators=(",", ":")) + "\n").encode())
|
|
os.fsync(fd)
|
|
finally:
|
|
os.close(fd)
|
|
|
|
|
|
def _update_handle(path: Path, state: dict) -> None:
|
|
fd = _open_handle(path, os.O_WRONLY)
|
|
try:
|
|
os.lseek(fd, 0, os.SEEK_SET)
|
|
os.ftruncate(fd, 0)
|
|
_write_all(fd, (json.dumps(state, sort_keys=True, separators=(",", ":")) + "\n").encode())
|
|
os.fsync(fd)
|
|
finally:
|
|
os.close(fd)
|
|
|
|
|
|
def _read_handle(path: Path) -> dict:
|
|
fd = _open_handle(path, os.O_RDONLY)
|
|
try:
|
|
chunks = []
|
|
while True:
|
|
chunk = os.read(fd, 65536)
|
|
if not chunk:
|
|
break
|
|
chunks.append(chunk)
|
|
data = json.loads(b"".join(chunks).decode())
|
|
finally:
|
|
os.close(fd)
|
|
if not isinstance(data, dict) or not isinstance(data.get("token"), str) or not isinstance(data.get("socket"), str):
|
|
raise ValueError("invalid handle")
|
|
return data
|
|
|
|
|
|
def read_handle_metadata(path: os.PathLike[str] | str) -> dict:
|
|
"""Return non-secret, authenticated-handle identity metadata only."""
|
|
data = _read_handle(_abs(str(path), "handle"))
|
|
return {key: value for key, value in data.items() if key != "token"}
|
|
|
|
|
|
class Broker:
|
|
"""Broker locks are trusted by OS UID; random token authenticates the holder.
|
|
|
|
Same-UID role strings are metadata, not authorization: the same OS UID is the
|
|
security boundary and an unpredictable token rejects non-holder requests.
|
|
"""
|
|
def __init__(self, args: argparse.Namespace):
|
|
if args.role not in ROLES:
|
|
raise ValueError("invalid role")
|
|
if not isinstance(args.run_id, str) or not args.run_id or "\x00" in args.run_id or "/" in args.run_id or "\\" in args.run_id:
|
|
raise ValueError("invalid run_id")
|
|
if not isinstance(args.issue, int) or args.issue <= 0:
|
|
raise ValueError("issue must be positive")
|
|
if not 1 <= args.ttl_seconds <= MAX_TTL:
|
|
raise ValueError("ttl_seconds out of bounds")
|
|
self.root = _abs(args.root, "root")
|
|
self.handle = _abs(args.handle, "handle")
|
|
self.sock_path = _abs(args.socket, "socket")
|
|
locks = self.root / "repositories" / args.owner / args.repository / "locks"
|
|
_safe_dir(locks, "repository locks")
|
|
_safe_dir(self.handle.parent, "handle")
|
|
_safe_dir(self.sock_path.parent, "socket")
|
|
self.token = secrets.token_hex(32)
|
|
self.deadline = time.monotonic() + args.ttl_seconds
|
|
self.set = ManagedLockSet.for_issue(self.root, args.owner, args.repository, args.issue, timeout=args.acquire_timeout)
|
|
self.meta = {"root": str(self.root), "owner": args.owner, "repository": args.repository,
|
|
"issue": args.issue, "role": args.role, "run_id": args.run_id,
|
|
"pid": os.getpid(), "deadline": self.deadline}
|
|
self.state = {**self.meta, "socket": str(self.sock_path), "status": "held", "token": self.token}
|
|
self.server: socket.socket | None = None
|
|
self.socket_identity: tuple[int, int] | None = None
|
|
self.socket_parent_fd: int | None = None
|
|
self.stopping = False
|
|
self.handle_created = False
|
|
self.handle_identity: tuple[int, int] | None = None
|
|
self.started = False
|
|
|
|
def _record_handle_identity(self, identity: tuple[int, int]) -> None:
|
|
self.handle_identity = identity
|
|
self.handle_created = True
|
|
|
|
def persist(self, status: str, *, retain_token: bool = True) -> None:
|
|
self.state["status"] = status
|
|
self.state["deadline"] = self.deadline
|
|
persisted = dict(self.state)
|
|
if not retain_token:
|
|
persisted.pop("token", None)
|
|
_update_handle(self.handle, persisted)
|
|
|
|
def cleanup(self, status: str) -> bool:
|
|
released = False
|
|
try:
|
|
if self.set.repository.held:
|
|
self.set.repository.release()
|
|
self.set.release_issue()
|
|
released = True
|
|
except Exception:
|
|
released = False
|
|
if released and self.handle_created:
|
|
try:
|
|
self.persist(status, retain_token=False)
|
|
except Exception:
|
|
released = False
|
|
if self.server:
|
|
self.server.close()
|
|
self.server = None
|
|
if self.socket_identity is not None:
|
|
try:
|
|
if self.socket_parent_fd is not None:
|
|
st = os.stat(self.sock_path.name, dir_fd=self.socket_parent_fd, follow_symlinks=False)
|
|
if (st.st_dev, st.st_ino) == self.socket_identity:
|
|
os.unlink(self.sock_path.name, dir_fd=self.socket_parent_fd)
|
|
except (FileNotFoundError, OSError):
|
|
pass
|
|
if self.handle_created and not self.started:
|
|
try:
|
|
parent_fd, name = _open_parent(self.handle)
|
|
try:
|
|
st = os.stat(name, dir_fd=parent_fd, follow_symlinks=False)
|
|
if self.handle_identity == (st.st_dev, st.st_ino):
|
|
try:
|
|
self.persist(status, retain_token=False)
|
|
except Exception:
|
|
os.unlink(name, dir_fd=parent_fd)
|
|
finally:
|
|
os.close(parent_fd)
|
|
except OSError:
|
|
pass
|
|
if self.socket_parent_fd is not None:
|
|
try:
|
|
os.close(self.socket_parent_fd)
|
|
except OSError:
|
|
pass
|
|
self.socket_parent_fd = None
|
|
return released
|
|
|
|
def authenticate(self, req: dict) -> bool:
|
|
supplied = req.get("token")
|
|
return isinstance(supplied, str) and secrets.compare_digest(supplied, self.token)
|
|
|
|
def handle_request(self, req: dict) -> dict:
|
|
if not self.authenticate(req):
|
|
return {"error": "unauthorized"}
|
|
if time.monotonic() >= self.deadline:
|
|
self.cleanup("expired")
|
|
self.stopping = True
|
|
return {"error": "expired"}
|
|
action = req.get("action")
|
|
if action == "status":
|
|
return {"status": self.state["status"], "deadline": self.deadline,
|
|
"repository_held": self.set.repository.held, "root": self.state["root"],
|
|
"owner": self.state["owner"], "repository": self.state["repository"],
|
|
"issue": self.state["issue"], "role": self.state["role"],
|
|
"pid": self.state["pid"]}
|
|
if action == "renew":
|
|
ttl = req.get("ttl_seconds")
|
|
if isinstance(ttl, bool) or not isinstance(ttl, (int, float)) or not 1 <= ttl <= MAX_TTL:
|
|
return {"error": "invalid ttl_seconds"}
|
|
self.deadline = time.monotonic() + ttl
|
|
self.persist("repository-held" if self.set.repository.held else "held")
|
|
return {"status": self.state["status"], "deadline": self.deadline}
|
|
if action == "acquire_repository":
|
|
try:
|
|
self.set.repository.acquire()
|
|
except LockTimeoutError:
|
|
return {"error": "repository-busy"}
|
|
self.persist("repository-held")
|
|
return {"status": "repository-held"}
|
|
if action == "release_repository":
|
|
if self.set.repository.held:
|
|
self.set.repository.release()
|
|
self.persist("held")
|
|
return {"status": "held"}
|
|
if action == "release":
|
|
if not self.cleanup("released"):
|
|
return {"error": "release-failed"}
|
|
self.stopping = True
|
|
return {"status": "released"}
|
|
return {"error": "invalid action"}
|
|
|
|
def run(self) -> None:
|
|
try:
|
|
self.set.issue.acquire()
|
|
signal.signal(signal.SIGTERM, lambda *_: setattr(self, "stopping", True))
|
|
signal.signal(signal.SIGINT, lambda *_: setattr(self, "stopping", True))
|
|
self.set.issue.write_metadata(self.meta)
|
|
try:
|
|
_write_handle(self.handle, self.state, exclusive=True,
|
|
on_created=self._record_handle_identity)
|
|
except BaseException:
|
|
raise
|
|
st = self.handle.lstat()
|
|
self.handle_identity = (st.st_dev, st.st_ino)
|
|
self.handle_created = True
|
|
self.socket_parent_fd, socket_name = _open_parent(self.sock_path)
|
|
try:
|
|
os.stat(socket_name, dir_fd=self.socket_parent_fd, follow_symlinks=False)
|
|
raise ValueError("socket path already exists")
|
|
except FileNotFoundError:
|
|
pass
|
|
self.server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
self.server.bind(f"/proc/self/fd/{self.socket_parent_fd}/{socket_name}")
|
|
sock_st = os.fstat(self.server.fileno())
|
|
self.socket_identity = (sock_st.st_dev, sock_st.st_ino)
|
|
os.fchmod(self.server.fileno(), 0o600)
|
|
self.server.listen(8); self.server.settimeout(0.2)
|
|
self.started = True
|
|
print(json.dumps({"status": "held", "handle": str(self.handle), "socket": str(self.sock_path)}, separators=(",", ":")), flush=True)
|
|
while not self.stopping:
|
|
if time.monotonic() >= self.deadline:
|
|
self.stopping = True
|
|
self.cleanup("expired"); break
|
|
try:
|
|
conn, _ = self.server.accept()
|
|
except socket.timeout:
|
|
continue
|
|
with conn:
|
|
conn.settimeout(1)
|
|
line = conn.makefile("rb").readline(65536)
|
|
try:
|
|
req = json.loads(line.decode())
|
|
response = self.handle_request(req) if isinstance(req, dict) else {"error": "invalid request"}
|
|
except Exception:
|
|
response = {"error": "invalid request"}
|
|
conn.sendall((json.dumps(response, separators=(",", ":")) + "\n").encode())
|
|
finally:
|
|
if self.set.issue.held or self.handle_created or self.server:
|
|
self.cleanup("released" if not self.stopping else "expired")
|
|
|
|
|
|
def serve(args: argparse.Namespace) -> int:
|
|
broker = Broker(args)
|
|
broker.run()
|
|
return 0
|
|
|
|
|
|
def request_handle(handle_path: str | Path, action: str, ttl: float | None = None) -> dict:
|
|
"""Send an authenticated request using a validated handle, without printing."""
|
|
handle = _abs(str(handle_path), "handle")
|
|
data = _read_handle(handle)
|
|
sock_path = _abs(data["socket"], "socket")
|
|
req = {"action": action, "token": data["token"]}
|
|
if ttl is not None:
|
|
req["ttl_seconds"] = ttl
|
|
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
|
|
sock.settimeout(3); sock.connect(str(sock_path)); sock.sendall((json.dumps(req) + "\n").encode())
|
|
with sock.makefile("rb") as response_file:
|
|
response = json.loads(response_file.readline())
|
|
if not isinstance(response, dict):
|
|
raise ValueError("invalid broker response")
|
|
return response
|
|
|
|
|
|
def request(args: argparse.Namespace) -> int:
|
|
response = request_handle(args.handle, args.action, args.ttl_seconds)
|
|
print(json.dumps(response, separators=(",", ":")))
|
|
return 0 if "error" not in response else 1
|
|
|
|
|
|
def probe(args: argparse.Namespace) -> int:
|
|
root = _abs(args.root, "root")
|
|
issue_path, repository_path = ManagedLockSet.paths_for_issue(
|
|
root, args.owner, args.repository, args.issue)
|
|
_safe_dir(issue_path.parent, "repository locks")
|
|
issue_free = probe_free(issue_path)
|
|
repository_free = probe_free(repository_path)
|
|
if args.resource == "issue":
|
|
free = issue_free
|
|
elif args.resource == "repository":
|
|
free = repository_free
|
|
else:
|
|
free = issue_free and repository_free
|
|
print(json.dumps({"free": free, "issue_free": issue_free,
|
|
"repository_free": repository_free}, separators=(",", ":")))
|
|
return 0 if free else 1
|
|
|
|
|
|
def parser() -> argparse.ArgumentParser:
|
|
p = argparse.ArgumentParser()
|
|
sub = p.add_subparsers(dest="command", required=True)
|
|
s = sub.add_parser("serve")
|
|
s.add_argument("--root", required=True); s.add_argument("--owner", required=True); s.add_argument("--repository", required=True); s.add_argument("--issue", required=True, type=int); s.add_argument("--role", required=True); s.add_argument("--run-id", required=True); s.add_argument("--handle", required=True); s.add_argument("--socket", required=True); s.add_argument("--ttl-seconds", required=True, type=float); s.add_argument("--acquire-timeout", type=float, default=0.0)
|
|
r = sub.add_parser("request")
|
|
r.add_argument("--handle", required=True); r.add_argument("--action", required=True, choices=["status", "renew", "acquire_repository", "release_repository", "release"]); r.add_argument("--ttl-seconds", type=float)
|
|
check = sub.add_parser("probe")
|
|
check.add_argument("--root", required=True); check.add_argument("--owner", required=True); check.add_argument("--repository", required=True); check.add_argument("--issue", required=True, type=int)
|
|
check.add_argument("--resource", choices=["issue", "repository", "all"], default="all")
|
|
return p
|
|
|
|
|
|
def main() -> int:
|
|
try:
|
|
args = parser().parse_args()
|
|
if args.command == "serve":
|
|
return serve(args)
|
|
if args.command == "request":
|
|
return request(args)
|
|
return probe(args)
|
|
except (ValueError, OSError, LockTimeoutError, json.JSONDecodeError, KeyError, TypeError) as exc:
|
|
print(f"error: {type(exc).__name__}", file=sys.stderr)
|
|
return 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|