gitea-dynamic-runner/src/gitea_dynamic_runner/controller.py

#!/usr/bin/env python3
"""Gitea workflow_job webhook to NATS JetStream producer."""

import hashlib
import hmac
import json
import logging
import os
import ssl
from pathlib import Path

import nats
from aiohttp import web
from nats.js.api import DiscardPolicy, RetentionPolicy, StorageType, StreamConfig
from nats.js.errors import NotFoundError

from .models import IdentityBinding, RunnerRequest


LOG = logging.getLogger(__name__)
SUBJECT_PREFIX = os.environ.get("NATS_SUBJECT_PREFIX", "ci.runner")
STREAM = os.environ.get("NATS_STREAM", "CI_RUNNER")
NATS_URL = os.environ.get("NATS_URL", "tls://nats.ad.ddupan.top:4222")
NATS_USER = os.environ.get("NATS_USER", "ci-producer")
NATS_PASSWORD_FILE = Path(os.environ.get("NATS_PASSWORD_FILE", "/run/secrets/nats/password"))
NATS_CA_FILE = os.environ.get("NATS_CA_FILE", "/etc/ssl/certs/ca-certificates.crt")
WEBHOOK_SECRET_FILE = Path(os.environ.get("WEBHOOK_SECRET_FILE", "/run/secrets/gitea/webhook-secret"))


def accepts(payload: object) -> tuple[bool, str | None]:
    """Compatibility helper for callers that only need acceptance and identity."""
    request = RunnerRequest.from_webhook(payload)
    return (request is not None, str(request.job_id) if request else None)


def valid_signature(body: bytes, signature: str) -> bool:
    expected = hmac.new(WEBHOOK_SECRET_FILE.read_bytes().strip(), body, hashlib.sha256).hexdigest()
    return hmac.compare_digest(signature.removeprefix("sha256="), expected)


async def ensure_stream(js: object) -> None:
    config = StreamConfig(
        name=STREAM,
        subjects=[f"{SUBJECT_PREFIX}.>"],
        retention=RetentionPolicy.WORK_QUEUE,
        storage=StorageType.FILE,
        discard=DiscardPolicy.OLD,
        max_age=24 * 60 * 60,
        max_msgs=10_000,
        max_bytes=256 * 1024 * 1024,
        duplicate_window=24 * 60 * 60,
    )
    try:
        await js.stream_info(STREAM)
    except NotFoundError:
        await js.add_stream(config=config)
    else:
        await js.update_stream(config=config)


async def webhook(request: web.Request) -> web.Response:
    body = await request.read()
    if not valid_signature(body, request.headers.get("X-Gitea-Signature", "")):
        raise web.HTTPUnauthorized()
    try:
        payload = json.loads(body)
    except json.JSONDecodeError as error:
        raise web.HTTPBadRequest(text="invalid JSON\n") from error
    context = _webhook_context(payload)
    LOG.info("workflow_job webhook received %s", context)

    runner_request = RunnerRequest.from_webhook(payload)
    if runner_request is not None:
        subject = f"{SUBJECT_PREFIX}.{runner_request.backend}"
        await request.app["js"].publish(
            subject,
            runner_request.to_json(),
            headers={"Nats-Msg-Id": f"gitea-workflow-job-{runner_request.job_id}-queued"},
        )
        LOG.info(
            "runner request published subject=%s job_id=%d run_id=%d backend=%s "
            "repository=%s job_name=%r labels=%s",
            subject,
            runner_request.job_id,
            runner_request.run_id,
            runner_request.backend,
            runner_request.repository,
            runner_request.job_name,
            runner_request.labels,
        )
        return web.Response(status=202, text="queued\n")

    binding = IdentityBinding.from_webhook(payload)
    if binding is not None:
        subject = f"{SUBJECT_PREFIX}.{binding.backend}.binding"
        await request.app["js"].publish(
            subject,
            binding.to_json(),
            headers={"Nats-Msg-Id": f"gitea-workflow-job-{binding.job_id}-in-progress"},
        )
        LOG.info(
            "identity binding published subject=%s job_id=%d run_id=%d backend=%s "
            "runner_name=%s repository=%s job_name=%r",
            subject,
            binding.job_id,
            binding.run_id,
            binding.backend,
            binding.runner_name,
            binding.repository,
            binding.job_name,
        )
        return web.Response(status=202, text="binding queued\n")

    LOG.warning("workflow_job webhook ignored %s", context)
    return web.Response(status=204)


def _webhook_context(payload: object) -> str:
    if not isinstance(payload, dict):
        return f"payload_type={type(payload).__name__}"
    job = payload.get("workflow_job")
    repository = payload.get("repository")
    job = job if isinstance(job, dict) else {}
    repository = repository if isinstance(repository, dict) else {}
    return (
        f"action={payload.get('action')!r} job_id={job.get('id')!r} "
        f"run_id={job.get('run_id')!r} runner_name={job.get('runner_name')!r} "
        f"repository={repository.get('full_name')!r} job_name={job.get('name')!r} "
        f"labels={job.get('labels')!r}"
    )


async def health(request: web.Request) -> web.Response:
    connected = request.app["nc"].is_connected
    return web.Response(text="ok\n" if connected else "disconnected\n", status=200 if connected else 503)


async def nats_context(app: web.Application):
    tls = ssl.create_default_context(cafile=NATS_CA_FILE)
    nc = await nats.connect(
        NATS_URL,
        user=NATS_USER,
        password=NATS_PASSWORD_FILE.read_text().strip(),
        tls=tls,
        name="microvm-runner-controller",
    )
    app["nc"] = nc
    app["js"] = nc.jetstream()
    await ensure_stream(app["js"])
    yield
    await nc.drain()


def create_app() -> web.Application:
    app = web.Application(client_max_size=1024 * 1024)
    app.cleanup_ctx.append(nats_context)
    app.router.add_post("/webhook", webhook)
    app.router.add_get("/healthz", health)
    return app


def main() -> None:
    logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO"))
    web.run_app(create_app(), host=os.environ.get("LISTEN", "0.0.0.0"), port=int(os.environ.get("PORT", "8787")))


if __name__ == "__main__":
    main()