diff --git a/tests/test_litellm/proxy/auth_v2/auth_v2_helpers.py b/tests/test_litellm/proxy/auth_v2/auth_v2_helpers.py index d2e503ea93c..7a22d6fe2ac 100644 --- a/tests/test_litellm/proxy/auth_v2/auth_v2_helpers.py +++ b/tests/test_litellm/proxy/auth_v2/auth_v2_helpers.py @@ -17,7 +17,7 @@ class _StaticSigningKey: class FakeJwksClient: """Stands in for PyJWKClient. Returns one fixed key for every token so - JwtVerifier performs a real RS256 signature check against it via PyJWT.""" + JWTVerifier performs a real RS256 signature check against it via PyJWT.""" def __init__(self, public_key: Any) -> None: self._public_key = public_key diff --git a/tests/test_litellm/proxy/auth_v2/conftest.py b/tests/test_litellm/proxy/auth_v2/conftest.py index fe4df1984e0..2131ab25c78 100644 --- a/tests/test_litellm/proxy/auth_v2/conftest.py +++ b/tests/test_litellm/proxy/auth_v2/conftest.py @@ -6,8 +6,8 @@ import pytest from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa -from litellm.proxy.auth_v2.authenticators import JwtVerifier -from litellm.proxy.auth_v2.config import OidcProviderConfig +from litellm.proxy.auth_v2.authenticators import JWTVerifier +from litellm.proxy.auth_v2.config import OIDCProviderConfig from auth_v2_helpers import TEST_AUDIENCE, TEST_ISSUER, FakeJwksClient, TokenFactory @@ -39,13 +39,13 @@ def token_factory(rsa_keypair: Tuple[bytes, Any]) -> TokenFactory: @pytest.fixture -def oidc_provider() -> OidcProviderConfig: - return OidcProviderConfig(issuer=TEST_ISSUER, audience=[TEST_AUDIENCE]) +def oidc_provider() -> OIDCProviderConfig: + return OIDCProviderConfig(issuer=TEST_ISSUER, audience=[TEST_AUDIENCE]) @pytest.fixture def jwt_verifier( - rsa_keypair: Tuple[bytes, Any], oidc_provider: OidcProviderConfig -) -> JwtVerifier: + rsa_keypair: Tuple[bytes, Any], oidc_provider: OIDCProviderConfig +) -> JWTVerifier: _, public_key = rsa_keypair - return JwtVerifier(oidc_provider, jwks_client=FakeJwksClient(public_key)) + return JWTVerifier(oidc_provider, jwks_client=FakeJwksClient(public_key)) diff --git a/tests/test_litellm/proxy/auth_v2/test_authenticators.py b/tests/test_litellm/proxy/auth_v2/test_authenticators.py index 30a18d46192..b3da31202ce 100644 --- a/tests/test_litellm/proxy/auth_v2/test_authenticators.py +++ b/tests/test_litellm/proxy/auth_v2/test_authenticators.py @@ -7,13 +7,13 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.proxy.auth_v2.authenticators import ( - ApiKeyAuthenticator, + APIKeyAuthenticator, HttpAuthenticator, InMemoryBasicAuthStore, - JwtVerifier, - MutualTlsAuthenticator, + JWTVerifier, + MutualTLSAuthenticator, OAuth2Authenticator, - OidcAuthenticator, + OIDCAuthenticator, build_authenticators, hash_basic_password, ) @@ -21,13 +21,13 @@ from litellm.proxy.auth_v2.config import ( ApiKeySchemeConfig, AuthConfig, HttpBasicConfig, - MutualTlsConfig, + MutualTLSConfig, OAuth2IntrospectionConfig, - OidcProviderConfig, + OIDCProviderConfig, TrustedProxyConfig, ) from litellm.proxy.auth_v2.errors import AuthError -from litellm.proxy.auth_v2.models import AuthMethod, SecuritySchemeType +from litellm.proxy.auth_v2.models import AuthMethod from auth_v2_helpers import ( TEST_AUDIENCE, @@ -37,7 +37,7 @@ from auth_v2_helpers import ( ) # --------------------------------------------------------------------------- # -# JwtVerifier: every RFC 7519 check must be enforced. +# JWTVerifier: every RFC 7519 check must be enforced. # --------------------------------------------------------------------------- # @@ -51,7 +51,7 @@ def test_jwt_verifier_rejects_bad_signature( rsa_keypair, other_rsa_keypair, oidc_provider, token_factory ): _, public_key = rsa_keypair - verifier = JwtVerifier(oidc_provider, jwks_client=FakeJwksClient(public_key)) + verifier = JWTVerifier(oidc_provider, jwks_client=FakeJwksClient(public_key)) other_pem, _ = other_rsa_keypair forged = token_factory.mint(private_pem=other_pem) with pytest.raises(AuthError) as exc: @@ -99,12 +99,12 @@ def test_jwt_verifier_enforces_at_jwt_typ(jwt_verifier, token_factory): # --------------------------------------------------------------------------- # -# ApiKeyAuthenticator +# APIKeyAuthenticator # --------------------------------------------------------------------------- # async def test_api_key_authenticator_extracts_header(): - auth = ApiKeyAuthenticator(ApiKeySchemeConfig(header_name="x-litellm-api-key")) + auth = APIKeyAuthenticator(ApiKeySchemeConfig(header_name="x-litellm-api-key")) request = make_request(headers={"x-litellm-api-key": "sk-secret-value"}) credential = await auth.authenticate(request) assert credential is not None @@ -115,7 +115,7 @@ async def test_api_key_authenticator_extracts_header(): async def test_api_key_authenticator_returns_none_when_absent(): - auth = ApiKeyAuthenticator(ApiKeySchemeConfig()) + auth = APIKeyAuthenticator(ApiKeySchemeConfig()) assert await auth.authenticate(make_request()) is None @@ -127,8 +127,8 @@ async def test_api_key_authenticator_returns_none_when_absent(): def _http_auth( public_key: Any, *, basic: HttpBasicConfig = None, basic_verifier=None ) -> HttpAuthenticator: - verifier = JwtVerifier( - OidcProviderConfig(issuer=TEST_ISSUER, audience=[TEST_AUDIENCE]), + verifier = JWTVerifier( + OIDCProviderConfig(issuer=TEST_ISSUER, audience=[TEST_AUDIENCE]), jwks_client=FakeJwksClient(public_key), ) return HttpAuthenticator( @@ -265,8 +265,8 @@ def test_hash_basic_password_is_salted_and_verifiable(): def _oauth2(public_key: Any) -> OAuth2Authenticator: - verifier = JwtVerifier( - OidcProviderConfig(issuer=TEST_ISSUER, audience=[TEST_AUDIENCE]), + verifier = JWTVerifier( + OIDCProviderConfig(issuer=TEST_ISSUER, audience=[TEST_AUDIENCE]), jwks_client=FakeJwksClient(public_key), ) return OAuth2Authenticator([verifier], introspection=None) @@ -373,16 +373,16 @@ async def test_oauth2_introspection_non_200_raises(): # --------------------------------------------------------------------------- # -# OidcAuthenticator +# OIDCAuthenticator # --------------------------------------------------------------------------- # -def _oidc(public_key: Any) -> OidcAuthenticator: - verifier = JwtVerifier( - OidcProviderConfig(issuer=TEST_ISSUER, audience=[TEST_AUDIENCE]), +def _oidc(public_key: Any) -> OIDCAuthenticator: + verifier = JWTVerifier( + OIDCProviderConfig(issuer=TEST_ISSUER, audience=[TEST_AUDIENCE]), jwks_client=FakeJwksClient(public_key), ) - return OidcAuthenticator([verifier]) + return OIDCAuthenticator([verifier]) async def test_oidc_valid_token_sets_oidc_method(rsa_keypair, token_factory): @@ -405,7 +405,7 @@ async def test_oidc_unknown_issuer_raises(rsa_keypair, token_factory): # --------------------------------------------------------------------------- # -# MutualTlsAuthenticator +# MutualTLSAuthenticator # --------------------------------------------------------------------------- # @@ -413,12 +413,12 @@ async def test_oidc_unknown_issuer_raises(rsa_keypair, token_factory): _TRUSTED_NET = TrustedProxyConfig(trusted_proxy_cidrs=["203.0.113.0/24"]) -def _mtls(config: MutualTlsConfig, network: TrustedProxyConfig = None): - return MutualTlsAuthenticator(config, network or _TRUSTED_NET) +def _mtls(config: MutualTLSConfig, network: TrustedProxyConfig = None): + return MutualTLSAuthenticator(config, network or _TRUSTED_NET) async def test_mtls_reads_forwarded_subject_header_from_trusted_peer(): - auth = _mtls(MutualTlsConfig(enabled=True, forwarded_subject_header="x-client-dn")) + auth = _mtls(MutualTLSConfig(enabled=True, forwarded_subject_header="x-client-dn")) request = make_request(headers={"x-client-dn": "CN=svc-a,O=Co,C=US"}) credential = await auth.authenticate(request) assert credential is not None @@ -429,7 +429,7 @@ async def test_mtls_reads_forwarded_subject_header_from_trusted_peer(): async def test_mtls_forwarded_header_from_untrusted_peer_is_ignored(): # spoofing guard: a client that is not a trusted proxy cannot forge the DN header - auth = _mtls(MutualTlsConfig(enabled=True, forwarded_subject_header="x-client-dn")) + auth = _mtls(MutualTLSConfig(enabled=True, forwarded_subject_header="x-client-dn")) request = make_request( headers={"x-client-dn": "CN=attacker"}, client=("8.8.8.8", 4444) ) @@ -437,12 +437,12 @@ async def test_mtls_forwarded_header_from_untrusted_peer_is_ignored(): async def test_mtls_forwarded_header_absent_returns_none(): - auth = _mtls(MutualTlsConfig(enabled=True, forwarded_subject_header="x-client-dn")) + auth = _mtls(MutualTLSConfig(enabled=True, forwarded_subject_header="x-client-dn")) assert await auth.authenticate(make_request()) is None async def test_mtls_reads_asgi_tls_extension(): - auth = _mtls(MutualTlsConfig(enabled=True)) + auth = _mtls(MutualTLSConfig(enabled=True)) request = make_request( scope_extra={"extensions": {"tls": {"client_cert_name": "CN=from-asgi"}}} ) @@ -452,7 +452,7 @@ async def test_mtls_reads_asgi_tls_extension(): async def test_mtls_no_cert_returns_none(): - auth = _mtls(MutualTlsConfig(enabled=True)) + auth = _mtls(MutualTLSConfig(enabled=True)) assert await auth.authenticate(make_request()) is None @@ -463,24 +463,23 @@ async def test_mtls_no_cert_returns_none(): def test_build_authenticators_follows_scheme_order(): config = AuthConfig() - authenticators = build_authenticators(config) - schemes = [a.scheme for a in authenticators] + types = [type(a) for a in build_authenticators(config)] # mutual_tls disabled by default -> excluded - assert schemes == [ - SecuritySchemeType.API_KEY, - SecuritySchemeType.HTTP, - SecuritySchemeType.OPENID_CONNECT, - SecuritySchemeType.OAUTH2, + assert types == [ + APIKeyAuthenticator, + HttpAuthenticator, + OIDCAuthenticator, + OAuth2Authenticator, ] def test_build_authenticators_omits_api_key_when_unconfigured(): config = AuthConfig(api_key=None) - schemes = [a.scheme for a in build_authenticators(config)] - assert SecuritySchemeType.API_KEY not in schemes + types = [type(a) for a in build_authenticators(config)] + assert APIKeyAuthenticator not in types def test_build_authenticators_includes_mtls_when_enabled(): - config = AuthConfig(mutual_tls=MutualTlsConfig(enabled=True)) - schemes = [a.scheme for a in build_authenticators(config)] - assert SecuritySchemeType.MUTUAL_TLS == schemes[-1] + config = AuthConfig(mutual_tls=MutualTLSConfig(enabled=True)) + types = [type(a) for a in build_authenticators(config)] + assert types[-1] is MutualTLSAuthenticator diff --git a/tests/test_litellm/proxy/auth_v2/test_config.py b/tests/test_litellm/proxy/auth_v2/test_config.py index 86da462fff7..fd4c26100be 100644 --- a/tests/test_litellm/proxy/auth_v2/test_config.py +++ b/tests/test_litellm/proxy/auth_v2/test_config.py @@ -5,23 +5,23 @@ from pydantic import ValidationError from litellm.proxy.auth_v2.config import ( OAuth2IntrospectionConfig, - OidcProviderConfig, - SamlConfig, + OIDCProviderConfig, + SAMLConfig, ) def test_saml_config_requires_idp_metadata_when_enabled(): with pytest.raises(ValidationError): - SamlConfig(enabled=True, entity_id="sp", acs_url="https://sp/acs") + SAMLConfig(enabled=True, entity_id="sp", acs_url="https://sp/acs") def test_saml_config_allows_empty_metadata_when_disabled(): - config = SamlConfig(enabled=False, entity_id="sp", acs_url="https://sp/acs") + config = SAMLConfig(enabled=False, entity_id="sp", acs_url="https://sp/acs") assert config.idp_metadata == "" def test_saml_config_accepts_inline_metadata(): - config = SamlConfig( + config = SAMLConfig( enabled=True, entity_id="sp", acs_url="https://sp/acs", @@ -32,11 +32,11 @@ def test_saml_config_accepts_inline_metadata(): def test_oidc_provider_requires_audience(): with pytest.raises(ValidationError): - OidcProviderConfig(issuer="https://idp.example.com") + OIDCProviderConfig(issuer="https://idp.example.com") def test_oidc_provider_defaults_to_rs256(): - provider = OidcProviderConfig(issuer="https://idp.example.com", audience=["x"]) + provider = OIDCProviderConfig(issuer="https://idp.example.com", audience=["x"]) assert provider.algorithms == ["RS256"] assert provider.require_at_jwt is False diff --git a/tests/test_litellm/proxy/auth_v2/test_oidc.py b/tests/test_litellm/proxy/auth_v2/test_oidc.py index 52218423d6e..6c816dd338f 100644 --- a/tests/test_litellm/proxy/auth_v2/test_oidc.py +++ b/tests/test_litellm/proxy/auth_v2/test_oidc.py @@ -1,6 +1,6 @@ from __future__ import annotations -from litellm.proxy.auth_v2.config import OidcProviderConfig +from litellm.proxy.auth_v2.config import OIDCProviderConfig from litellm.proxy.auth_v2.oidc import _provider_key, _user_from_userinfo from litellm.proxy.auth_v2.resolver import InMemoryIdentityStore @@ -26,7 +26,7 @@ def test_userinfo_falls_back_to_email_when_no_preferred_username(): def test_provider_key_sanitizes_issuer_url(): key = _provider_key( - OidcProviderConfig(issuer="https://Login.Example.com/realm", audience=["x"]) + OIDCProviderConfig(issuer="https://Login.Example.com/realm", audience=["x"]) ) assert key == "https-login-example-com-realm" assert " " not in key diff --git a/tests/test_litellm/proxy/auth_v2/test_rbac.py b/tests/test_litellm/proxy/auth_v2/test_rbac.py index d16432b9fb3..d0094593995 100644 --- a/tests/test_litellm/proxy/auth_v2/test_rbac.py +++ b/tests/test_litellm/proxy/auth_v2/test_rbac.py @@ -4,7 +4,7 @@ import pytest from fastapi.security import SecurityScopes from litellm.proxy.auth_v2.models import AuthMethod, Principal, PrincipalType -from litellm.proxy.auth_v2.rbac import RbacEngine, Role, has_required_scopes +from litellm.proxy.auth_v2.rbac import RBACEngine, Role, has_required_scopes def _principal(*, scopes=None, roles=None) -> Principal: @@ -38,13 +38,13 @@ def test_empty_required_scopes_always_passes(): # --------------------------------------------------------------------------- # -# RbacEngine.has_role honors the role hierarchy (Casbin g-rules) +# RBACEngine.has_role honors the role hierarchy (Casbin g-rules) # --------------------------------------------------------------------------- # @pytest.fixture -def engine() -> RbacEngine: - return RbacEngine() +def engine() -> RBACEngine: + return RBACEngine() @pytest.mark.parametrize( @@ -61,7 +61,7 @@ def engine() -> RbacEngine: ], ) def test_has_role_inherits_down_the_hierarchy(engine, held, gate): - assert engine.has_role(_principal(roles=[held]), (gate,)) + assert engine.has_any_role(_principal(roles=[held]), (gate,)) @pytest.mark.parametrize( @@ -73,15 +73,15 @@ def test_has_role_inherits_down_the_hierarchy(engine, held, gate): ], ) def test_has_role_does_not_climb_the_hierarchy(engine, held, gate): - assert not engine.has_role(_principal(roles=[held]), (gate,)) + assert not engine.has_any_role(_principal(roles=[held]), (gate,)) def test_has_role_false_without_roles(engine): - assert not engine.has_role(_principal(), (Role.TEAM_MEMBER,)) + assert not engine.has_any_role(_principal(), (Role.TEAM_MEMBER,)) # --------------------------------------------------------------------------- # -# RbacEngine.enforce against the default policy +# RBACEngine.enforce against the default policy # --------------------------------------------------------------------------- # @@ -115,7 +115,7 @@ def test_enforce_false_without_roles(engine): def test_csv_policy_overrides_defaults(tmp_path): policy = tmp_path / "policy.csv" policy.write_text("p, platform_viewer, /reports, POST\n") - engine = RbacEngine(policy_path=str(policy)) + engine = RBACEngine(policy_path=str(policy)) # the operator rule is honored assert engine.enforce(_principal(roles=[Role.PLATFORM_VIEWER]), "/reports", "POST") diff --git a/tests/test_litellm/proxy/auth_v2/test_saml.py b/tests/test_litellm/proxy/auth_v2/test_saml.py index e3c5f655fe0..f5ade9b0863 100644 --- a/tests/test_litellm/proxy/auth_v2/test_saml.py +++ b/tests/test_litellm/proxy/auth_v2/test_saml.py @@ -94,7 +94,7 @@ def saml_env(tmp_path: Path) -> SamlEnv: from saml2.saml import NAMEID_FORMAT_EMAILADDRESS from saml2.server import Server - from litellm.proxy.auth_v2.config import SamlConfig + from litellm.proxy.auth_v2.config import SAMLConfig idp_key, idp_cert = _gen_cert(tmp_path, "idp") sp_key, sp_cert = _gen_cert(tmp_path, "sp") @@ -142,7 +142,7 @@ def saml_env(tmp_path: Path) -> SamlEnv: idp = Server(config=idp_conf) idp_metadata = str(entity_descriptor(idp.config)) - config = SamlConfig( + config = SAMLConfig( enabled=True, entity_id=SP_ENTITY_ID, acs_url=ACS_URL, @@ -161,22 +161,17 @@ def _build_app(saml_env: SamlEnv): from litellm.proxy.auth_v2.config import AuthConfig from litellm.proxy.auth_v2.models import Principal from litellm.proxy.auth_v2.resolver import InMemoryIdentityStore - from litellm.proxy.auth_v2.security import get_current_principal, install_auth + from litellm.proxy.auth_v2.saml import build_saml_router + from litellm.proxy.auth_v2.security import AuthSecurity app = FastAPI() store = InMemoryIdentityStore() - install_auth( - app, - AuthConfig(saml=saml_env.config), - store, - mount_scim=False, - mount_oidc=False, - mount_saml=True, - ) + auth = AuthSecurity(AuthConfig(saml=saml_env.config), store) + app.include_router(build_saml_router(auth)) @app.get("/whoami") async def whoami( - principal: "Principal" = Security(get_current_principal), + principal: "Principal" = Security(auth.principal), ): return { "subject": principal.subject, @@ -226,7 +221,7 @@ def test_acs_accepts_signed_assertion_and_provisions_user(saml_env): follow_redirects=False, ) assert acs.status_code == 303 - assert "saml_session" in acs.cookies + assert "litellm_session" in acs.cookies # user was provisioned into the ProvisioningStore via the shared upsert seam users = list(store._users.values()) @@ -243,7 +238,7 @@ def test_session_cookie_authenticates_with_saml_method(saml_env): data={"SAMLResponse": saml_env.mint_response()}, follow_redirects=False, ) - client.cookies.set("saml_session", acs.cookies["saml_session"]) + client.cookies.set("litellm_session", acs.cookies["litellm_session"]) whoami = client.get("/whoami") assert whoami.status_code == 200 @@ -404,9 +399,9 @@ def test_user_from_mapped_builds_name_and_email(): ], ) def test_safe_relay_state_blocks_open_redirects(candidate, expected): - from litellm.proxy.auth_v2.saml import _safe_relay_state + from litellm.proxy.auth_v2.session import safe_relay_state - assert _safe_relay_state(candidate, "/") == expected + assert safe_relay_state(candidate, "/") == expected @pytest.mark.parametrize( @@ -431,28 +426,28 @@ def test_acs_session_cookie_is_secure(saml_env): data={"SAMLResponse": saml_env.mint_response()}, follow_redirects=False, ) - assert "saml_session" in acs.cookies + assert "litellm_session" in acs.cookies assert "secure" in acs.headers["set-cookie"].lower() # --------------------------------------------------------------------------- # -# SamlSessionStore TTL + size eviction (no xmlsec1 needed) +# SessionStore TTL + size eviction (no xmlsec1 needed) # --------------------------------------------------------------------------- # def test_session_store_expires_entries(): - from litellm.proxy.auth_v2.saml import SamlSessionStore + from litellm.proxy.auth_v2.session import SessionStore - store = SamlSessionStore(ttl_seconds=0) + store = SessionStore(ttl_seconds=0) session_id = store.create_session({"name_id": "alice@example.com"}) # ttl of 0 means the entry is already past its expiry on the next read assert store.get(session_id) is None def test_session_store_evicts_when_over_capacity(): - from litellm.proxy.auth_v2.saml import SamlSessionStore + from litellm.proxy.auth_v2.session import SessionStore - store = SamlSessionStore(max_size=3) + store = SessionStore(max_size=3) ids = [store.create_session({"name_id": f"user-{i}"}) for i in range(5)] live = [sid for sid in ids if store.get(sid) is not None] assert len(live) <= 3 diff --git a/tests/test_litellm/proxy/auth_v2/test_scim.py b/tests/test_litellm/proxy/auth_v2/test_scim.py index e05e9086115..ce463c0ed36 100644 --- a/tests/test_litellm/proxy/auth_v2/test_scim.py +++ b/tests/test_litellm/proxy/auth_v2/test_scim.py @@ -7,7 +7,8 @@ from fastapi.testclient import TestClient from litellm.proxy.auth_v2.config import AuthConfig from litellm.proxy.auth_v2.models import AuthMethod, Principal, PrincipalType from litellm.proxy.auth_v2.resolver import InMemoryIdentityStore, _hash_api_key -from litellm.proxy.auth_v2.security import install_auth +from litellm.proxy.auth_v2.scim import build_scim_router +from litellm.proxy.auth_v2.security import AuthSecurity USER_SCHEMA = "urn:ietf:params:scim:schemas:core:2.0:User" GROUP_SCHEMA = "urn:ietf:params:scim:schemas:core:2.0:Group" @@ -28,19 +29,14 @@ def _principal(subject: str, scopes: list) -> Principal: def _app() -> FastAPI: app = FastAPI() - install_auth( - app, - AuthConfig(), - InMemoryIdentityStore( - api_keys={ - _hash_api_key(SCIM_KEY): _principal("scim-writer", ["scim:write"]), - _hash_api_key(NOSCOPE_KEY): _principal("no-scope", []), - } - ), - mount_scim=True, - mount_oidc=False, - mount_saml=False, + store = InMemoryIdentityStore( + api_keys={ + _hash_api_key(SCIM_KEY): _principal("scim-writer", ["scim:write"]), + _hash_api_key(NOSCOPE_KEY): _principal("no-scope", []), + } ) + auth = AuthSecurity(AuthConfig(), store) + app.include_router(build_scim_router(auth)) return app diff --git a/tests/test_litellm/proxy/auth_v2/test_security.py b/tests/test_litellm/proxy/auth_v2/test_security.py index f9f92c1e0ed..96d240d86e5 100644 --- a/tests/test_litellm/proxy/auth_v2/test_security.py +++ b/tests/test_litellm/proxy/auth_v2/test_security.py @@ -1,31 +1,26 @@ from __future__ import annotations -from typing import Annotated, Any, Tuple +from typing import Any, Tuple import pytest from fastapi import FastAPI, Security from fastapi.testclient import TestClient from litellm.proxy.auth_v2.authenticators import ( - ApiKeyAuthenticator, + APIKeyAuthenticator, HttpAuthenticator, - JwtVerifier, + JWTVerifier, ) from litellm.proxy.auth_v2.config import ( ApiKeySchemeConfig, AuthConfig, HttpBasicConfig, - OidcProviderConfig, + OIDCProviderConfig, ) from litellm.proxy.auth_v2.models import AuthMethod, Principal, PrincipalType -from litellm.proxy.auth_v2.rbac import RbacEngine, Role +from litellm.proxy.auth_v2.rbac import RBACEngine, Role from litellm.proxy.auth_v2.resolver import InMemoryIdentityStore, _hash_api_key -from litellm.proxy.auth_v2.security import ( - AuthContext, - get_current_principal, - require_permission, - require_roles, -) +from litellm.proxy.auth_v2.security import AuthSecurity from auth_v2_helpers import TEST_AUDIENCE, TEST_ISSUER, FakeJwksClient @@ -46,13 +41,15 @@ def _principal(subject: str, *, scopes=None, roles=None) -> Principal: ) -def _build_app(public_key: Any) -> Tuple[FastAPI, InMemoryIdentityStore]: - verifier = JwtVerifier( - OidcProviderConfig(issuer=TEST_ISSUER, audience=[TEST_AUDIENCE]), +def _build_app( + public_key: Any, *, rbac: RBACEngine = None +) -> Tuple[FastAPI, InMemoryIdentityStore]: + verifier = JWTVerifier( + OIDCProviderConfig(issuer=TEST_ISSUER, audience=[TEST_AUDIENCE]), jwks_client=FakeJwksClient(public_key), ) authenticators = [ - ApiKeyAuthenticator(ApiKeySchemeConfig()), + APIKeyAuthenticator(ApiKeySchemeConfig()), HttpAuthenticator(HttpBasicConfig(), [verifier]), ] resolver = InMemoryIdentityStore( @@ -72,15 +69,17 @@ def _build_app(public_key: Any) -> Tuple[FastAPI, InMemoryIdentityStore]: ), } ) - ctx = AuthContext(AuthConfig(), authenticators, resolver) + auth = AuthSecurity( + AuthConfig(), resolver, rbac=rbac, authenticators=authenticators + ) app = FastAPI() - app.state.auth_v2 = ctx + # default-value Security() style: the dependency marker stays a real object even + # under `from __future__ import annotations`, where Annotated[...] would be a string + # that FastAPI re-evaluates in module globals (the closure-local `auth` is invisible) @app.get("/open") - async def open_route( - principal: Annotated[Principal, Security(get_current_principal)], - ): + async def open_route(principal: Principal = Security(auth.principal)): return { "subject": principal.subject, "auth_method": principal.auth_method.value, @@ -89,23 +88,19 @@ def _build_app(public_key: Any) -> Tuple[FastAPI, InMemoryIdentityStore]: @app.get("/scoped") async def scoped_route( - principal: Annotated[ - Principal, Security(get_current_principal, scopes=["models:read"]) - ], + principal: Principal = Security(auth.principal, scopes=["models:read"]), ): return {"subject": principal.subject} @app.get("/admin") async def admin_route( - principal: Annotated[Principal, Security(require_roles(Role.ORG_ADMIN))], + principal: Principal = Security(auth.require_roles(Role.ORG_ADMIN)), ): return {"subject": principal.subject} @app.post("/perm-widgets") async def widgets_route( - principal: Annotated[ - Principal, Security(require_permission("/widgets", "POST")) - ], + principal: Principal = Security(auth.require_permission("/widgets", "POST")), ): return {"subject": principal.subject} @@ -229,7 +224,7 @@ def test_required_role_honors_hierarchy(client): # --------------------------------------------------------------------------- # -# Permission enforcement (require_permission -> RbacEngine.enforce) +# Permission enforcement (require_permission -> RBACEngine.enforce) # --------------------------------------------------------------------------- # @@ -262,8 +257,7 @@ def test_injected_rbac_engine_overrides_default_policy(rsa_keypair, tmp_path): policy.write_text("p, platform_viewer, /widgets, POST\n") _, public_key = rsa_keypair - app, _ = _build_app(public_key) - app.state.auth_v2.rbac = RbacEngine(policy_path=str(policy)) + app, _ = _build_app(public_key, rbac=RBACEngine(policy_path=str(policy))) client = TestClient(app) # viewer now passes, platform_admin (default grant removed) now fails