Kustt/mbappe_hermes_hf2

1#!/usr/bin/env python32"""HTTP reverse proxy with Host header rewrite + Basic Auth for HF Spaces.3Listens on 0.0.0.0:7860 and forwards to 127.0.0.1:7861.4Uses raw sockets for HTTP but transparently forwards WebSocket upgrades.5"""6 7import socket8import threading9import sys10import os11import base6412import secrets13 14LISTEN_HOST = os.environ.get("PROXY_LISTEN_HOST", "0.0.0.0")15LISTEN_PORT = int(os.environ.get("PROXY_LISTEN_PORT", "7860"))16BACKEND_HOST = os.environ.get("PROXY_BACKEND_HOST", "127.0.0.1")17BACKEND_PORT = int(os.environ.get("PROXY_BACKEND_PORT", "7861"))18 19# Basic Auth gate — read from HF Space Secrets (env) so secrets never land20# in /opt/data/.env (which is synced to the public dataset).21BASIC_AUTH_USER = os.environ.get("PROXY_BASIC_AUTH_USER", "")22BASIC_AUTH_PASSWORD = os.environ.get("PROXY_BASIC_AUTH_PASSWORD", "")23BASIC_AUTH_REALM = os.environ.get("PROXY_BASIC_AUTH_REALM", "luke-starwar")24 25BASIC_AUTH_ENABLED = bool(BASIC_AUTH_USER and BASIC_AUTH_PASSWORD)26if BASIC_AUTH_ENABLED:27    # Never log the actual creds. Log only whether auth is on.28    print(f"[proxy] Basic Auth enabled (user='{BASIC_AUTH_USER}', realm='{BASIC_AUTH_REALM}')", flush=True)29else:30    print("[proxy] WARNING: Basic Auth DISABLED — set PROXY_BASIC_AUTH_USER + PROXY_BASIC_AUTH_PASSWORD", flush=True)31 32 33def _check_basic_auth(data: bytes) -> bool:34    """Return True if request passes the Basic Auth gate (or auth disabled).35 36    Parses the Authorization header from the already-buffered request prefix.37    Uses hmac.compare_digest for constant-time comparison so a timing side38    channel can't leak the password byte-by-byte.39    """40    if not BASIC_AUTH_ENABLED:41        return True42    try:43        # Find end-of-headers marker to get the complete headers section44        headers_end = data.find(b"\r\n\r\n")45        if headers_end == -1:46            # No complete headers in buffer — can't auth, reject47            return False48        headers_section = data[:headers_end]49        text = headers_section.decode("iso-8859-1", errors="replace")50        # Walk header lines (skip the request line, text.split("\r\n")[0])51        for line in text.split("\r\n")[1:]:52            if line.lower().startswith("authorization:"):53                value = line.split(":", 1)[1].strip()54                if not value.lower().startswith("basic "):55                    return False56                encoded = value[6:].strip()57                try:58                    decoded = base64.b64decode(encoded).decode("utf-8", errors="replace")59                except Exception:60                    return False61                if ":" not in decoded:62                    return False63                user, _, password = decoded.partition(":")64                return (65                    secrets.compare_digest(user, BASIC_AUTH_USER)66                    and secrets.compare_digest(password, BASIC_AUTH_PASSWORD)67                )68    except Exception:69        return False70    return False71 72 73def _send_401(client_sock, data: bytes) -> None:74    """Emit a minimal 401 + WWW-Authenticate response. Use the same content-75    length framing pattern the dashboard expects so HF Spaces edge caching76    doesn't intercept."""77    body = b"Unauthorized\n"78    reason = b"Unauthorized"79    # Find request line to echo the HTTP version (1.0 vs 1.1)80    try:81        first = data.split(b"\r\n", 1)[0]82        # "GET /foo HTTP/1.1" -> "HTTP/1.1"83        parts = first.split(b" ")84        version = parts[2] if len(parts) >= 3 else b"HTTP/1.1"85    except Exception:86        version = b"HTTP/1.1"87    resp = (88        version + b" 401 " + reason + b"\r\n"89        b"WWW-Authenticate: Basic realm=\"" + BASIC_AUTH_REALM.encode("utf-8") + b"\"\r\n"90        b"Content-Type: text/plain; charset=utf-8\r\n"91        b"Content-Length: " + str(len(body)).encode("ascii") + b"\r\n"92        b"Connection: close\r\n"93        b"\r\n" + body94    )95    try:96        client_sock.sendall(resp)97    except OSError:98        pass99    finally:100        try:101            client_sock.shutdown(socket.SHUT_WR)102        except OSError:103            pass104        try:105            client_sock.close()106        except OSError:107            pass108 109 110def relay(src, dst):111    """Bidirectional byte copy between two sockets."""112    try:113        while True:114            data = src.recv(8192)115            if not data:116                break117            dst.sendall(data)118    except (OSError, ConnectionResetError):119        pass120    finally:121        try:122            src.shutdown(socket.SHUT_RD)123        except OSError:124            pass125        try:126            dst.shutdown(socket.SHUT_WR)127        except OSError:128            pass129 130 131def handle_http(client_sock, initial_data):132    """Handle HTTP/HTTPS request: check auth, rewrite Host header, forward to backend."""133    # Basic Auth gate — checked against the buffered request prefix so we134    # can fail fast without consuming any more bytes from the client.135    if not _check_basic_auth(initial_data):136        _send_401(client_sock, initial_data)137        return138    backend = None139    try:140        # Find end of headers141        data = initial_data142        if b"\r\n\r\n" in data:143            headers_end = data.index(b"\r\n\r\n") + 4144        else:145            headers_end = 0146            client_sock.settimeout(2.0)147            try:148                while b"\r\n\r\n" not in data:149                    chunk = client_sock.recv(8192)150                    if not chunk:151                        break152                    data += chunk153                if b"\r\n\r\n" in data:154                    headers_end = data.index(b"\r\n\r\n") + 4155            except socket.timeout:156                pass157            finally:158                client_sock.settimeout(None)159 160        # Reconstruct request with rewritten Host + Origin headers161        request_text = data[:headers_end].decode("iso-8859-1", errors="replace")162        lines = request_text.split("\r\n")163        new_lines = []164        for line in lines:165            if line.lower().startswith("host:"):166                new_lines.append(f"Host: {BACKEND_HOST}:{BACKEND_PORT}")167            elif line.lower().startswith("origin:"):168                # Rewrite Origin to loopback so PTY WebSocket origin check passes169                new_lines.append(f"Origin: http://{BACKEND_HOST}:{BACKEND_PORT}")170            else:171                new_lines.append(line)172        new_request = "\r\n".join(new_lines).encode("iso-8859-1") + data[headers_end:]173 174        backend = socket.socket(socket.AF_INET, socket.SOCK_STREAM)175        backend.connect((BACKEND_HOST, BACKEND_PORT))176        backend.sendall(new_request)177 178        # Now relay in both directions179        t1 = threading.Thread(target=relay, args=(client_sock, backend), daemon=True)180        t2 = threading.Thread(target=relay, args=(backend, client_sock), daemon=True)181        t1.start()182        t2.start()183        t1.join()184        t2.join()185    except Exception:186        pass187    finally:188        try:189            client_sock.close()190        except OSError:191            pass192        if backend is not None:193            try:194                backend.close()195            except OSError:196                pass197 198 199def handle_client(client_sock):200    """Peek at initial bytes to handle HTTP."""201    try:202        client_sock.settimeout(2.0)203        initial = client_sock.recv(8192)204        client_sock.settimeout(None)205        if not initial:206            return207        if (initial.startswith(b"GET ") or initial.startswith(b"POST ") or208            initial.startswith(b"PUT ") or initial.startswith(b"DELETE ") or209            initial.startswith(b"HEAD ") or initial.startswith(b"OPTIONS ") or210            initial.startswith(b"PATCH ")):211            handle_http(client_sock, initial)212        else:213            # Plain TCP passthrough (e.g., HTTPS after CONNECT)214            backend = socket.socket(socket.AF_INET, socket.SOCK_STREAM)215            backend.connect((BACKEND_HOST, BACKEND_PORT))216            backend.sendall(initial)217            t1 = threading.Thread(target=relay, args=(client_sock, backend), daemon=True)218            t2 = threading.Thread(target=relay, args=(backend, client_sock), daemon=True)219            t1.start()220            t2.start()221            t1.join()222            t2.join()223            try:224                backend.close()225            except OSError:226                pass227    except (OSError, socket.timeout):228        pass229    finally:230        try:231            client_sock.close()232        except OSError:233            pass234 235 236def main():237    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)238    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)239    server.bind((LISTEN_HOST, LISTEN_PORT))240    server.listen(128)241    print(f"[proxy] Listening on {LISTEN_HOST}:{LISTEN_PORT} → {BACKEND_HOST}:{BACKEND_PORT}", flush=True)242 243    while True:244        client_sock, addr = server.accept()245        threading.Thread(target=handle_client, args=(client_sock,), daemon=True).start()246 247 248if __name__ == "__main__":249    try:250        main()251    except KeyboardInterrupt:252        sys.exit(0)253