atlas-iac/services/hermes/scm-common/scripts/scm_broker_server.py

#!/usr/bin/env python3
"""Concurrency-bounded threaded HTTP server for the Hermes SCM broker."""

from __future__ import annotations

import threading
import time
import logging
from http.server import ThreadingHTTPServer

MAX_CONCURRENT_REQUESTS = 3


def validate_ascii_headers(headers, *, maximum_count: int, maximum_bytes: int, error) -> None:
    """Reject oversized or ambiguous broker headers before a request body is read."""
    items = list(headers.items())
    if len(items) > maximum_count:
        raise error("SCM request has too many headers")
    total = 0
    for name, value in items:
        if not name.isascii() or not value.isascii():
            raise error("SCM request headers must be ASCII")
        total += len(name) + len(value) + 4
        if any(ord(character) < 32 and character != "\t" for character in value):
            raise error("SCM request header contains controls")
    if total > maximum_bytes:
        raise error("SCM request headers exceed the safe limit")
    if len(headers.get_all("Content-Length", [])) > 1:
        raise error("SCM request has duplicate Content-Length")


def _rejection_category(phase: str, error: BaseException) -> str:
    """Reduce known pack-policy failures to fixed, non-sensitive categories."""
    if error.__class__.__name__ != "PolicyError":
        return "upstream" if error.__class__.__name__ in {"TimeoutError", "URLError"} else "io"
    if phase != "pack":
        return "policy"
    message = str(error)
    if message == "Git command does not match its task grant":
        return "grant-command"
    if message == "Git update does not prove fast-forward ancestry":
        return "ancestry"
    if message == "Git delta base is outside the pack; push full packs":
        return "thin-pack"
    if "runtime credential material" in message:
        return "credential"
    if "credential-shaped content" in message:
        return "content"
    if "ref command" in message or "limited to namespaced" in message:
        return "ref"
    if "receive-pack" in message or "Git request has" in message:
        return "framing"
    return "policy"


def log_rejection(phase: str, error: BaseException) -> None:
    """Log a fixed operational category without request material or error text."""
    category = _rejection_category(phase, error)
    logging.warning("scm_rejected phase=%s category=%s", phase, category)


class _AbsoluteDeadlineReader:
    """Read header lines against one wall-clock deadline, not idle timeouts."""

    def __init__(self, stream, connection):
        self._stream = stream
        self._connection = connection
        self._deadline: float | None = None

    def begin(self, timeout: float) -> None:
        self._deadline = time.monotonic() + timeout

    def end(self) -> None:
        self._deadline = None

    def readline(self, limit: int = -1) -> bytes:
        if self._deadline is None:
            return self._stream.readline(limit)
        value = bytearray()
        while limit < 0 or len(value) < limit:
            remaining = self._deadline - time.monotonic()
            if remaining <= 0:
                raise TimeoutError("absolute SCM header deadline exceeded")
            self._connection.settimeout(remaining)
            character = self._stream.read(1)
            if not character:
                break
            value.extend(character)
            if character == b"\n":
                break
        return bytes(value)

    def __getattr__(self, name: str):
        return getattr(self._stream, name)


class AbsoluteHeaderDeadlineMixin:
    """Apply an absolute request-line plus header deadline to HTTP handlers."""

    header_deadline_seconds = 10.0

    def setup(self) -> None:
        super().setup()
        self._header_reader = _AbsoluteDeadlineReader(self.rfile, self.connection)
        self.rfile = self._header_reader

    def handle_one_request(self) -> None:
        try:
            self._header_reader.begin(self.header_deadline_seconds)
            try:
                self.raw_requestline = self.rfile.readline(65537)
                if len(self.raw_requestline) > 65536:
                    self.requestline = ""
                    self.request_version = ""
                    self.command = ""
                    self.send_error(414)
                    return
                if not self.raw_requestline:
                    self.close_connection = True
                    return
                if not self.parse_request():
                    return
            finally:
                self._header_reader.end()
            method_name = "do_" + self.command
            if not hasattr(self, method_name):
                self.send_error(501, "Unsupported method")
                return
            getattr(self, method_name)()
            self.wfile.flush()
        except TimeoutError:
            self.close_connection = True


class BoundedThreadingHTTPServer(ThreadingHTTPServer):
    """Cap active handlers so large Git exchanges cannot exhaust memory."""

    daemon_threads = True
    block_on_close = True
    request_queue_size = 16

    def __init__(self, *args, **kwargs):
        self._slots = threading.BoundedSemaphore(MAX_CONCURRENT_REQUESTS)
        super().__init__(*args, **kwargs)

    def process_request(self, request, client_address) -> None:
        if not self._slots.acquire(blocking=False):
            try:
                request.sendall(
                    b"HTTP/1.1 503 Service Unavailable\r\n"
                    b"Connection: close\r\nContent-Length: 0\r\n\r\n"
                )
            finally:
                self.shutdown_request(request)
            return
        try:
            super().process_request(request, client_address)
        except Exception:
            self._slots.release()
            self.shutdown_request(request)
            raise

    def process_request_thread(self, request, client_address) -> None:
        try:
            super().process_request_thread(request, client_address)
        finally:
            self._slots.release()