diff --git a/plugins/headroom-oauth2/src/headroom_oauth2/__init__.py b/plugins/headroom-oauth2/src/headroom_oauth2/__init__.py index e4ffe31d3..b202caf39 100644 --- a/plugins/headroom-oauth2/src/headroom_oauth2/__init__.py +++ b/plugins/headroom-oauth2/src/headroom_oauth2/__init__.py @@ -13,7 +13,7 @@ from typing import Any from .middleware import OAuth2Middleware from .provider import OAuth2ClientCredentials, OAuth2Error -__all__ = ["install", "OAuth2ClientCredentials", "OAuth2Error", "OAuth2Middleware", "parse_headers"] +__all__ = ["OAuth2ClientCredentials", "OAuth2Error", "OAuth2Middleware", "install", "parse_headers"] __version__ = "0.1.0" log = logging.getLogger("headroom_oauth2") @@ -116,7 +116,13 @@ def install(app: Any, config: Any) -> None: 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 + except ( + ImportError, + AttributeError, + OSError, + TypeError, + ValueError, + ) 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. diff --git a/plugins/headroom-oauth2/src/headroom_oauth2/provider.py b/plugins/headroom-oauth2/src/headroom_oauth2/provider.py index dbcf92df6..49f00d996 100644 --- a/plugins/headroom-oauth2/src/headroom_oauth2/provider.py +++ b/plugins/headroom-oauth2/src/headroom_oauth2/provider.py @@ -14,6 +14,7 @@ import threading import time import urllib.parse import urllib.request +from contextlib import suppress from urllib.error import HTTPError, URLError log = logging.getLogger("headroom_oauth2") @@ -128,10 +129,8 @@ class OAuth2ClientCredentials: with urllib.request.urlopen(req, timeout=self.timeout) as resp: payload = json.load(resp) except HTTPError as e: - try: + with suppress(Exception): 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 diff --git a/plugins/headroom-oauth2/tests/test_oauth2.py b/plugins/headroom-oauth2/tests/test_oauth2.py index 5372cb599..4a8b03d35 100644 --- a/plugins/headroom-oauth2/tests/test_oauth2.py +++ b/plugins/headroom-oauth2/tests/test_oauth2.py @@ -1,9 +1,11 @@ import asyncio import base64 import json +import logging import threading import time from http.server import BaseHTTPRequestHandler, HTTPServer +from urllib.error import HTTPError import pytest @@ -161,6 +163,26 @@ def test_error_on_bad_status_hides_body(idp): assert "SENSITIVE" not in str(ei.value) # IdP error body must not leak into the exception +def test_error_body_drain_failure_is_sanitized(monkeypatch, caplog): + class _UnreadableHTTPError(HTTPError): + def read(self): + raise RuntimeError("SENSITIVE") + + def fail(*_args, **_kwargs): + raise _UnreadableHTTPError("https://idp.example/token", 503, "unavailable", {}, None) + + monkeypatch.setattr("urllib.request.urlopen", fail) + caplog.set_level(logging.DEBUG, logger="headroom_oauth2") + p = OAuth2ClientCredentials( + token_url="https://idp.example/token", client_id="c", client_secret="s" + ) + + with pytest.raises(OAuth2Error, match="HTTP 503") as exc_info: + p.token() + assert "SENSITIVE" not in str(exc_info.value) + assert "SENSITIVE" not in caplog.text + + 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") @@ -492,3 +514,24 @@ def test_install_sets_static_headers(monkeypatch): install(App(), _cfg("litellm-openai")) assert fake.headers == {"X-App": "demo"} # valid header set on litellm; malformed key dropped + + +def test_install_handles_invalid_litellm_headers(monkeypatch, caplog): + import sys + import types + + fake = types.ModuleType("litellm") + fake.headers = object() + 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") + caplog.set_level(logging.WARNING, logger="headroom_oauth2") + + class App: + def add_middleware(self, *a, **k): + pass + + install(App(), _cfg("litellm-openai")) + assert "could not set litellm.headers" in caplog.text