apex.ermnvldmr.com/assets/composition-readiness/adapter.py

"""Authenticated readiness endpoint for a fixed, bounded application probe."""
import argparse
import hmac
import http.server
import json
import os
import pathlib
import re
import signal
import subprocess
import threading


class ReadinessServer(http.server.ThreadingHTTPServer):
    daemon_threads = True
    block_on_close = False
    allow_reuse_address = True

    def __init__(self, address, *, token, path, command, timeout=1.5, slots=4):
        if not isinstance(token, str) or not re.fullmatch(r'[!-~]{16,512}', token):
            raise ValueError('A private visible-ASCII token is required')
        if not isinstance(path, str) or not re.fullmatch(r'/[A-Za-z0-9._~/-]+', path) or any(part in ('', '.', '..') for part in path[1:].split('/')):
            raise ValueError('A canonical absolute readiness path is required')
        if not isinstance(command, list) or not command or not all(isinstance(arg, str) and '\0' not in arg for arg in command) or not os.path.isabs(command[0]):
            raise ValueError('A fixed absolute executable and argv are required')
        if isinstance(timeout, bool) or not 0 < timeout <= 1.5 or isinstance(slots, bool) or not isinstance(slots, int) or not 1 <= slots <= 16:
            raise ValueError('Invalid probe bounds')
        self.token = token.encode('ascii')
        self.ready_path = path
        self.command = command
        self.timeout = timeout
        self.probe_slots = threading.BoundedSemaphore(slots)
        self.connection_slots = threading.BoundedSemaphore(slots * 2 + 4)
        self.active = set()
        self.active_lock = threading.Lock()
        self.closing = False
        super().__init__(address, Handler)

    def get_request(self):
        connection, address = super().get_request()
        connection.settimeout(2)
        return connection, address

    def process_request(self, request, address):
        if not self.connection_slots.acquire(blocking=False):
            try:
                request.sendall(b'HTTP/1.1 503 Service Unavailable\r\nCache-Control: no-store\r\nContent-Length: 0\r\nConnection: close\r\n\r\n')
            except OSError:
                pass
            finally:
                self.shutdown_request(request)
            return
        try:
            super().process_request(request, address)
        except BaseException:
            self.connection_slots.release()
            raise

    def process_request_thread(self, request, address):
        try:
            super().process_request_thread(request, address)
        finally:
            self.connection_slots.release()

    def handle_error(self, request, address):
        pass

    @staticmethod
    def stop_probe(process):
        try:
            os.killpg(process.pid, signal.SIGKILL)
        except ProcessLookupError:
            pass
        process.wait()

    def probe(self):
        if not self.probe_slots.acquire(blocking=False):
            return False
        process = None
        try:
            with self.active_lock:
                if self.closing:
                    return False
                process = subprocess.Popen(self.command, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True, close_fds=True)
                self.active.add(process)
            try:
                return process.wait(timeout=self.timeout) == 0
            except subprocess.TimeoutExpired:
                self.stop_probe(process)
                return False
        except (OSError, ValueError):
            return False
        finally:
            if process is not None:
                with self.active_lock:
                    self.active.discard(process)
            self.probe_slots.release()

    def server_close(self):
        with self.active_lock:
            self.closing = True
            active = list(self.active)
        for process in active:
            self.stop_probe(process)
        super().server_close()


class Handler(http.server.BaseHTTPRequestHandler):
    server_version = 'APEX'
    sys_version = ''

    def log_message(self, *args):
        pass

    def send_error(self, code, message=None, explain=None):
        self.respond(code)

    def respond(self, status):
        if self.request_version == 'HTTP/0.9':
            self.request_version = 'HTTP/1.0'
        self.send_response(status)
        self.send_header('Cache-Control', 'no-store')
        self.send_header('Content-Length', '0')
        self.send_header('Connection', 'close')
        self.end_headers()
        self.close_connection = True

    def dispatch(self):
        target = self.requestline.split()[1]
        if target != self.server.ready_path or self.path != self.server.ready_path:
            return self.respond(404)
        if self.command != 'GET':
            return self.respond(405)
        headers = self.headers.get_all('X-APEX-Ready-Token', [])
        supplied = headers[0].encode('latin-1') if len(headers) == 1 else b''
        if len(headers) != 1 or not hmac.compare_digest(supplied, self.server.token):
            return self.respond(403)
        return self.respond(204 if self.server.probe() else 503)

    do_GET = dispatch
    do_HEAD = dispatch
    do_POST = dispatch
    do_PUT = dispatch
    do_DELETE = dispatch
    do_OPTIONS = dispatch
    do_PATCH = dispatch


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', required=True)
    args = parser.parse_args()
    config = json.loads(pathlib.Path(args.config).read_text())
    token = pathlib.Path(config.pop('token_file')).read_text().strip()
    address = (config.pop('host', '127.0.0.1'), config.pop('port', 9092))
    server = ReadinessServer(address, token=token, **config)
    def terminate(signum, frame):
        raise SystemExit(0)
    signal.signal(signal.SIGTERM, terminate)
    try:
        server.serve_forever()
    finally:
        server.server_close()


if __name__ == '__main__':
    main()