Merge pull request #1836 from BerriAI/litellm_fix_litellm_dashboard_keys

[Fix] UI - Security - Litellm UI Keys meant for litellm-dashboard shouldn't be allowed to make non-management related requests

litellm/proxy/proxy_server.py

@@ -636,6 +636,31 @@ async def user_api_key_auth(
                     raise Exception(
                         f"Only master key can be used to generate, delete, update or get info for new keys/users. Value of allow_user_auth={allow_user_auth}"
                     )
+
+        # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions
+        # sso/login, ui/login, /key functions and /user functions
+        # this will never be allowed to call /chat/completions
+        token_team = getattr(valid_token, "team_id", None)
+        if token_team is not None:
+            if token_team == "litellm-dashboard":
+                # this token is only used for managing the ui
+                allowed_routes = [
+                    "/sso",
+                    "/login",
+                    "/key",
+                    "/spend",
+                    "/user",
+                ]
+                # check if the current route startswith any of the allowed routes
+                if any(
+                    route.startswith(allowed_route) for allowed_route in allowed_routes
+                ):
+                    # Do something if the current route starts with any of the allowed routes
+                    pass
+                else:
+                    raise Exception(
+                        f"This key is made for LiteLLM UI, Tried to access route: {route}. Not allowed"
+                    )
             return UserAPIKeyAuth(api_key=api_key, **valid_token_dict)
         else:
             raise Exception(f"Invalid Key Passed to LiteLLM Proxy")

litellm/tests/test_key_generate_prisma.py

@@ -44,6 +44,7 @@ from litellm.proxy.proxy_server import (
     info_key_fn,
     update_key_fn,
     generate_key_fn,
+    generate_key_helper_fn,
     spend_user_fn,
     spend_key_fn,
     view_spend_logs,
@@ -1378,3 +1379,35 @@ async def test_user_api_key_auth_without_master_key(prisma_client):
     except Exception as e:
         print("Got Exception", e)
         pytest.fail(f"Got exception {e}")
+
+
+@pytest.mark.asyncio
+async def test_key_with_no_permissions(prisma_client):
+    """
+    - create key
+    - get key info
+    - assert key_name is null
+    """
+    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    setattr(litellm.proxy.proxy_server, "general_settings", {"allow_user_auth": False})
+    await litellm.proxy.proxy_server.prisma_client.connect()
+    try:
+        response = await generate_key_helper_fn(
+            **{"duration": "1hr", "key_max_budget": 0, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": "ishaan", "team_id": "litellm-dashboard"}  # type: ignore
+        )
+
+        print(response)
+        key = response["token"]
+
+        # make a /chat/completions call -> it should fail
+        request = Request(scope={"type": "http"})
+        request._url = URL(url="/chat/completions")
+
+        # use generated key to auth in
+        result = await user_api_key_auth(request=request, api_key="Bearer " + key)
+        print("result from user auth with new key", result)
+        pytest.fail(f"This should have failed!. IT's an invalid key")
+    except Exception as e:
+        print("Got Exception", e)
+        print(e.message)