tai/examples/reverse-proxy.py

#!/usr/bin/env python3
"""Round-robin reverse proxy met healthchecks. Stdlib only, single file, 3.6+.

Run:
    python3 reverse-proxy.py --backends http://127.0.0.1:9001,http://127.0.0.1:9002
    python3 reverse-proxy.py --backends http://127.0.0.1:9001 --port 8080 --health-path /healthz

Spec (dit bestand is de spec):

    Gedrag:
    - Round-robin over gezonde backends (alleen http, geen https-targets).
    - Actieve healthcheck: GET {health-path} elke --health-interval s;
      2xx/3xx is gezond. Ongezonde backend tijdelijk uit rotatie.
    - Passieve healthcheck: connect/read-fout markeert backend direct down
      en de request failovert naar de volgende gezonde backend.
    - Retries: maximaal N pogingen (N = aantal backends), nooit een body
      tweemaal sturen bij POST/PUT tenzij klein genoeg om te bufferen.
    - Hop-header hygiene: Connection, Keep-Alive, Proxy-*, TE, Trailer,
      Transfer-Encoding en Upgrade worden nooit doorgestuurd.
    - X-Forwarded-For (client-IP appended), X-Forwarded-Proto, X-Forwarded-Host
      en Via worden gezet.
    - GET /__proxy__/status: JSON met backend-staten (alleen dit pad is
      van de proxy zelf; de rest wordt 1-op-1 doorgezet).

Caps: MAX_BODY 10 MiB, TIMEOUT 10 s (connect) / 30 s (read),
HEALTH_INTERVAL 10 s, HEALTH_TIMEOUT 3 s, MAX_BACKENDS 32.
"""

import argparse
import json
import socketserver
import threading
import time
from http.client import HTTPConnection, HTTPException
from http.server import BaseHTTPRequestHandler, HTTPServer
from urllib.parse import urlparse

HOP_HEADERS = frozenset([
    "connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
    "te", "trailer", "transfer-encoding", "upgrade",
])
MAX_BODY = 10 * 1024 * 1024
CONNECT_TIMEOUT = 10
READ_TIMEOUT = 30


class Backend(object):
    """Een upstream + zijn gezondheidsstaat."""

    def __init__(self, url):
        parsed = urlparse(url.strip())
        if parsed.scheme != "http" or not parsed.hostname:
            raise ValueError("backend must be http://host[:port][/prefix]: %s"
                             % url)
        self.name = url.strip()
        self.host = parsed.hostname
        self.port = parsed.port or 80
        self.prefix = parsed.path.rstrip("/")
        self._lock = threading.RLock()
        self.healthy = True
        self.last_check = 0.0
        self.last_error = ""
        self.requests = 0
        self.failures = 0

    def mark(self, healthy, error=""):
        with self._lock:
            self.healthy = healthy
            self.last_check = time.time()
            self.last_error = error[:200] if error else ""
            if healthy:
                self.requests += 1
            else:
                self.failures += 1

    def status(self):
        with self._lock:
            return {"name": self.name, "healthy": self.healthy,
                    "last_check": self.last_check,
                    "last_error": self.last_error,
                    "requests": self.requests, "failures": self.failures}


class BackendPool(object):
    """Round-robin over gezonde backends. Thread-safe."""

    def __init__(self, backends):
        self.backends = backends
        self._lock = threading.RLock()
        self._next = 0

    def pick_order(self):
        """Volgorde voor deze request: rotatie-start, gezonde eerst."""
        with self._lock:
            backends = list(self.backends)
            start = self._next % len(backends)
            self._next += 1
        ordered = backends[start:] + backends[:start]
        healthy = [b for b in ordered if b.healthy]
        return healthy + [b for b in ordered if not b.healthy]

    def status(self):
        return {"backends": [b.status() for b in self.backends]}


class HealthChecker(threading.Thread):
    """Achtergrondlus: pollt elke interval het health-path per backend."""

    daemon = True

    def __init__(self, pool, path="/", interval=10, timeout=3):
        super(HealthChecker, self).__init__(name="healthchecker")
        self.pool = pool
        self.path = path or "/"
        self.interval = max(1, interval)
        self.timeout = timeout
        self._stop = threading.Event()

    def stop(self):
        self._stop.set()

    def run(self):
        while not self._stop.wait(self.interval):
            for backend in self.pool.backends:
                self.check_once(backend)

    def check_once(self, backend):
        try:
            conn = HTTPConnection(backend.host, backend.port,
                                  timeout=self.timeout)
            try:
                conn.request("GET", backend.prefix + self.path,
                             headers={"Connection": "close",
                                      "Host": backend.host})
                resp = conn.getresponse()
                resp.read(65536)
                ok = 200 <= resp.status < 400
                backend.mark(True if ok else False,
                             "" if ok else "health %d" % resp.status)
            finally:
                conn.close()
        except Exception as exc:
            backend.mark(False, "%s: %s" % (type(exc).__name__, exc))


