phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(fix) Fix - don't allow viewer roles to create virtual keys (#6764)

Browse source 
* fix ui route permissions

* fix test_is_ui_route_allowed

* fix test_is_ui_route_allowed

* test_user_role_permissions
This commit is contained in:
Ishaan Jaff 2024-11-15 18:02:13 -08:00 • committed by GitHub
parent 9ba8f40bd1
commit 6d2ee70a9a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 148 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/_types.py
View file

@ -278,6 +278,7 @@ class LiteLLMRoutes(enum.Enum):
management_routes = [ # key
"/key/generate",
"/key/{token_id}/regenerate",
"/key/update",
"/key/delete",
"/key/info",
@ -339,11 +340,7 @@ class LiteLLMRoutes(enum.Enum):
"/sso",
"/sso/get/ui_settings",
"/login",
"/key/generate",
"/key/{token_id}/regenerate",
"/key/update",
"/key/info",
"/key/delete",
"/config",
"/spend",
"/user",
@ -364,6 +361,7 @@ class LiteLLMRoutes(enum.Enum):
internal_user_routes = (
[
"/key/generate",
"/key/{token_id}/regenerate",
"/key/update",
"/key/delete",
"/key/health",

43
litellm/proxy/auth/route_checks.py
View file

@ -1,5 +1,5 @@
import re
from typing import Optional
from typing import List, Optional
from fastapi import HTTPException, Request, status
@ -80,7 +80,9 @@ class RouteChecks:
status_code=status.HTTP_403_FORBIDDEN,
detail=f"user not allowed to access this OpenAI routes, role= {_user_role}",
)
if route in LiteLLMRoutes.management_routes.value:
if RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.management_routes.value
):
# the Admin Viewer is only allowed to call /user/update for their own user_id and can only update
if route == "/user/update":
@ -101,21 +103,27 @@ class RouteChecks:
elif (
_user_role == LitellmUserRoles.INTERNAL_USER.value
and route in LiteLLMRoutes.internal_user_routes.value
and RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.internal_user_routes.value
)
):
pass
elif (
_user_is_org_admin(request_data=request_data, user_object=user_obj)
and route in LiteLLMRoutes.org_admin_allowed_routes.value
elif _user_is_org_admin(
request_data=request_data, user_object=user_obj
) and RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.org_admin_allowed_routes.value
):
pass
elif (
_user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value
and route in LiteLLMRoutes.internal_user_view_only_routes.value
and RouteChecks.check_route_access(
route=route,
allowed_routes=LiteLLMRoutes.internal_user_view_only_routes.value,
)
):
pass
elif (
route in LiteLLMRoutes.self_managed_routes.value
elif RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.self_managed_routes.value
): # routes that manage their own allowed/disallowed logic
pass
else:
@ -207,3 +215,20 @@ class RouteChecks:
if re.match(pattern, route):
return True
return False
@staticmethod
def check_route_access(route: str, allowed_routes: List[str]) -> bool:
"""
Check if a route has access by checking both exact matches and patterns
Args:
route (str): The route to check
allowed_routes (list): List of allowed routes/patterns
Returns:
bool: True if route is allowed, False otherwise
"""
return route in allowed_routes or any( # Check exact match
RouteChecks._route_matches_pattern(route=route, pattern=allowed_route)
for allowed_route in allowed_routes
) # Check pattern match

18
litellm/proxy/auth/user_api_key_auth.py
View file

