litellm/tests/integration/_support/oauth_server.py

"""OAuth 2.1 authorization-server double: metadata, DCR, PKCE authorization code, refresh, client credentials,
token exchange and revocation, every request recorded."""

import base64
import hashlib
import json
import secrets
import threading
import uuid
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Final
from urllib.parse import parse_qs, urlencode, urlsplit

from integration._support.wire import Reply, Request, Wire, wire_server

TOKEN_EXCHANGE: Final = "urn:ietf:params:oauth:grant-type:token-exchange"


@dataclass(slots=True)
class AuthorizationServer:
    wire: Wire
    clients: dict[str, str] = field(default_factory=dict)
    codes: dict[str, dict[str, str]] = field(default_factory=dict)
    access_tokens: dict[str, dict[str, str]] = field(default_factory=dict)
    refresh_tokens: dict[str, dict[str, str]] = field(default_factory=dict)
    revoked: set[str] = field(default_factory=set)
    lock: threading.Lock = field(default_factory=threading.Lock)
    mint: Callable[[str], str] | None = None

    @property
    def issuer(self) -> str:
        return self.wire.url

    def drain(self) -> tuple[Request, ...]:
        return self.wire.drain()

    def token_requests(self) -> tuple[dict[str, str], ...]:
        return tuple(
            {name: values[0] for name, values in parse_qs(item.body.decode()).items()}
            for item in self.drain()
            if item.target.startswith("/token")
        )

    def is_live(self, token: str) -> bool:
        with self.lock:
            return token in self.access_tokens and token not in self.revoked

    def issue(self, grant: str, client_id: str, subject: str, scope: str) -> dict[str, object]:
        access: Final = self.mint(grant) if self.mint is not None else f"at-{grant}-{secrets.token_urlsafe(8)}"
        refresh: Final = f"rt-{secrets.token_urlsafe(8)}"
        with self.lock:
            self.access_tokens[access] = {"client_id": client_id, "subject": subject, "scope": scope, "grant": grant}
            self.refresh_tokens[refresh] = {"client_id": client_id, "subject": subject, "scope": scope}
        return {
            "access_token": access,
            "token_type": "Bearer",
            "expires_in": 3600,
            "refresh_token": refresh,
            "scope": scope,
        }


def _pkce_matches(challenge: str, verifier: str) -> bool:
    digest: Final = hashlib.sha256(verifier.encode()).digest()
    return base64.urlsafe_b64encode(digest).rstrip(b"=").decode() == challenge


def _json(status: int, body: dict[str, object]) -> Reply:
    return Reply(status=status, body=json.dumps(body).encode())


def _client_credentials(request: Request, form: dict[str, str]) -> tuple[str, str | None]:
    header: Final = request.headers.get("authorization", "")
    if header.lower().startswith("basic "):
        decoded: Final = base64.b64decode(header.split(" ", 1)[1]).decode()
        client_id, _, secret = decoded.partition(":")
        return client_id, secret
    return form.get("client_id", ""), form.get("client_secret")