class ProxyHandler(BaseHTTPRequestHandler):
    server_version = "TaiProxy/1.0"
    protocol_version = "HTTP/1.1"

    def log_message(self, fmt, *args):
        pass

    def _send(self, code, body=b"", ctype="text/plain; charset=utf-8"):
        if isinstance(body, str):
            body = body.encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        if self.command != "HEAD" and body:
            self.wfile.write(body)

    def _forward_headers(self):
        client_ip = self.client_address[0]
        fwd = self.headers.get("X-Forwarded-For", "")
        out = {}
        for key, value in self.headers.items():
            if key.lower() in HOP_HEADERS:
                continue
            out[key] = value
        out["X-Forwarded-For"] = (fwd + ", " + client_ip) if fwd else client_ip
        out["X-Forwarded-Proto"] = "http"
        out["X-Forwarded-Host"] = self.headers.get("Host", "")
        via = self.headers.get("Via", "")
        out["Via"] = (via + ", " if via else "") + "1.1 TaiProxy"
        out["Connection"] = "close"
        return out

    def _read_body(self):
        try:
            length = int(self.headers.get("Content-Length", "0") or "0")
        except ValueError:
            return None
        if length < 0 or length > MAX_BODY:
            return None
        if length == 0:
            return b""
        return self.rfile.read(length)

    def _try_backend(self, backend, path, body, headers):
        target = backend.prefix + path
        conn = HTTPConnection(backend.host, backend.port,
                              timeout=CONNECT_TIMEOUT)
        try:
            out_headers = dict(headers)
            out_headers["Host"] = backend.host + (
                "" if backend.port == 80 else ":%d" % backend.port)
            conn.request(self.command, target, body=body,
                         headers=out_headers)
            conn.sock.settimeout(READ_TIMEOUT)
            resp = conn.getresponse()
            self.send_response(resp.status, resp.reason)
            for key, value in resp.getheaders():
                if key.lower() in HOP_HEADERS:
                    continue
                self.send_header(key, value)
            self.send_header("Via", "1.1 TaiProxy")
            self.end_headers()
            if self.command != "HEAD":
                while True:
                    chunk = resp.read(65536)
                    if not chunk:
                        break
                    self.wfile.write(chunk)
            backend.mark(True)
            return True
        finally:
            try:
                conn.close()
            except Exception:
                pass

    def _proxy(self):
        pool = self.server.pool  # type: BackendPool
        if self.path == "/__proxy__/status" and self.command == "GET":
            self._send(200, json.dumps(pool.status()), "application/json")
            return
        body = self._read_body()
        if body is None:
            self._send(413, "body too large or bad length (max 10 MiB)")
            return
        headers = self._forward_headers()
        errors = []
        for backend in pool.pick_order():
            try:
                if self._try_backend(backend, self.path, body, headers):
                    return
            except (HTTPException, OSError) as exc:
                backend.mark(False, "%s: %s" % (type(exc).__name__, exc))
                errors.append("%s: %s" % (backend.name, exc))
                continue
            except BrokenPipeError:
                return
        self._send(502, "no healthy backend (%s)" % "; ".join(errors))

    def _dispatch(self):
        try:
            self._proxy()
        except BrokenPipeError:
            pass
        except Exception as exc:
            try:
                self._send(500, "proxy error: %s" % exc)
            except (BrokenPipeError, OSError):
                pass

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


class ProxyServer(socketserver.ThreadingMixIn, HTTPServer):
    daemon_threads = True
    allow_reuse_address = True


def main(argv=None):
    parser = argparse.ArgumentParser(
        description="Round-robin reverse proxy met healthchecks.")
    parser.add_argument("--backends", required=True,
                        help="comma-separated http://host:port[/prefix]")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8080)
    parser.add_argument("--health-path", default="/",
                        help="pad voor healthchecks (default /)")
    parser.add_argument("--health-interval", type=int, default=10,
                        help="seconden tussen healthchecks (default 10)")
    args = parser.parse_args(argv)
    urls = [u for u in args.backends.split(",") if u.strip()]
    if not urls or len(urls) > 32:
        parser.error("need 1..32 backends")
    try:
        backends = [Backend(u) for u in urls]
    except ValueError as exc:
        parser.error(str(exc))
    pool = BackendPool(backends)
    server = ProxyServer((args.host, args.port), ProxyHandler)
    server.pool = pool
    checker = HealthChecker(pool, path=args.health_path,
                            interval=args.health_interval)
    for backend in backends:
        checker.check_once(backend)
    checker.start()
    print("proxy: http://%s:%d/ -> %s" % (args.host, args.port,
                                          ", ".join(urls)))
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        checker.stop()


if __name__ == "__main__":
    main()