@ -111,12 +111,12 @@ def _get_bearer_token(
return api_key
def _is_ui_route_allowed(
def _is_ui_route(
route: str,
user_obj: Optional[LiteLLM_UserTable] = None,
) -> bool:
"""
- Route b/w ui token check and normal token check
- Check if the route is a UI used route
"""
# this token is only used for managing the ui
allowed_routes = LiteLLMRoutes.ui_routes.value
@ -133,15 +133,7 @@ def _is_ui_route_allowed(
for allowed_route in allowed_routes
):
return True
else:
if user_obj is not None and _is_user_proxy_admin(user_obj=user_obj):
return True
elif _has_user_setup_sso() and route in LiteLLMRoutes.sso_only_routes.value:
return True
else:
raise Exception(
f"This key is made for LiteLLM UI, Tried to access route: {route}. Not allowed"
)
return False
def _is_api_route_allowed(
@ -185,8 +177,8 @@ def _is_allowed_route(
"""
- Route b/w ui token check and normal token check
"""
if token_type == "ui":
return _is_ui_route_allowed(route=route, user_obj=user_obj)
if token_type == "ui" and _is_ui_route(route=route, user_obj=user_obj):
return True
else:
return _is_api_route_allowed(
route=route,

16
tests/local_testing/test_user_api_key_auth.py
View file

@ -305,14 +305,14 @@ async def test_auth_with_allowed_routes(route, should_raise_error):
[
# Proxy Admin checks
("/global/spend/logs", "proxy_admin", True),
("/key/delete", "proxy_admin", True),
("/key/generate", "proxy_admin", True),
("/key/regenerate", "proxy_admin", True),
("/key/delete", "proxy_admin", False),
("/key/generate", "proxy_admin", False),
("/key/regenerate", "proxy_admin", False),
# Internal User checks - allowed routes
("/global/spend/logs", "internal_user", True),
("/key/delete", "internal_user", True),
("/key/generate", "internal_user", True),
("/key/82akk800000000jjsk/regenerate", "internal_user", True),
("/key/delete", "internal_user", False),
("/key/generate", "internal_user", False),
("/key/82akk800000000jjsk/regenerate", "internal_user", False),
# Internal User Viewer
("/key/generate", "internal_user_viewer", False),
# Internal User checks - disallowed routes
@ -320,7 +320,7 @@ async def test_auth_with_allowed_routes(route, should_raise_error):
],
)
def test_is_ui_route_allowed(route, user_role, expected_result):
from litellm.proxy.auth.user_api_key_auth import _is_ui_route_allowed
from litellm.proxy.auth.user_api_key_auth import _is_ui_route
from litellm.proxy._types import LiteLLM_UserTable
user_obj = LiteLLM_UserTable(
@ -342,7 +342,7 @@ def test_is_ui_route_allowed(route, user_role, expected_result):
"user_obj": user_obj,
}
try:
assert _is_ui_route_allowed(**received_args) == expected_result
assert _is_ui_route(**received_args) == expected_result
except Exception as e:
# If expected result is False, we expect an error
if expected_result is False:

91
tests/proxy_admin_ui_tests/test_role_based_access.py
View file

@ -437,3 +437,94 @@ async def test_org_admin_create_user_team_wrong_org_permissions(prisma_client):
"You do not have the required role to call" in e.message
and org2_id in e.message
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"route, user_role, expected_result",
[
# Proxy Admin checks
("/global/spend/logs", LitellmUserRoles.PROXY_ADMIN, True),
("/key/delete", LitellmUserRoles.PROXY_ADMIN, True),
("/key/generate", LitellmUserRoles.PROXY_ADMIN, True),
("/key/regenerate", LitellmUserRoles.PROXY_ADMIN, True),
# # Internal User checks - allowed routes
("/global/spend/logs", LitellmUserRoles.INTERNAL_USER, True),
("/key/delete", LitellmUserRoles.INTERNAL_USER, True),
("/key/generate", LitellmUserRoles.INTERNAL_USER, True),
("/key/82akk800000000jjsk/regenerate", LitellmUserRoles.INTERNAL_USER, True),
# Internal User Viewer
("/key/generate", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, False),
(
"/key/82akk800000000jjsk/regenerate",
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
False,
),
("/key/delete", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, False),
("/team/new", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, False),
("/team/delete", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, False),
("/team/update", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, False),
# Proxy Admin Viewer
("/global/spend/logs", LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, True),
("/key/delete", LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, False),
("/key/generate", LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, False),
(
"/key/82akk800000000jjsk/regenerate",
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
False,
),
("/team/new", LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, False),
("/team/delete", LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, False),
("/team/update", LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, False),
# Internal User checks - disallowed routes
("/organization/member_add", LitellmUserRoles.INTERNAL_USER, False),
],
)
async def test_user_role_permissions(prisma_client, route, user_role, expected_result):
"""Test user role based permissions for different routes"""
try:
# Setup
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
# Admin - admin creates a new user
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-1234",
user_id="1234",
)
request = NewUserRequest(user_role=user_role)
new_user_response = await new_user(request, user_api_key_dict=user_api_key_dict)
user_id = new_user_response.user_id
# Generate key for new user with team_id="litellm-dashboard"
key_response = await generate_key_fn(
data=GenerateKeyRequest(user_id=user_id, team_id="litellm-dashboard"),
user_api_key_dict=user_api_key_dict,
)
generated_key = key_response.key
bearer_token = "Bearer " + generated_key
# Create request with route
request = Request(scope={"type": "http"})
request._url = URL(url=route)
# Test authorization
if expected_result is True:
# Should pass without error
result = await user_api_key_auth(request=request, api_key=bearer_token)
print(f"Auth passed as expected for {route} with role {user_role}")
else:
# Should raise an error
with pytest.raises(Exception) as exc_info:
await user_api_key_auth(request=request, api_key=bearer_token)
print(f"Auth failed as expected for {route} with role {user_role}")
print(f"Error message: {str(exc_info.value)}")
except Exception as e:
if expected_result:
pytest.fail(f"Expected success but got exception: {str(e)}")
else:
print(f"Got expected exception: {str(e)}")

16
tests/proxy_unit_tests/test_user_api_key_auth.py
View file

@ -305,14 +305,14 @@ async def test_auth_with_allowed_routes(route, should_raise_error):
[
# Proxy Admin checks
("/global/spend/logs", "proxy_admin", True),
("/key/delete", "proxy_admin", True),
("/key/generate", "proxy_admin", True),
("/key/regenerate", "proxy_admin", True),
("/key/delete", "proxy_admin", False),
("/key/generate", "proxy_admin", False),
("/key/regenerate", "proxy_admin", False),
# Internal User checks - allowed routes
("/global/spend/logs", "internal_user", True),
("/key/delete", "internal_user", True),
("/key/generate", "internal_user", True),
("/key/82akk800000000jjsk/regenerate", "internal_user", True),
("/key/delete", "internal_user", False),
("/key/generate", "internal_user", False),
("/key/82akk800000000jjsk/regenerate", "internal_user", False),
# Internal User Viewer
("/key/generate", "internal_user_viewer", False),
# Internal User checks - disallowed routes
@ -320,7 +320,7 @@ async def test_auth_with_allowed_routes(route, should_raise_error):
],
)
def test_is_ui_route_allowed(route, user_role, expected_result):
from litellm.proxy.auth.user_api_key_auth import _is_ui_route_allowed
from litellm.proxy.auth.user_api_key_auth import _is_ui_route
from litellm.proxy._types import LiteLLM_UserTable
user_obj = LiteLLM_UserTable(
@ -342,7 +342,7 @@ def test_is_ui_route_allowed(route, user_role, expected_result):
"user_obj": user_obj,
}
try:
assert _is_ui_route_allowed(**received_args) == expected_result
assert _is_ui_route(**received_args) == expected_result
except Exception as e:
# If expected result is False, we expect an error
if expected_result is False: