From 0df36bb139375481ab831849242f58875c8a179f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 28 Feb 2026 02:31:21 +0000 Subject: [PATCH] Add integration demo script for JWT auth fixes (OIDC discovery, array roles, dot-notation hints) Co-authored-by: Ishaan Jaff --- demo_jwt_fixes.py | 329 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 329 insertions(+) create mode 100644 demo_jwt_fixes.py diff --git a/demo_jwt_fixes.py b/demo_jwt_fixes.py new file mode 100644 index 00000000000..a2edaa8663a --- /dev/null +++ b/demo_jwt_fixes.py @@ -0,0 +1,329 @@ +#!/usr/bin/env python3 +""" +Integration demo for JWT auth fixes on branch fix/jwt-auth-oidc-array-roles. + +Spins up: + - Mock JWKS + OIDC discovery server on port 19900 + - Fake LLM backend on port 19901 + - LiteLLM proxy on port 19902 (premium_user patched to bypass license) + +Tests three scenarios: + Fix 1 OIDC discovery URL resolution (valid JWT → 200, tampered → 401) + Fix 2 roles claim as JSON array (["team-beta","team-gamma"] → 200, [] → 401) + Fix 3 Helpful error for dot-notation (roles.0 → 401 with hint) +""" + +import asyncio, base64, json, os, signal, subprocess, sys, textwrap, time, threading +from http.server import HTTPServer, BaseHTTPRequestHandler + +from cryptography.hazmat.primitives.asymmetric import rsa +from cryptography.hazmat.primitives import serialization +import jwt as pyjwt + +# ── 1. RSA key pair ────────────────────────────────────────────────────────── + +RSA_KEY = rsa.generate_private_key(public_exponent=65537, key_size=2048) +RSA_PUB = RSA_KEY.public_key() +PUB_NUMBERS = RSA_PUB.public_numbers() + + +def _b64url(n: int, length: int) -> str: + return base64.urlsafe_b64encode(n.to_bytes(length, "big")).rstrip(b"=").decode() + + +KID = "demo-kid-1" +JWKS = { + "keys": [{ + "kty": "RSA", "kid": KID, "use": "sig", "alg": "RS256", + "n": _b64url(PUB_NUMBERS.n, 256), + "e": _b64url(PUB_NUMBERS.e, 3), + }] +} + +PEM_PRIVATE = RSA_KEY.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), +) + +ISSUER = "https://demo-idp.example.com" +AUDIENCE = "litellm-proxy-demo" + + +def sign_token(claims: dict) -> str: + payload = {"iss": ISSUER, "aud": AUDIENCE, "exp": int(time.time()) + 3600, **claims} + return pyjwt.encode(payload, PEM_PRIVATE, algorithm="RS256", headers={"kid": KID}) + + +# ── 2. Mock JWKS + OIDC discovery server (port 19900) ──────────────────────── + +class JWKSHandler(BaseHTTPRequestHandler): + def do_GET(self): + if self.path == "/.well-known/openid-configuration": + body = json.dumps({ + "issuer": ISSUER, + "jwks_uri": "http://127.0.0.1:19900/jwks", + "authorization_endpoint": "https://example.com/authorize", + "token_endpoint": "https://example.com/token", + }).encode() + elif self.path == "/jwks": + body = json.dumps(JWKS).encode() + else: + self.send_response(404); self.end_headers(); return + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(body) + + def log_message(self, *_): + pass + + +# ── 3. Fake LLM backend (port 19901) ───────────────────────────────────────── + +class FakeLLMHandler(BaseHTTPRequestHandler): + def do_POST(self): + body = json.dumps({ + "id": "chatcmpl-fake", "object": "chat.completion", "created": 1700000000, + "model": "fake-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello from fake LLM!"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12}, + }).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + if "/models" in self.path: + body = json.dumps({"data": [{"id": "fake-model", "object": "model"}], "object": "list"}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(body) + else: + self.send_response(404); self.end_headers() + + def log_message(self, *_): + pass + + +# ── 4. Proxy startup that patches premium_user ─────────────────────────────── + +PROXY_PORT = 19902 + +CONFIG_FIX12 = textwrap.dedent("""\ + model_list: + - model_name: fake-model + litellm_params: + model: openai/fake-model + api_key: fake-key + api_base: http://127.0.0.1:19901/ + general_settings: + enable_jwt_auth: true + litellm_jwtauth: + team_id_jwt_field: "roles" + team_id_default: "" + litellm_settings: + drop_params: true +""") + +CONFIG_FIX3 = textwrap.dedent("""\ + model_list: + - model_name: fake-model + litellm_params: + model: openai/fake-model + api_key: fake-key + api_base: http://127.0.0.1:19901/ + general_settings: + enable_jwt_auth: true + litellm_jwtauth: + team_id_jwt_field: "roles.0" + team_id_default: "" + litellm_settings: + drop_params: true +""") + + +def _write_config(content: str, path: str): + with open(path, "w") as f: + f.write(content) + + +PROXY_WRAPPER = textwrap.dedent("""\ + import sys, os, unittest.mock + os.environ["JWT_PUBLIC_KEY_URL"] = "http://127.0.0.1:19900/.well-known/openid-configuration" + os.environ["JWT_AUDIENCE"] = "{audience}" + os.environ["DISABLE_SCHEMA_UPDATE"] = "true" + # Patch premium_user BEFORE proxy imports the variable + import litellm.proxy.proxy_server as _ps + _ps.premium_user = True + from litellm.proxy.proxy_server import app, initialize + import asyncio, uvicorn + asyncio.get_event_loop().run_until_complete(initialize(config="{config_path}")) + uvicorn.run(app, host="127.0.0.1", port={port}, log_level="warning") +""") + + +def start_proxy(config_path: str) -> subprocess.Popen: + wrapper_code = PROXY_WRAPPER.format( + audience=AUDIENCE, + config_path=config_path, + port=PROXY_PORT, + ) + wrapper_path = "/tmp/_proxy_wrapper.py" + with open(wrapper_path, "w") as f: + f.write(wrapper_code) + proc = subprocess.Popen( + [sys.executable, wrapper_path], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + ) + return proc + + +def wait_for_proxy(timeout: int = 90) -> bool: + import urllib.request + deadline = time.time() + timeout + while time.time() < deadline: + try: + r = urllib.request.urlopen(f"http://127.0.0.1:{PROXY_PORT}/health/readiness", timeout=2) + if r.status == 200: + return True + except Exception: + pass + time.sleep(1) + return False + + +def stop_proxy(proc: subprocess.Popen): + proc.send_signal(signal.SIGTERM) + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + + +def proxy_request(token: str, model: str = "fake-model") -> tuple: + import urllib.request, urllib.error + data = json.dumps({"model": model, "messages": [{"role": "user", "content": "Hi"}]}).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{PROXY_PORT}/v1/chat/completions", + data=data, + headers={"Content-Type": "application/json", "Authorization": f"Bearer {token}"}, + ) + try: + r = urllib.request.urlopen(req, timeout=15) + return r.status, r.read().decode() + except urllib.error.HTTPError as e: + return e.code, e.read().decode() + + +# ── 5. Pretty output ───────────────────────────────────────────────────────── + +GREEN = "\033[92m" +RED = "\033[91m" +BOLD = "\033[1m" +RESET = "\033[0m" +results = [] + + +def header(text: str): + print(f"\n{BOLD}{'='*72}{RESET}") + print(f"{BOLD} {text}{RESET}") + print(f"{BOLD}{'='*72}{RESET}") + + +def check(label: str, condition: bool, detail: str = ""): + mark = f"{GREEN}PASS{RESET}" if condition else f"{RED}FAIL{RESET}" + print(f" [{mark}] {label}") + if detail: + for line in detail.strip().split("\n"): + print(f" {line}") + results.append((label, condition)) + + +# ── 6. Main ────────────────────────────────────────────────────────────────── + +def main(): + _write_config(CONFIG_FIX12, "/tmp/config_fix12.yaml") + _write_config(CONFIG_FIX3, "/tmp/config_fix3.yaml") + + print(f"{BOLD}Starting mock JWKS + OIDC discovery server on :19900 ...{RESET}") + jwks_srv = HTTPServer(("127.0.0.1", 19900), JWKSHandler) + threading.Thread(target=jwks_srv.serve_forever, daemon=True).start() + + print(f"{BOLD}Starting fake LLM backend on :19901 ...{RESET}") + llm_srv = HTTPServer(("127.0.0.1", 19901), FakeLLMHandler) + threading.Thread(target=llm_srv.serve_forever, daemon=True).start() + + # ── Fix 1 + 2 ──────────────────────────────────────────────────────── + header("Fix 1 + Fix 2: OIDC Discovery & Array Roles") + proxy = start_proxy("/tmp/config_fix12.yaml") + try: + print(f" Waiting for proxy on :{PROXY_PORT} (may take ~30s for Prisma) ...") + if not wait_for_proxy(): + print(f"{RED}Proxy failed to start!{RESET}") + out = proxy.stdout.read().decode() if proxy.stdout else "" + print(out[-3000:]) + raise SystemExit(1) + print(f" {GREEN}Proxy is up.{RESET}\n") + + # Fix 1a + token_ok = sign_token({"sub": "user-1", "roles": ["team-beta", "team-gamma"]}) + status, body = proxy_request(token_ok) + check("Fix 1 — OIDC discovery: valid JWT → HTTP 200", status == 200, f"status={status}") + + # Fix 1b + tampered = token_ok[:-5] + "XXXXX" + status, body = proxy_request(tampered) + check("Fix 1 — OIDC discovery: tampered JWT → HTTP 401", status == 401, f"status={status}") + + # Fix 2a + token_arr = sign_token({"sub": "user-2", "roles": ["team-beta", "team-gamma"]}) + status, body = proxy_request(token_arr) + check("Fix 2 — roles array ['team-beta','team-gamma'] → HTTP 200", status == 200, f"status={status}") + + # Fix 2b + token_empty = sign_token({"sub": "user-3", "roles": []}) + status, body = proxy_request(token_empty) + check("Fix 2 — empty roles [] → HTTP 401", status == 401, f"status={status}") + + finally: + stop_proxy(proxy) + time.sleep(2) + + # ── Fix 3 ───────────────────────────────────────────────────────────── + header("Fix 3: Dot-notation hint (roles.0)") + proxy = start_proxy("/tmp/config_fix3.yaml") + try: + print(f" Waiting for proxy on :{PROXY_PORT} ...") + if not wait_for_proxy(): + print(f"{RED}Proxy failed to start!{RESET}") + out = proxy.stdout.read().decode() if proxy.stdout else "" + print(out[-3000:]) + raise SystemExit(1) + print(f" {GREEN}Proxy is up.{RESET}\n") + + token_dot = sign_token({"sub": "user-4", "roles": ["team-alpha"]}) + status, body = proxy_request(token_dot) + hint_text = "Use 'roles' instead" + has_hint = hint_text in body + check( + "Fix 3 — roles.0 config → HTTP 401 with helpful hint", + status == 401 and has_hint, + f"status={status}, hint_present={has_hint}\n" + f"body excerpt: {body[:400]}", + ) + + finally: + stop_proxy(proxy) + + # ── Summary ─────────────────────────────────────────────────────────── + all_pass = all(ok for _, ok in results) + header("ALL TESTS PASSED" if all_pass else "SOME TESTS FAILED") + print() + raise SystemExit(0 if all_pass else 1) + + +if __name__ == "__main__": + main()