@contextmanager
def oauth_server(
    *, scopes: tuple[str, ...] = ("tools.read", "tools.call"), mint: Callable[[str], str] | None = None
) -> Iterator[AuthorizationServer]:
    """``mint(grant)``, when given, chooses each issued access token instead of a random one."""
    holder: list[AuthorizationServer] = []

    def respond(request: Request) -> Reply:
        server: Final = holder[0]
        path: Final = urlsplit(request.target).path
        query: Final = {name: values[0] for name, values in parse_qs(urlsplit(request.target).query).items()}
        form: Final = {name: values[0] for name, values in parse_qs(request.body.decode()).items()}
        if path.startswith("/.well-known/oauth-authorization-server") or path == "/.well-known/openid-configuration":
            return _json(
                200,
                {
                    "issuer": server.issuer,
                    "authorization_endpoint": server.issuer + "/authorize",
                    "token_endpoint": server.issuer + "/token",
                    "registration_endpoint": server.issuer + "/register",
                    "revocation_endpoint": server.issuer + "/revoke",
                    "introspection_endpoint": server.issuer + "/introspect",
                    "scopes_supported": list(scopes),
                    "response_types_supported": ["code"],
                    "grant_types_supported": [
                        "authorization_code",
                        "refresh_token",
                        "client_credentials",
                        TOKEN_EXCHANGE,
                    ],
                    "code_challenge_methods_supported": ["S256"],
                    "token_endpoint_auth_methods_supported": ["client_secret_post", "client_secret_basic", "none"],
                },
            )
        if path == "/register" and request.method == "POST":
            metadata: Final = json.loads(request.body or b"{}")
            client_id: Final = f"dcr-{uuid.uuid4().hex[:12]}"
            secret: Final = f"secret-{secrets.token_urlsafe(8)}"
            with server.lock:
                server.clients[client_id] = secret
            return _json(
                201,
                {
                    "client_id": client_id,
                    "client_secret": secret,
                    "client_id_issued_at": 0,
                    "redirect_uris": metadata.get("redirect_uris", []),
                    "grant_types": metadata.get("grant_types", ["authorization_code"]),
                    "token_endpoint_auth_method": metadata.get("token_endpoint_auth_method", "client_secret_post"),
                },
            )
        if path == "/authorize" and request.method == "GET":
            missing: Final = tuple(
                name for name in ("client_id", "redirect_uri", "code_challenge", "state") if name not in query
            )
            if missing or query.get("code_challenge_method", "S256") != "S256" or query.get("response_type") != "code":
                return _json(400, {"error": "invalid_request", "missing": list(missing), "received": query})
            code: Final = f"code-{secrets.token_urlsafe(8)}"
            with server.lock:
                server.codes[code] = {
                    "client_id": query["client_id"],
                    "redirect_uri": query["redirect_uri"],
                    "code_challenge": query["code_challenge"],
                    "scope": query.get("scope", " ".join(scopes)),
                }
            location: Final = (
                query["redirect_uri"]
                + ("&" if "?" in query["redirect_uri"] else "?")
                + urlencode({"code": code, "state": query["state"]})
            )
            return Reply(status=302, body=b"", headers={"location": location})
        if path == "/token" and request.method == "POST":
            grant: Final = form.get("grant_type", "")
            client_id, client_secret = _client_credentials(request, form)
            if grant == "authorization_code":
                with server.lock:
                    issued: Final = server.codes.pop(form.get("code", ""), None)
                if issued is None:
                    return _json(400, {"error": "invalid_grant", "error_description": "unknown or reused code"})
                if issued["client_id"] != client_id:
                    return _json(400, {"error": "invalid_client", "error_description": "code issued to another client"})
                if not _pkce_matches(issued["code_challenge"], form.get("code_verifier", "")):
                    return _json(400, {"error": "invalid_grant", "error_description": "pkce verifier mismatch"})
                return _json(200, server.issue("authorization_code", client_id, "integration-user", issued["scope"]))
            if grant == "refresh_token":
                with server.lock:
                    known: Final = server.refresh_tokens.pop(form.get("refresh_token", ""), None)
                if known is None:
                    return _json(400, {"error": "invalid_grant", "error_description": "unknown refresh token"})
                return _json(200, server.issue("refresh_token", known["client_id"], known["subject"], known["scope"]))
            if grant == "client_credentials":
                with server.lock:
                    expected: Final = server.clients.get(client_id)
                if not client_id or (expected is not None and expected != client_secret) or not client_secret:
                    return _json(401, {"error": "invalid_client"})
                return _json(200, server.issue("client_credentials", client_id, client_id, form.get("scope", "")))
            if grant == TOKEN_EXCHANGE:
                subject: Final = form.get("subject_token", "")
                if not subject:
                    return _json(400, {"error": "invalid_request", "error_description": "subject_token required"})
                if not client_id:
                    return _json(401, {"error": "invalid_client"})
                token: Final = server.issue("token_exchange", client_id, f"exchanged:{subject}", form.get("scope", ""))
                return _json(200, {**token, "issued_token_type": "urn:ietf:params:oauth:token-type:access_token"})
            return _json(400, {"error": "unsupported_grant_type", "grant_type": grant})
        if path == "/revoke" and request.method == "POST":
            with server.lock:
                server.revoked.add(form.get("token", ""))
            return Reply(status=200, body=b"{}")
        if path == "/introspect" and request.method == "POST":
            token: Final = form.get("token", "")
            with server.lock:
                info: Final = server.access_tokens.get(token)
                active: Final = info is not None and token not in server.revoked
            return _json(200, {"active": active, **(info or {})})
        return _json(404, {"error": "not_found", "path": path, "method": request.method})

    with wire_server(respond) as wire:
        holder.append(AuthorizationServer(wire, mint=mint))
        yield holder[0]