mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
## What & why
Adds **`headroom-oauth2`** under `plugins/` — a generic, vendor-neutral
proxy extension that mints an OAuth2 **client-credentials** (RFC 6749
§4.4) bearer from a configured token endpoint and injects it as the
upstream `Authorization` on each proxied request, via the opt-in
`headroom.proxy_extension` seam. **No core changes.**
This lets headroom front any gateway that requires a *minted,
short-lived machine token* rather than a static API key. It complements
**#510** (env-var/static-key auth) rather than replacing it.
Implements **#778** (feature request). Opening the implementation
alongside the issue so there's something concrete to react to — **happy
to hold/rework pending a 👍 from a maintainer**, per CONTRIBUTING.
## Spec
Full spec in
[`plugins/headroom-oauth2/SPEC.md`](plugins/headroom-oauth2/SPEC.md)
(API surface, behavior/compat, user stories, failure modes, resilience
incl. multi-process, security, observability, rollback). Highlights:
- **Opt-in & no-op by default:** dormant until `--proxy-extension
oauth2`, and a no-op unless `HEADROOM_OAUTH2_TOKEN_URL` is set. No
change to defaults, body, routing, or compression.
- **Config is 100% env** (no new CLI flags):
token_url/client_id/secret/scopes/audience, RFC 8707 `resource`,
`post`|`basic` auth style, static upstream headers, timeout/skew.
- **Token caching + single-flight refresh**; `expires_in` clamped to a
positive TTL.
- **Fails closed** on misconfig; returns `502 upstream_auth_error` on
mint failure **without leaking the IdP error body**; `token_url` is
**https-enforced** (loopback exempt for tests).
- **Standard-library only** (token minted via `urllib` → system cert
store, so it works behind corporate SSL inspection). `litellm` is
touched only for static headers and is an optional extra, not a core
dep.
- **Effective for** OpenAI-compatible / passthrough litellm backends.
`bedrock`/`vertex`/`sagemaker` auth from env and ignore a forwarded
bearer → the extension **warns loudly** and is a no-op there.
## Tests
37 tests covering behavior **and** failure modes (`ruff check`/`format`
clean, **98% coverage**): post/basic mint, caching, single-flight (cold
+ on-refresh, exact mint counts under concurrency), https enforcement +
`localhost` rejection + `::1`, `expires_in`
clamp/float/missing/non-numeric, `extra_params` cannot clobber canonical
fields, bad-status/non-JSON/no-token/unreachable (asserting no
secret/body leak), ASGI middleware (inject, non-http passthrough, 502 +
`no-store`, missing `headers` key), and `install()`
(no-op/fail-closed/bad-timeout/env-auth-backend-warning/static-headers).
## Real behavior proof
- **Setup:** Linux aarch64, Python 3.13.5, `headroom-ai` 0.23.0, real
`headroom proxy` process.
- **Steps:** started `headroom proxy --backend litellm-openai
--proxy-extension oauth2` with `HEADROOM_OAUTH2_*` env pointed at a
local OAuth2 token endpoint; an upstream echo server captured what the
backend received; sent two `/v1/messages` requests through the proxy.
- **Observed (copied output):**
```
PROXY: headroom-oauth2: client-credentials auth installed
(token_url=…/token, style=post)
MINTS (across 2 requests): 1 # token cached + reused -> 1 mint for 2
requests
UPSTREAM RECEIVED: auth=Bearer MINTED-FROM-IDP-…
static=generic-static-header
SECRET LEAK CHECK (client_secret in proxy logs): 0
```
→ The minted bearer (not the placeholder backend key) and the configured
static header reached the upstream; the client's inbound credential was
replaced; the client secret never appeared in logs; caching worked.
- **What I did *not* test:** a live commercial IdP
(Entra/Okta/Auth0/etc.) and a live cloud gateway — the token endpoint
and upstream here are local stand-ins. Also not tested: multi-worker
(gunicorn) deployment, and Python 3.10/3.11 (developed on 3.13).
## Placement
Proposed as a standalone installable package under
`plugins/headroom-oauth2/` (registers via the entry-point seam; `pip
install -e plugins/headroom-oauth2`). Open to baking it into core or
publishing it separately — maintainer's call.
This commit is contained in:
Khalid Shaikh
GitHub
parent
5dfb446da1
commit
eb2e50feb2
9 changed files with 1066 additions and 0 deletions
16
plugins/headroom-oauth2/CHANGELOG.md
Normal file
16
plugins/headroom-oauth2/CHANGELOG.md
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
# Changelog
|
||||
|
||||
## 0.1.0
|
||||
|
||||
Initial release — generic OAuth2 client-credentials upstream-auth extension for the Headroom proxy.
|
||||
|
||||
- Mints an OAuth2 client-credentials (RFC 6749 §4.4) bearer from a configurable token endpoint and
|
||||
injects it as the upstream `Authorization` on each proxied request, via Headroom's opt-in
|
||||
`headroom.proxy_extension` seam (`--proxy-extension oauth2`). No core changes; vendor-neutral.
|
||||
- `post` and `basic` client-auth styles; scopes, `audience`, RFC 8707 `resource`, static upstream
|
||||
headers, configurable timeout/skew — all env-driven.
|
||||
- Token caching with single-flight refresh and pre-expiry skew; `expires_in` clamped to a positive
|
||||
TTL.
|
||||
- Fails closed on misconfiguration; returns `502 upstream_auth_error` on mint failure without
|
||||
leaking the IdP error body. `token_url` https-enforced (loopback exempt). Std-lib only (system
|
||||
cert store -> works behind corporate SSL inspection).
|
||||
3
plugins/headroom-oauth2/LICENSE
Normal file
3
plugins/headroom-oauth2/LICENSE
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
SPDX-License-Identifier: Apache-2.0
|
||||
Apache License 2.0 — full text: https://www.apache.org/licenses/LICENSE-2.0
|
||||
(Matches the upstream Headroom license; full text bundled at publish time.)
|
||||
43
plugins/headroom-oauth2/README.md
Normal file
43
plugins/headroom-oauth2/README.md
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
# headroom-oauth2
|
||||
|
||||
Generic **OAuth2 client-credentials** upstream-auth extension for the
|
||||
[Headroom](https://github.com/chopratejas/headroom) proxy.
|
||||
|
||||
When Headroom routes to an OpenAI-compatible backend that is protected by an
|
||||
OAuth2 client-credentials flow (enterprise AI gateways, Azure AD / Entra, Okta,
|
||||
Auth0, Keycloak, Cognito, …), this extension mints a bearer token from a
|
||||
configurable token endpoint, caches + refreshes it (single-flight), and injects
|
||||
`Authorization: Bearer <token>` on each upstream request. Optional static upstream
|
||||
headers are sent via litellm. **Fully vendor-neutral — no provider is hard-coded.**
|
||||
|
||||
It plugs into Headroom's public `headroom.proxy_extension` entry-point seam, so it
|
||||
is fully out-of-tree and opt-in.
|
||||
|
||||
## Install & enable
|
||||
```bash
|
||||
pip install headroom-oauth2
|
||||
headroom proxy --backend litellm-openai --proxy-extension oauth2
|
||||
```
|
||||
|
||||
## Configure (env; no-op unless HEADROOM_OAUTH2_TOKEN_URL is set)
|
||||
| Env | Meaning |
|
||||
|-----|---------|
|
||||
| `HEADROOM_OAUTH2_TOKEN_URL` | token endpoint (client_credentials grant) |
|
||||
| `HEADROOM_OAUTH2_CLIENT_ID` / `_CLIENT_SECRET` | credentials (secrets) |
|
||||
| `HEADROOM_OAUTH2_SCOPES` | space/comma-separated scopes |
|
||||
| `HEADROOM_OAUTH2_AUDIENCE` | optional audience |
|
||||
| `HEADROOM_OAUTH2_GRANT_TYPE` | default `client_credentials` |
|
||||
| `HEADROOM_OAUTH2_AUTH_STYLE` | `post` (form creds) or `basic` (HTTP Basic) |
|
||||
| `HEADROOM_OAUTH2_HEADERS` | static upstream headers, `K=V,K2=V2` |
|
||||
|
||||
Tokens are minted with the standard library (`urllib`, system cert store), which
|
||||
works behind corporate SSL-inspection where bundled-root TLS stacks fail.
|
||||
|
||||
|
||||
**Effective backends:** the injected bearer reaches the upstream only for OpenAI-compatible /
|
||||
passthrough litellm providers. `bedrock` / `vertex` / `sagemaker` authenticate from env and
|
||||
ignore it, so this extension is a no-op there (it logs a warning at startup).
|
||||
|
||||
**Transport:** `token_url` must be `https` (loopback `http` is allowed for tests; set
|
||||
`HEADROOM_OAUTH2_ALLOW_INSECURE=1` to override). Tokens are minted with the standard library
|
||||
(`urllib`, system cert store), so a corporate-injected CA is trusted without bundling roots.
|
||||
115
plugins/headroom-oauth2/SPEC.md
Normal file
115
plugins/headroom-oauth2/SPEC.md
Normal file
|
|
@ -0,0 +1,115 @@
|
|||
# Spec — `headroom-oauth2` (generic OAuth2 client-credentials upstream auth)
|
||||
|
||||
Feature request: [chopratejas/headroom#778](https://github.com/chopratejas/headroom/issues/778).
|
||||
Status: implementation ready; awaiting maintainer 👍 before merge.
|
||||
|
||||
## Summary
|
||||
|
||||
A vendor-neutral proxy extension (registers on Headroom's `headroom.proxy_extension` seam) that
|
||||
mints an OAuth2 **client-credentials** (RFC 6749 §4.4) bearer from a configured token endpoint and
|
||||
injects it as the upstream `Authorization` on every proxied request. Lets Headroom front any
|
||||
gateway that requires a minted machine token (not a static API key) — with **zero core changes**
|
||||
and **no vendor specifics** (the gateway is entirely config/env).
|
||||
|
||||
It complements `#510` (env-var auth), which assumes a long-lived static key; this covers the
|
||||
"mint-then-refresh a short-lived token" case.
|
||||
|
||||
## API surface (config / CLI / env)
|
||||
|
||||
Opt-in only, via Headroom's existing flags — **no new CLI flags**:
|
||||
|
||||
headroom proxy --backend litellm-openai --proxy-extension oauth2
|
||||
# or HEADROOM_PROXY_EXTENSIONS=oauth2
|
||||
|
||||
All configuration is env (12-factor; nothing baked in):
|
||||
|
||||
| Env var | Required | Meaning |
|
||||
|---|---|---|
|
||||
| `HEADROOM_OAUTH2_TOKEN_URL` | yes (else no-op) | OAuth2 token endpoint; must be `https` (loopback `http` allowed for tests) |
|
||||
| `HEADROOM_OAUTH2_CLIENT_ID` / `_CLIENT_SECRET` | yes | client credentials |
|
||||
| `HEADROOM_OAUTH2_SCOPES` | no | space/comma-separated scopes |
|
||||
| `HEADROOM_OAUTH2_AUDIENCE` | no | `audience` form param |
|
||||
| `HEADROOM_OAUTH2_RESOURCE` | no | RFC 8707 target `resource` form param |
|
||||
| `HEADROOM_OAUTH2_GRANT_TYPE` | no | default `client_credentials` |
|
||||
| `HEADROOM_OAUTH2_AUTH_STYLE` | no | `post` (form creds) or `basic` (HTTP Basic) |
|
||||
| `HEADROOM_OAUTH2_HEADERS` | no | static upstream headers, `K=V,K2=V2` (control chars rejected) |
|
||||
| `HEADROOM_OAUTH2_TIMEOUT` / `_SKEW` | no | token request timeout / pre-expiry refresh skew (s) |
|
||||
| `HEADROOM_OAUTH2_ALLOW_INSECURE` | no | `1` to allow a non-loopback `http` token_url (discouraged) |
|
||||
|
||||
Public Python API: `OAuth2ClientCredentials`, `OAuth2Middleware`, `OAuth2Error`, `install`,
|
||||
`provider_from_env`, `parse_headers`.
|
||||
|
||||
## Changes to existing behavior / defaults / compatibility
|
||||
|
||||
- **None unless explicitly enabled.** The entry point is dormant until `--proxy-extension oauth2`
|
||||
is passed, and even then a **no-op** unless `HEADROOM_OAUTH2_TOKEN_URL` is set.
|
||||
- When active, it **overwrites the request `Authorization` header** with the minted bearer before
|
||||
the backend runs. The client's own `Authorization`/`x-api-key` is intentionally replaced (the
|
||||
proxy authenticates to the gateway on the client's behalf). *Compatibility note:* because the
|
||||
request then carries a bearer, Headroom classifies it as OAuth-mode auth — same as supplying a
|
||||
bearer yourself; no new classification path.
|
||||
- No change to defaults, the request/response body, model routing, or compression.
|
||||
|
||||
## User stories (Given / When / Then)
|
||||
|
||||
- **Golden path** — *Given* a proxy started with `--proxy-extension oauth2` and valid
|
||||
`TOKEN_URL`/`CLIENT_ID`/`CLIENT_SECRET`, *When* a client sends `/v1/messages`, *Then* the
|
||||
extension mints (or reuses a cached) bearer and the upstream receives `Authorization: Bearer
|
||||
<minted>` plus any static headers; the client never sees the secret.
|
||||
- **Edge: token endpoint down** — *Given* an unreachable/erroring `TOKEN_URL`, *When* a request
|
||||
arrives, *Then* the proxy returns `502 upstream_auth_error` (no upstream call, no secret/body
|
||||
leak) and stays up; the next request retries.
|
||||
- **Edge: wrong backend** — *Given* `--backend bedrock` (env-auth), *When* the extension installs,
|
||||
*Then* it logs a loud warning that the injected bearer will have no effect and to use an
|
||||
OpenAI-compatible/passthrough backend.
|
||||
|
||||
## Failure modes & recovery
|
||||
|
||||
| Failure | Behavior |
|
||||
|---|---|
|
||||
| Missing/invalid config at startup | `install()` raises -> proxy **fails closed** (won't start mis-auth'd) |
|
||||
| Token endpoint unreachable / non-2xx / non-JSON / no `access_token` | `OAuth2Error` -> `502`, per-request, proxy stays up, retried next request |
|
||||
| `expires_in` = 0/negative/absent | clamped to a positive TTL (never stale, never per-request mint) |
|
||||
| Concurrent first requests | single-flight lock -> exactly one mint per refresh |
|
||||
| Malformed `HEADROOM_OAUTH2_HEADERS` (CR/LF) | offending pair dropped with a warning (no header injection) |
|
||||
|
||||
## Resilience (Docker / native / wrappers / providers / multi-process)
|
||||
|
||||
- **Native & Docker:** identical; pure env-driven, std-lib only. Token minted via `urllib` against
|
||||
the **system cert store**, so a corporate-injected CA is trusted with no bundled roots (works in
|
||||
SSL-inspection networks).
|
||||
- **Wrappers (`headroom wrap`, agent hooks):** the extension lives at the proxy layer, so anything
|
||||
routed through the proxy inherits it transparently.
|
||||
- **Providers:** effective for OpenAI-compatible / passthrough litellm backends (those that forward
|
||||
the request bearer upstream). `bedrock`/`vertex`/`sagemaker` authenticate from env and ignore the
|
||||
bearer -> the extension warns and is a no-op there.
|
||||
- **Multi-process (multiple workers):** the token cache is per-process; each worker mints/refreshes
|
||||
independently. Acceptable for client-credentials (idempotent, low rate); no shared state, no
|
||||
cross-process lock needed. Documented so operators can size token-endpoint rate limits.
|
||||
|
||||
## Security & privacy
|
||||
|
||||
- Secrets are env-only; **never logged** and **never returned** to the client.
|
||||
- The IdP error body is **drained, not surfaced** (may echo sensitive context).
|
||||
- `token_url` is **https-enforced** (loopback exception for tests; explicit opt-out env).
|
||||
- The minted bearer is sent only to the configured upstream; the client's inbound credential is
|
||||
replaced, not forwarded onward.
|
||||
|
||||
## Observability / logging / telemetry
|
||||
|
||||
- `INFO` on install (token_url + auth style, no secrets) and on each mint (`ttl`, scopes).
|
||||
- `WARNING` on mint failure, env-auth-backend no-op, and dropped malformed static headers.
|
||||
- No metrics/telemetry emitted; piggybacks on Headroom's existing request logging. (A future
|
||||
counter for mint/refresh/failure could be added if maintainers want it.)
|
||||
|
||||
## Rollback / migration
|
||||
|
||||
- **No migration** — additive and opt-in; existing deployments are unaffected.
|
||||
- **Instant rollback:** drop `--proxy-extension oauth2` (or unset `HEADROOM_PROXY_EXTENSIONS`), or
|
||||
uninstall the package. No state to clean up, no config format changes.
|
||||
|
||||
## Dependencies
|
||||
|
||||
- **Runtime:** standard library only (no new core dependency). `litellm` is touched **only** if
|
||||
`HEADROOM_OAUTH2_HEADERS` is set, and it is already a Headroom backend dependency — declared here
|
||||
as the optional `[litellm]` extra, not a hard requirement.
|
||||
48
plugins/headroom-oauth2/pyproject.toml
Normal file
48
plugins/headroom-oauth2/pyproject.toml
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
[build-system]
|
||||
requires = ["setuptools>=61"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "headroom-oauth2"
|
||||
version = "0.1.0"
|
||||
description = "Generic OAuth2 client-credentials upstream-auth extension for the Headroom proxy"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
license = { text = "Apache-2.0" }
|
||||
authors = [{ name = "Khalid Shaikh", email = "43288811+mkhalid-s@users.noreply.github.com" }]
|
||||
keywords = ["headroom", "oauth2", "client-credentials", "proxy", "llm-gateway"]
|
||||
dependencies = []
|
||||
classifiers = [
|
||||
"Development Status :: 4 - Beta",
|
||||
"Intended Audience :: Developers",
|
||||
"License :: OSI Approved :: Apache Software License",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Topic :: Internet :: Proxy Servers",
|
||||
"Topic :: Security",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
# Only needed if you set HEADROOM_OAUTH2_HEADERS (static upstream headers via litellm.headers).
|
||||
# The token-minting + injection path is standard-library only and needs none of this.
|
||||
litellm = ["litellm>=1.40"]
|
||||
dev = ["pytest>=7", "ruff"]
|
||||
test = ["pytest>=7"]
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://github.com/chopratejas/headroom"
|
||||
Repository = "https://github.com/chopratejas/headroom"
|
||||
Issues = "https://github.com/chopratejas/headroom/issues"
|
||||
|
||||
# Registers with Headroom's opt-in proxy-extension seam (headroom/proxy/extensions.py).
|
||||
# Enable at runtime with `--proxy-extension oauth2` or HEADROOM_PROXY_EXTENSIONS=oauth2.
|
||||
[project.entry-points."headroom.proxy_extension"]
|
||||
oauth2 = "headroom_oauth2:install"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 100
|
||||
136
plugins/headroom-oauth2/src/headroom_oauth2/__init__.py
Normal file
136
plugins/headroom-oauth2/src/headroom_oauth2/__init__.py
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
"""Generic OAuth2 client-credentials upstream-auth extension for the Headroom proxy.
|
||||
|
||||
Enable: `--proxy-extension oauth2` (or HEADROOM_PROXY_EXTENSIONS=oauth2).
|
||||
No-op unless HEADROOM_OAUTH2_TOKEN_URL is set. See README for env config.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from .middleware import OAuth2Middleware
|
||||
from .provider import OAuth2ClientCredentials, OAuth2Error
|
||||
|
||||
__all__ = ["install", "OAuth2ClientCredentials", "OAuth2Error", "OAuth2Middleware", "parse_headers"]
|
||||
__version__ = "0.1.0"
|
||||
log = logging.getLogger("headroom_oauth2")
|
||||
|
||||
|
||||
def _split(s):
|
||||
return [x.strip() for x in (s or "").replace(",", " ").split() if x.strip()]
|
||||
|
||||
|
||||
def _ctrl(s):
|
||||
return any(ord(c) < 32 or ord(c) == 127 for c in s)
|
||||
|
||||
|
||||
def parse_headers(s: str | None) -> dict[str, str]:
|
||||
"""Parse ``K=V,K2=V2`` into a dict. Drops pairs whose key/value contain control
|
||||
characters, or whose key contains a space or colon -- prevents HTTP header injection
|
||||
from a malformed env value."""
|
||||
out: dict[str, str] = {}
|
||||
for pair in (s or "").split(","):
|
||||
if "=" not in pair:
|
||||
continue
|
||||
k, v = (x.strip() for x in pair.split("=", 1))
|
||||
if not k:
|
||||
continue
|
||||
if _ctrl(k) or _ctrl(v) or " " in k or ":" in k:
|
||||
log.warning("headroom-oauth2: dropping malformed static header: %r", k)
|
||||
continue
|
||||
out[k] = v
|
||||
return out
|
||||
|
||||
|
||||
def _int(env, key):
|
||||
raw = env.get(key)
|
||||
if raw is None or not str(raw).strip():
|
||||
return None
|
||||
try:
|
||||
return int(raw)
|
||||
except ValueError:
|
||||
raise ValueError(f"{key}={raw!r} is not an integer") from None
|
||||
|
||||
|
||||
def provider_from_env(env: dict | None = None) -> OAuth2ClientCredentials | None:
|
||||
"""Build a provider from ``HEADROOM_OAUTH2_*`` env vars, or None if TOKEN_URL is unset.
|
||||
|
||||
Raises ValueError on malformed config so callers can fail closed.
|
||||
"""
|
||||
env = os.environ if env is None else env
|
||||
token_url = env.get("HEADROOM_OAUTH2_TOKEN_URL")
|
||||
if not token_url:
|
||||
return None
|
||||
allow_insecure = env.get("HEADROOM_OAUTH2_ALLOW_INSECURE", "").strip().lower() in (
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
)
|
||||
if allow_insecure:
|
||||
log.warning("headroom-oauth2: ALLOW_INSECURE set -- token endpoint TLS check disabled")
|
||||
resource = env.get("HEADROOM_OAUTH2_RESOURCE") # RFC 8707 target service
|
||||
timeout = _int(env, "HEADROOM_OAUTH2_TIMEOUT")
|
||||
skew = _int(env, "HEADROOM_OAUTH2_SKEW")
|
||||
kwargs = {}
|
||||
if timeout is not None:
|
||||
kwargs["timeout_seconds"] = float(timeout)
|
||||
if skew is not None:
|
||||
kwargs["skew_seconds"] = skew
|
||||
return OAuth2ClientCredentials(
|
||||
token_url=token_url,
|
||||
client_id=env.get("HEADROOM_OAUTH2_CLIENT_ID", ""),
|
||||
client_secret=env.get("HEADROOM_OAUTH2_CLIENT_SECRET", ""),
|
||||
scopes=_split(env.get("HEADROOM_OAUTH2_SCOPES")),
|
||||
audience=env.get("HEADROOM_OAUTH2_AUDIENCE") or None,
|
||||
grant_type=env.get("HEADROOM_OAUTH2_GRANT_TYPE", "client_credentials"),
|
||||
auth_style=env.get("HEADROOM_OAUTH2_AUTH_STYLE", "post"),
|
||||
extra_params={"resource": resource} if resource else None,
|
||||
allow_insecure=allow_insecure,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def install(app: Any, config: Any) -> None:
|
||||
"""Headroom proxy-extension entry point: install(app, config) -> None."""
|
||||
try:
|
||||
provider = provider_from_env()
|
||||
except ValueError as e:
|
||||
raise RuntimeError(f"headroom-oauth2 misconfigured: {e}") from None # fail-closed
|
||||
if provider is None:
|
||||
log.info("headroom-oauth2 loaded but HEADROOM_OAUTH2_TOKEN_URL unset; no-op")
|
||||
return
|
||||
static = parse_headers(os.environ.get("HEADROOM_OAUTH2_HEADERS"))
|
||||
if static:
|
||||
try:
|
||||
# litellm's import runs load_dotenv and can inject .env values into os.environ;
|
||||
# snapshot and restore so we never leak unrelated keys into the process env.
|
||||
_before = dict(os.environ)
|
||||
import litellm
|
||||
|
||||
# drop keys litellm/load_dotenv added, restore any it changed (no empty-env window)
|
||||
for k in list(os.environ):
|
||||
if k not in _before:
|
||||
del os.environ[k]
|
||||
os.environ.update(_before)
|
||||
litellm.headers = {**(getattr(litellm, "headers", None) or {}), **static}
|
||||
log.info("headroom-oauth2: static upstream headers: %s", list(static))
|
||||
except Exception as e: # pragma: no cover
|
||||
log.warning("headroom-oauth2: could not set litellm.headers: %s", e)
|
||||
# The litellm backend auths bedrock/vertex/sagemaker from env and ignores a forwarded
|
||||
# bearer, so this extension is a no-op there -- warn loudly rather than silently do nothing.
|
||||
backend = str(getattr(config, "backend", "") or "").lower()
|
||||
if any(p in backend for p in ("bedrock", "vertex", "sagemaker")):
|
||||
log.warning(
|
||||
"headroom-oauth2: backend %r authenticates from env (bedrock/vertex/sagemaker) and "
|
||||
"ignores the injected bearer -- this extension will have NO effect. Use an "
|
||||
"OpenAI-compatible / passthrough backend (e.g. --backend litellm-openai).",
|
||||
backend or "<default>",
|
||||
)
|
||||
app.add_middleware(OAuth2Middleware, provider=provider)
|
||||
log.info(
|
||||
"headroom-oauth2: client-credentials auth installed (token_url=%s, style=%s)",
|
||||
provider.token_url,
|
||||
provider.auth_style,
|
||||
)
|
||||
61
plugins/headroom-oauth2/src/headroom_oauth2/middleware.py
Normal file
61
plugins/headroom-oauth2/src/headroom_oauth2/middleware.py
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
"""ASGI middleware that injects a refreshed OAuth2 bearer on each upstream request.
|
||||
|
||||
Headroom's litellm backend forwards the request's `Authorization` bearer to the
|
||||
upstream as the API key, so setting it here makes the minted token reach the
|
||||
backend with no core changes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
|
||||
from .provider import OAuth2Error
|
||||
|
||||
log = logging.getLogger("headroom_oauth2")
|
||||
|
||||
|
||||
class OAuth2Middleware:
|
||||
"""ASGI middleware that replaces the request Authorization with a minted bearer."""
|
||||
|
||||
def __init__(self, app, provider):
|
||||
self.app = app
|
||||
self.provider = provider
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
if scope.get("type") != "http":
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
# Hot path: a cached, still-valid token needs no thread hop. Only mint (blocking
|
||||
# urllib) off the event loop when the cache is empty/expired.
|
||||
token = self.provider.cached()
|
||||
if token is None:
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
token = await loop.run_in_executor(None, self.provider.token)
|
||||
except OAuth2Error as e:
|
||||
log.warning("oauth2: token mint failed: %s", e)
|
||||
await self._error(
|
||||
send, 502, "upstream_auth_error", "could not obtain upstream credentials"
|
||||
)
|
||||
return
|
||||
headers = [(k, v) for (k, v) in scope.get("headers", []) if k.lower() != b"authorization"]
|
||||
headers.append((b"authorization", b"Bearer " + token.encode()))
|
||||
await self.app(dict(scope, headers=headers), receive, send)
|
||||
|
||||
@staticmethod
|
||||
async def _error(send, status, etype, message):
|
||||
body = json.dumps({"type": "error", "error": {"type": etype, "message": message}}).encode()
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": status,
|
||||
"headers": [
|
||||
(b"content-type", b"application/json"),
|
||||
(b"content-length", str(len(body)).encode()),
|
||||
(b"cache-control", b"no-store"),
|
||||
],
|
||||
}
|
||||
)
|
||||
await send({"type": "http.response.body", "body": body})
|
||||
150
plugins/headroom-oauth2/src/headroom_oauth2/provider.py
Normal file
150
plugins/headroom-oauth2/src/headroom_oauth2/provider.py
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
"""Generic OAuth2 client-credentials token provider (RFC 6749 section 4.4).
|
||||
|
||||
Mints a bearer from a configurable token endpoint, caches it, and refreshes
|
||||
single-flight before expiry. Standard library only; `urllib` uses the system
|
||||
cert store (works behind corporate SSL inspection). No vendor specifics.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
from urllib.error import HTTPError, URLError
|
||||
|
||||
log = logging.getLogger("headroom_oauth2")
|
||||
|
||||
|
||||
def _https_or_local(url: str) -> bool:
|
||||
parts = urllib.parse.urlsplit(url)
|
||||
if parts.scheme == "https":
|
||||
return True
|
||||
# only numeric loopback -- "localhost" can be repointed via /etc/hosts or DNS rebinding
|
||||
return parts.scheme == "http" and (parts.hostname or "") in ("127.0.0.1", "::1")
|
||||
|
||||
|
||||
class OAuth2Error(RuntimeError):
|
||||
"""Raised when a token cannot be minted."""
|
||||
|
||||
|
||||
class OAuth2ClientCredentials:
|
||||
"""Mints and caches an OAuth2 client-credentials bearer (RFC 6749 section 4.4).
|
||||
|
||||
Thread-safe: ``token()`` refreshes single-flight before expiry; ``cached()`` is a
|
||||
lock-free read for the request hot path.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
token_url: str,
|
||||
client_id: str,
|
||||
client_secret: str,
|
||||
scopes=None,
|
||||
audience: str | None = None,
|
||||
grant_type: str = "client_credentials",
|
||||
auth_style: str = "post",
|
||||
extra_params=None,
|
||||
skew_seconds: int = 60,
|
||||
timeout_seconds: float = 30.0,
|
||||
allow_insecure: bool = False,
|
||||
):
|
||||
if not token_url:
|
||||
raise ValueError("token_url is required")
|
||||
if not client_id or not client_secret:
|
||||
raise ValueError("client_id and client_secret are required")
|
||||
if auth_style not in ("post", "basic"):
|
||||
raise ValueError("auth_style must be 'post' or 'basic'")
|
||||
if not allow_insecure and not _https_or_local(token_url):
|
||||
raise ValueError(
|
||||
"token_url must be https (loopback http allowed for tests; set "
|
||||
"allow_insecure=True / HEADROOM_OAUTH2_ALLOW_INSECURE=1 to override)"
|
||||
)
|
||||
self.token_url = token_url
|
||||
self.client_id = client_id
|
||||
self.client_secret = client_secret
|
||||
self.scopes = list(scopes or [])
|
||||
self.audience = audience
|
||||
self.grant_type = grant_type
|
||||
self.auth_style = auth_style
|
||||
self.extra_params = dict(extra_params or {})
|
||||
self.skew = max(0, int(skew_seconds))
|
||||
self.timeout = timeout_seconds
|
||||
self._lock = threading.Lock()
|
||||
self._token: str | None = None
|
||||
self._exp = 0.0
|
||||
self._eff_skew = self.skew
|
||||
|
||||
def _valid(self) -> bool:
|
||||
return self._token is not None and time.monotonic() < self._exp - self._eff_skew
|
||||
|
||||
def cached(self) -> str | None:
|
||||
"""Return the cached token if still valid, else None. No minting -- hot-path read."""
|
||||
return self._token if self._valid() else None
|
||||
|
||||
def token(self) -> str:
|
||||
"""Return a valid bearer, minting/refreshing single-flight if needed."""
|
||||
if self._valid():
|
||||
return self._token # type: ignore[return-value]
|
||||
with self._lock: # single-flight: one mint per burst
|
||||
if self._valid():
|
||||
return self._token # type: ignore[return-value]
|
||||
token, ttl = self._mint()
|
||||
# Publish _exp/_eff_skew BEFORE _token so a concurrent cached() reader never
|
||||
# sees a fresh token paired with a stale expiry.
|
||||
self._eff_skew = min(self.skew, max(0, ttl // 2))
|
||||
self._exp = time.monotonic() + ttl
|
||||
self._token = token
|
||||
return self._token
|
||||
|
||||
def _mint(self):
|
||||
form = dict(self.extra_params) # caller extras first; canonical fields below always win
|
||||
form["grant_type"] = self.grant_type
|
||||
if self.scopes:
|
||||
form["scope"] = " ".join(self.scopes)
|
||||
if self.audience:
|
||||
form["audience"] = self.audience
|
||||
headers = {
|
||||
"Content-Type": "application/x-www-form-urlencoded",
|
||||
"Accept": "application/json",
|
||||
}
|
||||
if self.auth_style == "basic":
|
||||
creds = base64.b64encode(f"{self.client_id}:{self.client_secret}".encode()).decode()
|
||||
headers["Authorization"] = "Basic " + creds
|
||||
else:
|
||||
form["client_id"] = self.client_id
|
||||
form["client_secret"] = self.client_secret
|
||||
req = urllib.request.Request(
|
||||
self.token_url,
|
||||
data=urllib.parse.urlencode(form).encode(),
|
||||
headers=headers,
|
||||
method="POST",
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=self.timeout) as resp:
|
||||
payload = json.load(resp)
|
||||
except HTTPError as e:
|
||||
try:
|
||||
e.read() # drain; do NOT surface the IdP body (may echo sensitive context)
|
||||
except Exception:
|
||||
pass
|
||||
raise OAuth2Error(f"token endpoint returned HTTP {e.code}") from None
|
||||
except (URLError, OSError) as e:
|
||||
raise OAuth2Error(f"token endpoint unreachable: {e}") from None
|
||||
except json.JSONDecodeError:
|
||||
raise OAuth2Error("token endpoint returned non-JSON") from None
|
||||
token = payload.get("access_token")
|
||||
if not token:
|
||||
raise OAuth2Error("token endpoint response had no access_token")
|
||||
raw = payload.get("expires_in")
|
||||
try:
|
||||
ttl = int(float(raw)) # tolerate "3600", "3600.0", 3600, or a JSON float
|
||||
except (TypeError, ValueError):
|
||||
ttl = 300
|
||||
ttl = max(1, ttl) # 0/negative would cause a stale token or per-request minting
|
||||
log.info("oauth2: minted token (ttl=%ss, scopes=%s)", ttl, self.scopes or "-")
|
||||
return token, ttl
|
||||
494
plugins/headroom-oauth2/tests/test_oauth2.py
Normal file
494
plugins/headroom-oauth2/tests/test_oauth2.py
Normal file
|
|
@ -0,0 +1,494 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom_oauth2 import (
|
||||
OAuth2ClientCredentials,
|
||||
OAuth2Error,
|
||||
OAuth2Middleware,
|
||||
_split,
|
||||
install,
|
||||
parse_headers,
|
||||
provider_from_env,
|
||||
)
|
||||
|
||||
|
||||
class _IdP(BaseHTTPRequestHandler):
|
||||
last_form = None
|
||||
last_auth = None
|
||||
status = 200
|
||||
tok = "TOK-1"
|
||||
expires_in = 3600
|
||||
mint_count = 0
|
||||
slow = False
|
||||
non_json = False
|
||||
omit_expires = False
|
||||
|
||||
def do_POST(self):
|
||||
n = int(self.headers.get("content-length", 0) or 0)
|
||||
_IdP.last_form = self.rfile.read(n).decode()
|
||||
_IdP.last_auth = self.headers.get("authorization")
|
||||
if _IdP.status != 200:
|
||||
self.send_response(_IdP.status)
|
||||
self.end_headers()
|
||||
self.wfile.write(b'{"error":"bad","error_description":"SENSITIVE"}')
|
||||
return
|
||||
if _IdP.slow:
|
||||
time.sleep(0.05) # widen the window so concurrent callers contend on the lock
|
||||
_IdP.mint_count += 1
|
||||
if _IdP.non_json:
|
||||
body = b"<html>not json SENSITIVE</html>"
|
||||
else:
|
||||
payload = {"access_token": _IdP.tok, "token_type": "Bearer"}
|
||||
if not _IdP.omit_expires:
|
||||
payload["expires_in"] = _IdP.expires_in
|
||||
body = json.dumps(payload).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("content-type", "application/json")
|
||||
self.send_header("content-length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def log_message(self, *a):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def idp():
|
||||
_IdP.last_form = _IdP.last_auth = None
|
||||
_IdP.status = 200
|
||||
_IdP.tok = "TOK-1"
|
||||
_IdP.expires_in = 3600
|
||||
_IdP.mint_count = 0
|
||||
_IdP.slow = False
|
||||
_IdP.non_json = False
|
||||
_IdP.omit_expires = False
|
||||
srv = HTTPServer(("127.0.0.1", 0), _IdP)
|
||||
threading.Thread(target=srv.serve_forever, daemon=True).start()
|
||||
yield f"http://127.0.0.1:{srv.server_address[1]}/token"
|
||||
srv.shutdown()
|
||||
|
||||
|
||||
# --- minimal ASGI test doubles -------------------------------------------------
|
||||
class _RecordingApp:
|
||||
def __init__(self):
|
||||
self.called = False
|
||||
self.scope = None
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
self.called = True
|
||||
self.scope = scope
|
||||
|
||||
|
||||
async def _recv():
|
||||
return {"type": "http.request"}
|
||||
|
||||
|
||||
async def _ignore(_msg):
|
||||
pass
|
||||
|
||||
|
||||
def _cfg(backend):
|
||||
return type("Cfg", (), {"backend": backend})()
|
||||
|
||||
|
||||
# --- provider: token minting ---------------------------------------------------
|
||||
def test_post_style_mint(idp):
|
||||
p = OAuth2ClientCredentials(
|
||||
token_url=idp, client_id="cid", client_secret="csec", scopes=["a", "b"], audience="aud"
|
||||
)
|
||||
assert p.token() == "TOK-1"
|
||||
assert "grant_type=client_credentials" in _IdP.last_form
|
||||
assert "scope=a+b" in _IdP.last_form
|
||||
assert "client_id=cid" in _IdP.last_form
|
||||
assert "audience=aud" in _IdP.last_form
|
||||
assert _IdP.last_auth is None
|
||||
|
||||
|
||||
def test_basic_style_mint(idp):
|
||||
p = OAuth2ClientCredentials(
|
||||
token_url=idp, client_id="cid", client_secret="csec", auth_style="basic"
|
||||
)
|
||||
p.token()
|
||||
assert _IdP.last_auth == "Basic " + base64.b64encode(b"cid:csec").decode()
|
||||
assert "client_secret" not in _IdP.last_form
|
||||
|
||||
|
||||
def test_cache_and_refresh(idp):
|
||||
p = OAuth2ClientCredentials(token_url=idp, client_id="c", client_secret="s")
|
||||
assert p.token() == "TOK-1"
|
||||
_IdP.tok = "TOK-2"
|
||||
_IdP.last_form = None
|
||||
assert p.token() == "TOK-1" # cached -> no re-mint
|
||||
assert _IdP.last_form is None
|
||||
p._exp = time.monotonic() - 1 # force expiry
|
||||
assert p.token() == "TOK-2" # re-minted
|
||||
|
||||
|
||||
def test_cached_fast_path(idp):
|
||||
p = OAuth2ClientCredentials(token_url=idp, client_id="c", client_secret="s")
|
||||
assert p.cached() is None # nothing minted yet -> middleware will mint off-loop
|
||||
p.token()
|
||||
assert p.cached() == "TOK-1" # now served without a token endpoint round-trip
|
||||
p._exp = time.monotonic() - 1
|
||||
assert p.cached() is None # expired -> forces a refresh
|
||||
|
||||
|
||||
def test_concurrent_single_flight(idp):
|
||||
p = OAuth2ClientCredentials(token_url=idp, client_id="c", client_secret="s")
|
||||
_IdP.slow = True
|
||||
out = []
|
||||
threads = [threading.Thread(target=lambda: out.append(p.token())) for _ in range(12)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
assert out == ["TOK-1"] * 12
|
||||
assert _IdP.mint_count == 1 # 12 concurrent callers -> exactly one mint
|
||||
|
||||
|
||||
# --- provider: failure modes ---------------------------------------------------
|
||||
def test_error_on_bad_status_hides_body(idp):
|
||||
_IdP.status = 401
|
||||
p = OAuth2ClientCredentials(token_url=idp, client_id="c", client_secret="s")
|
||||
with pytest.raises(OAuth2Error) as ei:
|
||||
p.token()
|
||||
assert "SENSITIVE" not in str(ei.value) # IdP error body must not leak into the exception
|
||||
|
||||
|
||||
def test_malformed_200_no_token(idp):
|
||||
_IdP.tok = None # HTTP 200 but no access_token field
|
||||
p = OAuth2ClientCredentials(token_url=idp, client_id="c", client_secret="s")
|
||||
with pytest.raises(OAuth2Error):
|
||||
p.token()
|
||||
|
||||
|
||||
def test_unreachable_token_url():
|
||||
p = OAuth2ClientCredentials(
|
||||
token_url="http://127.0.0.1:1/token", client_id="c", client_secret="s", timeout_seconds=1
|
||||
)
|
||||
with pytest.raises(OAuth2Error):
|
||||
p.token()
|
||||
|
||||
|
||||
def test_validation():
|
||||
with pytest.raises(ValueError):
|
||||
OAuth2ClientCredentials(token_url="", client_id="c", client_secret="s")
|
||||
with pytest.raises(ValueError):
|
||||
OAuth2ClientCredentials(token_url="u", client_id="", client_secret="s")
|
||||
with pytest.raises(ValueError):
|
||||
OAuth2ClientCredentials(token_url="u", client_id="c", client_secret="s", auth_style="x")
|
||||
|
||||
|
||||
def test_https_enforced():
|
||||
with pytest.raises(ValueError):
|
||||
OAuth2ClientCredentials(
|
||||
token_url="http://example.com/token", client_id="c", client_secret="s"
|
||||
)
|
||||
# loopback http allowed for local testing
|
||||
OAuth2ClientCredentials(token_url="http://127.0.0.1:1/token", client_id="c", client_secret="s")
|
||||
# explicit opt-out
|
||||
OAuth2ClientCredentials(
|
||||
token_url="http://example.com/token", client_id="c", client_secret="s", allow_insecure=True
|
||||
)
|
||||
|
||||
|
||||
def test_expires_in_clamp(idp):
|
||||
_IdP.expires_in = 0 # immediate-expiry -> must clamp to a positive ttl (not stale)
|
||||
p = OAuth2ClientCredentials(token_url=idp, client_id="c", client_secret="s")
|
||||
assert p.token() == "TOK-1"
|
||||
assert p._exp > time.monotonic()
|
||||
|
||||
_IdP.expires_in = -10 # negative -> must clamp (not perpetual re-mint)
|
||||
p2 = OAuth2ClientCredentials(token_url=idp, client_id="c", client_secret="s")
|
||||
assert p2.token() == "TOK-1"
|
||||
assert p2._exp > time.monotonic()
|
||||
|
||||
|
||||
# --- helpers / config ----------------------------------------------------------
|
||||
def test_helpers():
|
||||
assert _split("a, b c") == ["a", "b", "c"]
|
||||
assert parse_headers("X=1,Y=2") == {"X": "1", "Y": "2"}
|
||||
|
||||
|
||||
def test_parse_headers_rejects_control_chars():
|
||||
assert parse_headers("Good=ok,Bad=line\r\ninject") == {"Good": "ok"} # CRLF value dropped
|
||||
assert parse_headers("=novalue,K=v") == {"K": "v"} # empty key dropped
|
||||
assert parse_headers("") == {}
|
||||
|
||||
|
||||
def test_provider_from_env_wires_knobs(idp):
|
||||
env = {
|
||||
"HEADROOM_OAUTH2_TOKEN_URL": idp,
|
||||
"HEADROOM_OAUTH2_CLIENT_ID": "c",
|
||||
"HEADROOM_OAUTH2_CLIENT_SECRET": "s",
|
||||
"HEADROOM_OAUTH2_RESOURCE": "https://api.example",
|
||||
"HEADROOM_OAUTH2_TIMEOUT": "5",
|
||||
"HEADROOM_OAUTH2_SKEW": "10",
|
||||
}
|
||||
p = provider_from_env(env)
|
||||
assert p.extra_params == {"resource": "https://api.example"}
|
||||
assert p.timeout == 5.0
|
||||
assert p.skew == 10
|
||||
p.token()
|
||||
assert "resource=https" in _IdP.last_form
|
||||
|
||||
|
||||
def test_provider_from_env_none_when_unset():
|
||||
assert provider_from_env({}) is None
|
||||
|
||||
|
||||
# --- middleware (ASGI behavior) ------------------------------------------------
|
||||
def test_middleware_injects_bearer(idp):
|
||||
p = OAuth2ClientCredentials(token_url=idp, client_id="c", client_secret="s")
|
||||
app = _RecordingApp()
|
||||
mw = OAuth2Middleware(app, p)
|
||||
scope = {"type": "http", "headers": [(b"authorization", b"Bearer CLIENT"), (b"x-keep", b"1")]}
|
||||
asyncio.run(mw(scope, _recv, _ignore))
|
||||
hdrs = dict(app.scope["headers"])
|
||||
assert hdrs[b"authorization"] == b"Bearer TOK-1" # client creds replaced by minted token
|
||||
assert hdrs[b"x-keep"] == b"1" # other headers preserved
|
||||
|
||||
|
||||
def test_middleware_non_http_passthrough():
|
||||
app = _RecordingApp()
|
||||
mw = OAuth2Middleware(app, provider=object()) # provider must never be touched
|
||||
scope = {"type": "lifespan"}
|
||||
asyncio.run(mw(scope, _recv, _ignore))
|
||||
assert app.called and app.scope is scope
|
||||
|
||||
|
||||
def test_middleware_502_on_mint_failure():
|
||||
class _Bad:
|
||||
def cached(self):
|
||||
return None
|
||||
|
||||
def token(self):
|
||||
raise OAuth2Error("nope")
|
||||
|
||||
app = _RecordingApp()
|
||||
mw = OAuth2Middleware(app, _Bad())
|
||||
sent = []
|
||||
|
||||
async def send(msg):
|
||||
sent.append(msg)
|
||||
|
||||
asyncio.run(mw({"type": "http", "headers": []}, _recv, send))
|
||||
assert not app.called # request must not reach upstream without credentials
|
||||
assert sent[0]["status"] == 502
|
||||
assert json.loads(sent[1]["body"])["error"]["type"] == "upstream_auth_error"
|
||||
|
||||
|
||||
# --- install() (entry point) ---------------------------------------------------
|
||||
def test_install_noop_when_unset(monkeypatch):
|
||||
monkeypatch.delenv("HEADROOM_OAUTH2_TOKEN_URL", raising=False)
|
||||
|
||||
class App:
|
||||
def add_middleware(self, *a, **k):
|
||||
raise AssertionError("must not install middleware when unconfigured")
|
||||
|
||||
install(App(), _cfg("litellm-openai")) # no raise, no add_middleware
|
||||
|
||||
|
||||
def test_install_fail_closed_on_bad_config(monkeypatch):
|
||||
monkeypatch.setenv("HEADROOM_OAUTH2_TOKEN_URL", "https://idp.example.com/token")
|
||||
monkeypatch.setenv("HEADROOM_OAUTH2_CLIENT_ID", "") # missing -> ValueError -> RuntimeError
|
||||
with pytest.raises(RuntimeError):
|
||||
install(object(), _cfg("litellm-openai"))
|
||||
|
||||
|
||||
def test_install_warns_for_envauth_backend(monkeypatch, caplog):
|
||||
monkeypatch.setenv("HEADROOM_OAUTH2_TOKEN_URL", "https://idp.example.com/token")
|
||||
monkeypatch.setenv("HEADROOM_OAUTH2_CLIENT_ID", "c")
|
||||
monkeypatch.setenv("HEADROOM_OAUTH2_CLIENT_SECRET", "s")
|
||||
monkeypatch.delenv("HEADROOM_OAUTH2_HEADERS", raising=False)
|
||||
installed = []
|
||||
|
||||
class App:
|
||||
def add_middleware(self, *a, **k):
|
||||
installed.append(True)
|
||||
|
||||
with caplog.at_level("WARNING"):
|
||||
install(App(), _cfg("bedrock"))
|
||||
assert installed # still installs
|
||||
assert "NO effect" in caplog.text # but warns the bearer is ignored by env-auth backends
|
||||
|
||||
|
||||
def test_install_fail_closed_on_bad_timeout(monkeypatch):
|
||||
monkeypatch.setenv("HEADROOM_OAUTH2_TOKEN_URL", "https://idp.example.com/token")
|
||||
monkeypatch.setenv("HEADROOM_OAUTH2_CLIENT_ID", "c")
|
||||
monkeypatch.setenv("HEADROOM_OAUTH2_CLIENT_SECRET", "s")
|
||||
monkeypatch.setenv("HEADROOM_OAUTH2_TIMEOUT", "not-a-number") # invalid -> fail closed
|
||||
with pytest.raises(RuntimeError):
|
||||
install(object(), _cfg("litellm-openai"))
|
||||
|
||||
|
||||
def test_parse_headers_rejects_bad_keys():
|
||||
assert parse_headers("Bad Key=v,Ok=1") == {"Ok": "1"} # space in key dropped
|
||||
assert parse_headers("X:Y=v,Ok=1") == {"Ok": "1"} # colon in key dropped
|
||||
|
||||
|
||||
def test_expires_in_float(idp):
|
||||
_IdP.expires_in = 3599.9 # some IdPs return a JSON float -> must not fall back to 300
|
||||
p = OAuth2ClientCredentials(token_url=idp, client_id="c", client_secret="s")
|
||||
assert p.token() == "TOK-1"
|
||||
assert p._exp - time.monotonic() > 1000 # ~3599, not the 300 fallback
|
||||
|
||||
|
||||
def test_middleware_handles_missing_headers_key(idp):
|
||||
p = OAuth2ClientCredentials(token_url=idp, client_id="c", client_secret="s")
|
||||
app = _RecordingApp()
|
||||
mw = OAuth2Middleware(app, p)
|
||||
asyncio.run(mw({"type": "http"}, _recv, _ignore)) # scope without a "headers" key
|
||||
assert dict(app.scope["headers"])[b"authorization"] == b"Bearer TOK-1"
|
||||
|
||||
|
||||
def test_middleware_502_sets_no_store():
|
||||
class _Bad:
|
||||
def cached(self):
|
||||
return None
|
||||
|
||||
def token(self):
|
||||
raise OAuth2Error("nope")
|
||||
|
||||
app = _RecordingApp()
|
||||
mw = OAuth2Middleware(app, _Bad())
|
||||
sent = []
|
||||
|
||||
async def send(msg):
|
||||
sent.append(msg)
|
||||
|
||||
asyncio.run(mw({"type": "http", "headers": []}, _recv, send))
|
||||
hdrs = dict(sent[0]["headers"])
|
||||
assert hdrs[b"cache-control"] == b"no-store" # a fronting cache must not pin the 502
|
||||
|
||||
|
||||
# --- TLS / loopback edge cases -------------------------------------------------
|
||||
def test_localhost_rejected():
|
||||
# "localhost" is a name (DNS-rebinding / /etc/hosts risk) -> not a loopback exception
|
||||
with pytest.raises(ValueError):
|
||||
OAuth2ClientCredentials(
|
||||
token_url="http://localhost/token", client_id="c", client_secret="s"
|
||||
)
|
||||
|
||||
|
||||
def test_ipv6_loopback_allowed():
|
||||
OAuth2ClientCredentials(token_url="http://[::1]:1/token", client_id="c", client_secret="s")
|
||||
|
||||
|
||||
def test_allow_insecure_env_permits_nonloopback_http():
|
||||
p = provider_from_env(
|
||||
{
|
||||
"HEADROOM_OAUTH2_TOKEN_URL": "http://example.com/token",
|
||||
"HEADROOM_OAUTH2_CLIENT_ID": "c",
|
||||
"HEADROOM_OAUTH2_CLIENT_SECRET": "s",
|
||||
"HEADROOM_OAUTH2_ALLOW_INSECURE": "1",
|
||||
}
|
||||
)
|
||||
assert p.token_url == "http://example.com/token"
|
||||
|
||||
|
||||
# --- token-request form edge cases ---------------------------------------------
|
||||
def test_extra_params_cannot_override_canonical(idp):
|
||||
p = OAuth2ClientCredentials(
|
||||
token_url=idp,
|
||||
client_id="cid",
|
||||
client_secret="csec",
|
||||
scopes=["a", "b"],
|
||||
extra_params={"grant_type": "evil", "client_id": "evil", "scope": "evil", "resource": "r"},
|
||||
)
|
||||
p.token()
|
||||
assert "grant_type=client_credentials" in _IdP.last_form
|
||||
assert "client_id=cid" in _IdP.last_form
|
||||
assert "scope=a+b" in _IdP.last_form
|
||||
assert "resource=r" in _IdP.last_form # a benign extra still passes through
|
||||
assert "evil" not in _IdP.last_form # caller extras never clobber canonical fields
|
||||
|
||||
|
||||
def test_auth_style_basic_via_env(idp):
|
||||
p = provider_from_env(
|
||||
{
|
||||
"HEADROOM_OAUTH2_TOKEN_URL": idp,
|
||||
"HEADROOM_OAUTH2_CLIENT_ID": "cid",
|
||||
"HEADROOM_OAUTH2_CLIENT_SECRET": "csec",
|
||||
"HEADROOM_OAUTH2_AUTH_STYLE": "basic",
|
||||
}
|
||||
)
|
||||
assert p.auth_style == "basic"
|
||||
p.token()
|
||||
assert _IdP.last_auth == "Basic " + base64.b64encode(b"cid:csec").decode()
|
||||
|
||||
|
||||
def test_scopes_comma_separated_via_env(idp):
|
||||
p = provider_from_env(
|
||||
{
|
||||
"HEADROOM_OAUTH2_TOKEN_URL": idp,
|
||||
"HEADROOM_OAUTH2_CLIENT_ID": "c",
|
||||
"HEADROOM_OAUTH2_CLIENT_SECRET": "s",
|
||||
"HEADROOM_OAUTH2_SCOPES": "a, b ,c",
|
||||
}
|
||||
)
|
||||
assert p.scopes == ["a", "b", "c"]
|
||||
|
||||
|
||||
# --- expires_in edge cases -----------------------------------------------------
|
||||
def test_expires_in_missing_falls_back(idp):
|
||||
_IdP.omit_expires = True # no expires_in field -> default ttl, not stale
|
||||
p = OAuth2ClientCredentials(token_url=idp, client_id="c", client_secret="s")
|
||||
assert p.token() == "TOK-1"
|
||||
assert 100 < p._exp - time.monotonic() <= 300
|
||||
|
||||
|
||||
def test_expires_in_non_numeric_falls_back(idp):
|
||||
_IdP.expires_in = "not-a-number" # garbage -> default ttl, no crash
|
||||
p = OAuth2ClientCredentials(token_url=idp, client_id="c", client_secret="s")
|
||||
assert p.token() == "TOK-1"
|
||||
assert 100 < p._exp - time.monotonic() <= 300
|
||||
|
||||
|
||||
# --- response-shape failure modes ----------------------------------------------
|
||||
def test_non_json_200_raises_without_leak(idp):
|
||||
_IdP.non_json = True # HTTP 200 but body is not JSON
|
||||
p = OAuth2ClientCredentials(token_url=idp, client_id="c", client_secret="s")
|
||||
with pytest.raises(OAuth2Error) as ei:
|
||||
p.token()
|
||||
assert "SENSITIVE" not in str(ei.value) # body must not leak into the exception
|
||||
|
||||
|
||||
def test_single_flight_on_refresh(idp):
|
||||
p = OAuth2ClientCredentials(token_url=idp, client_id="c", client_secret="s")
|
||||
p.token() # cold mint #1
|
||||
assert _IdP.mint_count == 1
|
||||
p._exp = time.monotonic() - 1 # force expiry
|
||||
_IdP.slow = True
|
||||
threads = [threading.Thread(target=p.token) for _ in range(8)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
assert _IdP.mint_count == 2 # exactly one refresh despite 8 concurrent expired callers
|
||||
|
||||
|
||||
def test_install_sets_static_headers(monkeypatch):
|
||||
import sys
|
||||
import types
|
||||
|
||||
fake = types.ModuleType("litellm") # avoid importing the real (heavy) litellm
|
||||
fake.headers = {}
|
||||
monkeypatch.setitem(sys.modules, "litellm", fake)
|
||||
monkeypatch.setenv("HEADROOM_OAUTH2_TOKEN_URL", "https://idp.example.com/token")
|
||||
monkeypatch.setenv("HEADROOM_OAUTH2_CLIENT_ID", "c")
|
||||
monkeypatch.setenv("HEADROOM_OAUTH2_CLIENT_SECRET", "s")
|
||||
monkeypatch.setenv("HEADROOM_OAUTH2_HEADERS", "X-App=demo,Bad Key=x")
|
||||
|
||||
class App:
|
||||
def add_middleware(self, *a, **k):
|
||||
pass
|
||||
|
||||
install(App(), _cfg("litellm-openai"))
|
||||
assert fake.headers == {"X-App": "demo"} # valid header set on litellm; malformed key dropped
|
||||
Loading…
Add table
Add a link
Reference in a new issue