phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

(UI) - Security Improvement, move to JWT Auth for Admin UI Sessions (#8995)

Browse source 
* (UI) - Improvements to session handling logic  (#8970)

* add cookieUtils

* use utils for clearing cookies

* on logout use clearTokenCookies

* ui use correct clearTokenCookies

* navbar show userEmail on UserID page

* add timestamp on token cookie

* update generate_authenticated_redirect_response

* use common getAuthToken

* fix clearTokenCookies

* fixes for get auth token

* fix invitation link sign in logic

* Revert "fix invitation link sign in logic"

This reverts commit 30e5308cb3.

* fix getAuthToken

* update setAuthToken

* fix ui session handling

* fix ui session handler

* bug fix stop generating LiteLLM Virtual keys for access

* working JWT insert into cookies

* use central place to build UI JWT token

* add _validate_ui_token

* fix ui session handler

* fix fetchWithCredentials

* check allowed routes for ui session tokens

* expose validate_session endpoint

* validate session endpoint

* call sso/session/validate

* getUISessionDetails

* ui move to getUISessionDetails

* /sso/session/validate

* fix cookie utils

* use getUISessionDetails

* use ui_session_id

* "/spend/logs/ui" in spend_tracking_routes

* working sign in JWT flow for proxy admin

* allow proxy admin to access ui routes

* use check_route_access

* update types

* update login method

* fixes to ui session handler

* working flow for admin and internal users

* fixes for invite links

* use JWTs for SSO sign in

* fix /invitation/new flow

* fix code quality checks

* fix _get_ui_session_token_from_cookies

* /organization/list

* ui sso sign in

* TestUISessionHandler

* TestUISessionHandler
This commit is contained in:
Ishaan Jaff 2025-03-04 21:48:23 -08:00 • committed by GitHub
parent 42931638df
commit 01a44a4e47
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 1104 additions and 538 deletions
Download patch file Download diff file Expand all files Collapse all files

64
litellm/proxy/auth/handle_jwt.py
View file

@ -33,6 +33,7 @@ from litellm.proxy._types import (
ScopeMapping,
Span,
)
from litellm.proxy.management_helpers.ui_session_handler import UISessionHandler
from litellm.proxy.utils import PrismaClient, ProxyLogging
from .auth_checks import (
@ -406,10 +407,60 @@ class JWTHandler:
else:
return False
def _validate_ui_token(self, token: str) -> Optional[dict]:
"""
Helper function to validate tokens generated for the LiteLLM UI.
Returns the decoded payload if it's a valid UI token, None otherwise.
"""
import jwt
from litellm.proxy.proxy_server import master_key
try:
# Decode without verification to check if it's a UI token
unverified_payload = jwt.decode(token, options={"verify_signature": False})
# Check if this looks like a UI token (has specific claims that only UI tokens would have)
if UISessionHandler.is_ui_session_token(unverified_payload):
# This looks like a UI token, now verify it with the master key
if not master_key:
verbose_proxy_logger.debug(
"Missing LITELLM_MASTER_KEY for UI token validation"
)
return None
try:
payload = jwt.decode(
token,
master_key,
algorithms=["HS256"],
audience="litellm-ui",
leeway=self.leeway,
)
verbose_proxy_logger.debug(
"Successfully validated UI token for payload: %s",
json.dumps(payload, indent=4),
)
return payload
except jwt.InvalidTokenError as e:
verbose_proxy_logger.debug(f"Invalid UI token: {str(e)}")
raise ValueError(
f"Invalid UI token, Unable to validate token signature {str(e)}"
)
return None # Not a UI token
except Exception as e:
raise e
async def auth_jwt(self, token: str) -> dict:
# Supported algos: https://pyjwt.readthedocs.io/en/stable/algorithms.html
# "Warning: Make sure not to mix symmetric and asymmetric algorithms that interpret
# the key in different ways (e.g. HS* and RS*)."
ui_payload = self._validate_ui_token(token)
if ui_payload:
return ui_payload
algorithms = ["RS256", "RS384", "RS512", "PS256", "PS384", "PS512"]
audience = os.getenv("JWT_AUDIENCE")
@ -616,6 +667,7 @@ class JWTAuthManager:
user_id: Optional[str],
org_id: Optional[str],
api_key: str,
jwt_valid_token: dict,
) -> Optional[JWTAuthBuilderResult]:
"""Check admin status and route access permissions"""
if not jwt_handler.is_admin(scopes=scopes):
@ -625,6 +677,7 @@ class JWTAuthManager:
user_role=LitellmUserRoles.PROXY_ADMIN,
user_route=route,
litellm_proxy_roles=jwt_handler.litellm_jwtauth,
jwt_valid_token=jwt_valid_token,
)
if not is_allowed:
allowed_routes: List[Any] = jwt_handler.litellm_jwtauth.admin_allowed_routes
@ -698,6 +751,7 @@ class JWTAuthManager:
user_api_key_cache: DualCache,
parent_otel_span: Optional[Span],
proxy_logging_obj: ProxyLogging,
jwt_valid_token: dict,
) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]:
"""Find first team with access to the requested model"""
@ -730,6 +784,7 @@ class JWTAuthManager:
user_role=LitellmUserRoles.TEAM,
user_route=route,
litellm_proxy_roles=jwt_handler.litellm_jwtauth,
jwt_valid_token=jwt_valid_token,
)
if is_allowed:
return team_id, team_object
@ -920,7 +975,13 @@ class JWTAuthManager:
# Check admin access
admin_result = await JWTAuthManager.check_admin_access(
jwt_handler, scopes, route, user_id, org_id, api_key
jwt_handler=jwt_handler,
scopes=scopes,
route=route,
user_id=user_id,
org_id=org_id,
api_key=api_key,
jwt_valid_token=jwt_valid_token,
)
if admin_result:
return admin_result
@ -952,6 +1013,7 @@ class JWTAuthManager:
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
jwt_valid_token=jwt_valid_token,
)
# Get other objects