"""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()