diff --git a/.circleci/config.yml b/.circleci/config.yml index 6bbc95b5a423..8fcf51376c90 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -119,7 +119,7 @@ jobs: paths: - local_testing_coverage.xml - local_testing_coverage - ui_endpoint_testing: + auth_ui_unit_tests: docker: - image: cimg/python:3.11 auth: @@ -161,8 +161,8 @@ jobs: - run: name: Rename the coverage files command: | - mv coverage.xml ui_endpoint_testing_coverage.xml - mv .coverage ui_endpoint_testing_coverage + mv coverage.xml auth_ui_unit_tests_coverage.xml + mv .coverage auth_ui_unit_tests_coverage # Store test results - store_test_results: @@ -171,8 +171,8 @@ jobs: - persist_to_workspace: root: . paths: - - ui_endpoint_testing_coverage.xml - - ui_endpoint_testing_coverage + - auth_ui_unit_tests_coverage.xml + - auth_ui_unit_tests_coverage litellm_router_testing: # Runs all tests with the "router" keyword docker: - image: cimg/python:3.11 @@ -813,7 +813,7 @@ jobs: python -m venv venv . venv/bin/activate pip install coverage - coverage combine llm_translation_coverage logging_coverage litellm_router_coverage local_testing_coverage litellm_assistants_api_coverage ui_endpoint_testing_coverage + coverage combine llm_translation_coverage logging_coverage litellm_router_coverage local_testing_coverage litellm_assistants_api_coverage auth_ui_unit_tests_coverage coverage xml - codecov/upload: file: ./coverage.xml @@ -1013,7 +1013,7 @@ workflows: only: - main - /litellm_.*/ - - ui_endpoint_testing: + - auth_ui_unit_tests: filters: branches: only: @@ -1062,7 +1062,7 @@ workflows: - litellm_router_testing - local_testing - litellm_assistants_api_testing - - ui_endpoint_testing + - auth_ui_unit_tests - db_migration_disable_update_check: filters: branches: @@ -1090,7 +1090,7 @@ workflows: - logging_testing - litellm_router_testing - litellm_assistants_api_testing - - ui_endpoint_testing + - auth_ui_unit_tests - db_migration_disable_update_check - e2e_ui_testing - installing_litellm_on_python diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5201bfe1ed2d..206cb235e7ab 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2515,6 +2515,7 @@ def get_standard_logging_metadata( user_api_key_hash=None, user_api_key_alias=None, user_api_key_team_id=None, + user_api_key_org_id=None, user_api_key_user_id=None, user_api_key_team_alias=None, spend_logs_metadata=None, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 629e002b56f5..2da8674392c1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -340,6 +340,7 @@ class LiteLLMRoutes(enum.Enum): "/sso/get/ui_settings", "/login", "/key/generate", + "/key/{token_id}/regenerate", "/key/update", "/key/info", "/key/delete", diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 940c74b92c05..87a7b9ce274f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -28,7 +28,7 @@ LitellmUserRoles, UserAPIKeyAuth, ) -from litellm.proxy.auth.route_checks import is_llm_api_route +from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.utils import PrismaClient, ProxyLogging, log_to_opentelemetry from litellm.types.services import ServiceLoggerPayload, ServiceTypes @@ -138,7 +138,7 @@ def common_checks( # noqa: PLR0915 general_settings.get("enforce_user_param", None) is not None and general_settings["enforce_user_param"] is True ): - if is_llm_api_route(route=route) and "user" not in request_body: + if RouteChecks.is_llm_api_route(route=route) and "user" not in request_body: raise Exception( f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}" ) @@ -154,7 +154,7 @@ def common_checks( # noqa: PLR0915 + CommonProxyErrors.not_premium_user.value ) - if is_llm_api_route(route=route): + if RouteChecks.is_llm_api_route(route=route): # loop through each enforced param # example enforced_params ['user', 'metadata', 'metadata.generation_name'] for enforced_param in general_settings["enforced_params"]: @@ -182,7 +182,7 @@ def common_checks( # noqa: PLR0915 and global_proxy_spend is not None # only run global budget checks for OpenAI routes # Reason - the Admin UI should continue working if the proxy crosses it's global budget - and is_llm_api_route(route=route) + and RouteChecks.is_llm_api_route(route=route) and route != "/v1/models" and route != "/models" ): diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index cc8fd3113644..a237b0bdd60e 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -17,175 +17,199 @@ from .auth_utils import _has_user_setup_sso -def non_proxy_admin_allowed_routes_check( - user_obj: Optional[LiteLLM_UserTable], - _user_role: Optional[LitellmUserRoles], - route: str, - request: Request, - valid_token: UserAPIKeyAuth, - api_key: str, - request_data: dict, -): - """ - Checks if Non Proxy Admin User is allowed to access the route - """ - - # Check user has defined custom admin routes - custom_admin_only_route_check( - route=route, - ) - - if is_llm_api_route(route=route): - pass - elif ( - route in LiteLLMRoutes.info_routes.value - ): # check if user allowed to call an info route - if route == "/key/info": - # check if user can access this route - query_params = request.query_params - key = query_params.get("key") - if key is not None and hash_token(token=key) != api_key: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail="user not allowed to access this key's info", +class RouteChecks: + + @staticmethod + def non_proxy_admin_allowed_routes_check( + user_obj: Optional[LiteLLM_UserTable], + _user_role: Optional[LitellmUserRoles], + route: str, + request: Request, + valid_token: UserAPIKeyAuth, + api_key: str, + request_data: dict, + ): + """ + Checks if Non Proxy Admin User is allowed to access the route + """ + + # Check user has defined custom admin routes + RouteChecks.custom_admin_only_route_check( + route=route, + ) + + if RouteChecks.is_llm_api_route(route=route): + pass + elif ( + route in LiteLLMRoutes.info_routes.value + ): # check if user allowed to call an info route + if route == "/key/info": + # check if user can access this route + query_params = request.query_params + key = query_params.get("key") + if key is not None and hash_token(token=key) != api_key: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="user not allowed to access this key's info", + ) + elif route == "/user/info": + # check if user can access this route + query_params = request.query_params + user_id = query_params.get("user_id") + verbose_proxy_logger.debug( + f"user_id: {user_id} & valid_token.user_id: {valid_token.user_id}" ) - elif route == "/user/info": - # check if user can access this route - query_params = request.query_params - user_id = query_params.get("user_id") - verbose_proxy_logger.debug( - f"user_id: {user_id} & valid_token.user_id: {valid_token.user_id}" - ) - if user_id and user_id != valid_token.user_id: + if user_id and user_id != valid_token.user_id: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="key not allowed to access this user's info. user_id={}, key's user_id={}".format( + user_id, valid_token.user_id + ), + ) + elif route == "/model/info": + # /model/info just shows models user has access to + pass + elif route == "/team/info": + pass # handled by function itself + elif _has_user_setup_sso() and route in LiteLLMRoutes.sso_only_routes.value: + pass + elif ( + route in LiteLLMRoutes.global_spend_tracking_routes.value + and getattr(valid_token, "permissions", None) is not None + and "get_spend_routes" in getattr(valid_token, "permissions", []) + ): + + pass + elif _user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value: + if RouteChecks.is_llm_api_route(route=route): raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail="key not allowed to access this user's info. user_id={}, key's user_id={}".format( - user_id, valid_token.user_id - ), + detail=f"user not allowed to access this OpenAI routes, role= {_user_role}", ) - elif route == "/model/info": - # /model/info just shows models user has access to + if route in 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": + + # Check the Request params are valid for PROXY_ADMIN_VIEW_ONLY + if request_data is not None and isinstance(request_data, dict): + _params_updated = request_data.keys() + for param in _params_updated: + if param not in ["user_email", "password"]: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email and password can be updated", + ) + else: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", + ) + + elif ( + _user_role == LitellmUserRoles.INTERNAL_USER.value + and route in LiteLLMRoutes.internal_user_routes.value + ): pass - elif route == "/team/info": - pass # handled by function itself - elif _has_user_setup_sso() and route in LiteLLMRoutes.sso_only_routes.value: - pass - elif ( - route in LiteLLMRoutes.global_spend_tracking_routes.value - and getattr(valid_token, "permissions", None) is not None - and "get_spend_routes" in getattr(valid_token, "permissions", []) - ): - - pass - elif _user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value: - if is_llm_api_route(route=route): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this OpenAI routes, role= {_user_role}", + elif ( + _user_is_org_admin(request_data=request_data, user_object=user_obj) + and route in 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 + ): + pass + elif ( + route in LiteLLMRoutes.self_managed_routes.value + ): # routes that manage their own allowed/disallowed logic + pass + else: + user_role = "unknown" + user_id = "unknown" + if user_obj is not None: + user_role = user_obj.user_role or "unknown" + user_id = user_obj.user_id or "unknown" + raise Exception( + f"Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route={route}. Your role={user_role}. Your user_id={user_id}" ) - if route in 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": - - # Check the Request params are valid for PROXY_ADMIN_VIEW_ONLY - if request_data is not None and isinstance(request_data, dict): - _params_updated = request_data.keys() - for param in _params_updated: - if param not in ["user_email", "password"]: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email and password can be updated", - ) - else: + + @staticmethod + def custom_admin_only_route_check(route: str): + from litellm.proxy.proxy_server import general_settings, premium_user + + if "admin_only_routes" in general_settings: + if premium_user is not True: + verbose_proxy_logger.error( + f"Trying to use 'admin_only_routes' this is an Enterprise only feature. {CommonProxyErrors.not_premium_user.value}" + ) + return + if route in general_settings["admin_only_routes"]: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", + detail=f"user not allowed to access this route. Route={route} is an admin only route", ) - - elif ( - _user_role == LitellmUserRoles.INTERNAL_USER.value - and route in 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 - ): pass - elif ( - _user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value - and route in LiteLLMRoutes.internal_user_view_only_routes.value - ): - pass - elif ( - route in LiteLLMRoutes.self_managed_routes.value - ): # routes that manage their own allowed/disallowed logic - pass - else: - user_role = "unknown" - user_id = "unknown" - if user_obj is not None: - user_role = user_obj.user_role or "unknown" - user_id = user_obj.user_id or "unknown" - raise Exception( - f"Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route={route}. Your role={user_role}. Your user_id={user_id}" - ) - -def custom_admin_only_route_check(route: str): - from litellm.proxy.proxy_server import general_settings, premium_user - - if "admin_only_routes" in general_settings: - if premium_user is not True: - verbose_proxy_logger.error( - f"Trying to use 'admin_only_routes' this is an Enterprise only feature. {CommonProxyErrors.not_premium_user.value}" - ) - return - if route in general_settings["admin_only_routes"]: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route. Route={route} is an admin only route", - ) - pass - - -def is_llm_api_route(route: str) -> bool: - """ - Helper to checks if provided route is an OpenAI route - - - Returns: - - True: if route is an OpenAI route - - False: if route is not an OpenAI route - """ - - if route in LiteLLMRoutes.openai_routes.value: - return True - - if route in LiteLLMRoutes.anthropic_routes.value: - return True - - # fuzzy match routes like "/v1/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" - # Check for routes with placeholders - for openai_route in LiteLLMRoutes.openai_routes.value: - # Replace placeholders with regex pattern - # placeholders are written as "/threads/{thread_id}" - if "{" in openai_route: - pattern = re.sub(r"\{[^}]+\}", r"[^/]+", openai_route) - # Anchor the pattern to match the entire string - pattern = f"^{pattern}$" - if re.match(pattern, route): - return True - - # Pass through Bedrock, VertexAI, and Cohere Routes - if "/bedrock/" in route: - return True - if "/vertex-ai/" in route: - return True - if "/gemini/" in route: - return True - if "/cohere/" in route: - return True - if "/langfuse/" in route: - return True - return False + @staticmethod + def is_llm_api_route(route: str) -> bool: + """ + Helper to checks if provided route is an OpenAI route + + + Returns: + - True: if route is an OpenAI route + - False: if route is not an OpenAI route + """ + + if route in LiteLLMRoutes.openai_routes.value: + return True + + if route in LiteLLMRoutes.anthropic_routes.value: + return True + + # fuzzy match routes like "/v1/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" + # Check for routes with placeholders + for openai_route in LiteLLMRoutes.openai_routes.value: + # Replace placeholders with regex pattern + # placeholders are written as "/threads/{thread_id}" + if "{" in openai_route: + if RouteChecks._route_matches_pattern( + route=route, pattern=openai_route + ): + return True + + # Pass through Bedrock, VertexAI, and Cohere Routes + if "/bedrock/" in route: + return True + if "/vertex-ai/" in route: + return True + if "/gemini/" in route: + return True + if "/cohere/" in route: + return True + if "/langfuse/" in route: + return True + return False + + @staticmethod + def _route_matches_pattern(route: str, pattern: str) -> bool: + """ + Check if route matches the pattern placed in proxy/_types.py + + Example: + - pattern: "/threads/{thread_id}" + - route: "/threads/thread_49EIN5QF32s4mH20M7GFKdlZ" + - returns: True + + + - pattern: "/key/{token_id}/regenerate" + - route: "/key/regenerate/82akk800000000jjsk" + - returns: False, pattern is "/key/{token_id}/regenerate" + """ + pattern = re.sub(r"\{[^}]+\}", r"[^/]+", pattern) + # Anchor the pattern to match the entire string + pattern = f"^{pattern}$" + if re.match(pattern, route): + return True + return False diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 20b262d3176f..bbdddeee9b23 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -69,7 +69,7 @@ ) from litellm.proxy.auth.oauth2_check import check_oauth2_token from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request -from litellm.proxy.auth.route_checks import non_proxy_admin_allowed_routes_check +from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.service_account_checks import service_account_checks from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.utils import _to_ns @@ -122,6 +122,11 @@ def _is_ui_route_allowed( ): # Do something if the current route starts with any of the allowed routes return True + elif any( + RouteChecks._route_matches_pattern(route=route, pattern=allowed_route) + 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 @@ -150,7 +155,7 @@ def _is_api_route_allowed( raise Exception("Invalid proxy server token passed") if not _is_user_proxy_admin(user_obj=user_obj): # if non-admin - non_proxy_admin_allowed_routes_check( + RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=_user_role, route=route, diff --git a/tests/local_testing/test_proxy_routes.py b/tests/local_testing/test_proxy_routes.py index 41ea0e1b4bdf..31ff7d2ede20 100644 --- a/tests/local_testing/test_proxy_routes.py +++ b/tests/local_testing/test_proxy_routes.py @@ -22,7 +22,7 @@ import litellm from litellm.proxy._types import LiteLLMRoutes from litellm.proxy.auth.auth_utils import get_request_route -from litellm.proxy.auth.route_checks import is_llm_api_route +from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.proxy_server import app # Configure logging @@ -84,7 +84,7 @@ def test_routes_on_litellm_proxy(): ], ) def test_is_llm_api_route(route: str, expected: bool): - assert is_llm_api_route(route) == expected + assert RouteChecks.is_llm_api_route(route) == expected # Test-case for routes that are similar but should return False @@ -98,12 +98,12 @@ def test_is_llm_api_route(route: str, expected: bool): ], ) def test_is_llm_api_route_similar_but_false(route: str): - assert is_llm_api_route(route) == False + assert RouteChecks.is_llm_api_route(route) is False def test_anthropic_api_routes(): # allow non proxy admins to call anthropic api routes - assert is_llm_api_route(route="/v1/messages") is True + assert RouteChecks.is_llm_api_route(route="/v1/messages") is True def create_request(path: str, base_url: str = "http://testserver") -> Request: diff --git a/tests/local_testing/test_user_api_key_auth.py b/tests/local_testing/test_user_api_key_auth.py index 47f96ccf22f0..1baddc7dd813 100644 --- a/tests/local_testing/test_user_api_key_auth.py +++ b/tests/local_testing/test_user_api_key_auth.py @@ -293,26 +293,50 @@ async def test_auth_with_allowed_routes(route, should_raise_error): setattr(proxy_server, "general_settings", initial_general_settings) -@pytest.mark.parametrize("route", ["/global/spend/logs", "/key/delete"]) -def test_is_ui_route_allowed(route): +@pytest.mark.parametrize( + "route, user_role, expected_result", + [ + # Proxy Admin checks + ("/global/spend/logs", "proxy_admin", True), + ("/key/delete", "proxy_admin", True), + ("/key/generate", "proxy_admin", True), + ("/key/regenerate", "proxy_admin", True), + # 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), + # Internal User checks - disallowed routes + ("/organization/member_add", "internal_user", False), + ], +) +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._types import LiteLLM_UserTable + user_obj = LiteLLM_UserTable( + user_id="3b803c0e-666e-4e99-bd5c-6e534c07e297", + max_budget=None, + spend=0.0, + model_max_budget={}, + model_spend={}, + user_email="my-test-email@1234.com", + models=[], + tpm_limit=None, + rpm_limit=None, + user_role=user_role, + organization_memberships=[], + ) + received_args: dict = { "route": route, - "user_obj": LiteLLM_UserTable( - user_id="3b803c0e-666e-4e99-bd5c-6e534c07e297", - max_budget=None, - spend=0.0, - model_max_budget={}, - model_spend={}, - user_email="my-test-email@1234.com", - models=[], - tpm_limit=None, - rpm_limit=None, - user_role="internal_user", - organization_memberships=[], - ), + "user_obj": user_obj, } - - assert _is_ui_route_allowed(**received_args) + try: + assert _is_ui_route_allowed(**received_args) == expected_result + except Exception as e: + # If expected result is False, we expect an error + if expected_result is False: + pass + else: + raise e diff --git a/tests/proxy_admin_ui_tests/test_route_check_unit_tests.py b/tests/proxy_admin_ui_tests/test_route_check_unit_tests.py new file mode 100644 index 000000000000..9a4ec84671a5 --- /dev/null +++ b/tests/proxy_admin_ui_tests/test_route_check_unit_tests.py @@ -0,0 +1,199 @@ +import os +import sys +import traceback +import uuid +import datetime as dt +from datetime import datetime + +from dotenv import load_dotenv +from fastapi import Request +from fastapi.routing import APIRoute + +load_dotenv() +import io +import os +import time + + +# this file is to test litellm/proxy + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import asyncio +import logging + +from fastapi import HTTPException, Request +import pytest +from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth + +# Replace the actual hash_token function with our mock +import litellm.proxy.auth.route_checks + + +# Mock objects and functions +class MockRequest: + def __init__(self, query_params=None): + self.query_params = query_params or {} + + +def mock_hash_token(token): + return token + + +litellm.proxy.auth.route_checks.hash_token = mock_hash_token + + +# Test is_llm_api_route +def test_is_llm_api_route(): + assert RouteChecks.is_llm_api_route("/v1/chat/completions") is True + assert RouteChecks.is_llm_api_route("/v1/completions") is True + assert RouteChecks.is_llm_api_route("/v1/embeddings") is True + assert RouteChecks.is_llm_api_route("/v1/images/generations") is True + assert RouteChecks.is_llm_api_route("/v1/threads/thread_12345") is True + assert RouteChecks.is_llm_api_route("/bedrock/model/invoke") is True + assert RouteChecks.is_llm_api_route("/vertex-ai/text") is True + assert RouteChecks.is_llm_api_route("/gemini/generate") is True + assert RouteChecks.is_llm_api_route("/cohere/generate") is True + + # check non-matching routes + assert RouteChecks.is_llm_api_route("/some/random/route") is False + assert RouteChecks.is_llm_api_route("/key/regenerate/82akk800000000jjsk") is False + assert RouteChecks.is_llm_api_route("/key/82akk800000000jjsk/delete") is False + + +# Test _route_matches_pattern +def test_route_matches_pattern(): + # check matching routes + assert ( + RouteChecks._route_matches_pattern( + "/threads/thread_12345", "/threads/{thread_id}" + ) + is True + ) + assert ( + RouteChecks._route_matches_pattern( + "/key/regenerate/82akk800000000jjsk", "/key/{token_id}/regenerate" + ) + is False + ) + assert ( + RouteChecks._route_matches_pattern( + "/v1/chat/completions", "/v1/chat/completions" + ) + is True + ) + assert ( + RouteChecks._route_matches_pattern( + "/v1/models/gpt-4", "/v1/models/{model_name}" + ) + is True + ) + + # check non-matching routes + assert ( + RouteChecks._route_matches_pattern( + "/v1/chat/completionz/thread_12345", "/v1/chat/completions/{thread_id}" + ) + is False + ) + assert ( + RouteChecks._route_matches_pattern( + "/v1/{thread_id}/messages", "/v1/messages/thread_2345" + ) + is False + ) + + +@pytest.fixture +def route_checks(): + return RouteChecks() + + +def test_llm_api_route(route_checks): + """ + Internal User is allowed to access all LLM API routes + """ + assert ( + route_checks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/v1/chat/completions", + request=MockRequest(), + valid_token=UserAPIKeyAuth(api_key="test_key"), + api_key="test_key", + request_data={}, + ) + is None + ) + + +def test_key_info_route_allowed(route_checks): + """ + Internal User is allowed to access /key/info route + """ + assert ( + route_checks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/key/info", + request=MockRequest(query_params={"key": "test_key"}), + valid_token=UserAPIKeyAuth(api_key="test_key"), + api_key="test_key", + request_data={}, + ) + is None + ) + + +def test_key_info_route_forbidden(route_checks): + """ + Internal User is not allowed to access /key/info route for a key they're not using in Authenticated API Key + """ + with pytest.raises(HTTPException) as exc_info: + route_checks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/key/info", + request=MockRequest(query_params={"key": "wrong_key"}), + valid_token=UserAPIKeyAuth(api_key="test_key"), + api_key="test_key", + request_data={}, + ) + assert exc_info.value.status_code == 403 + + +def test_user_info_route_allowed(route_checks): + """ + Internal User is allowed to access /user/info route for their own user_id + """ + assert ( + route_checks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/user/info", + request=MockRequest(query_params={"user_id": "test_user"}), + valid_token=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + api_key="test_key", + request_data={}, + ) + is None + ) + + +def test_user_info_route_forbidden(route_checks): + """ + Internal User is not allowed to access /user/info route for a different user_id + """ + with pytest.raises(HTTPException) as exc_info: + route_checks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/user/info", + request=MockRequest(query_params={"user_id": "wrong_user"}), + valid_token=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + api_key="test_key", + request_data={}, + ) + assert exc_info.value.status_code == 403