Merge pull request #4357 from BerriAI/litellm_use_jwt_in_cookies

[Security Fix - Proxy Server ADMIN UI] - Store credentials in cookies + use strong JWT signing secret

litellm/proxy/proxy_server.py

@@ -7502,6 +7502,12 @@ async def login(request: Request):
             litellm_dashboard_ui += "/ui/"
         import jwt
 
+        if litellm_master_key_hash is None:
+            raise HTTPException(
+                status_code=500,
+                detail={"error": "No master key set, please set LITELLM_MASTER_KEY"},
+            )
+
         jwt_token = jwt.encode(
             {
                 "user_id": user_id,
@@ -7511,11 +7517,13 @@ async def login(request: Request):
                 "login_method": "username_password",
                 "premium_user": premium_user,
             },
-            "secret",
+            litellm_master_key_hash,
             algorithm="HS256",
         )
-        litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwt_token
+        litellm_dashboard_ui += "?userID=" + user_id
-        return RedirectResponse(url=litellm_dashboard_ui, status_code=303)
+        redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
+        redirect_response.set_cookie(key="token", value=jwt_token)
+        return redirect_response
     elif _user_row is not None:
         """
         When sharing invite links
@@ -7564,6 +7572,14 @@ async def login(request: Request):
                 litellm_dashboard_ui += "/ui/"
             import jwt
 
+            if litellm_master_key_hash is None:
+                raise HTTPException(
+                    status_code=500,
+                    detail={
+                        "error": "No master key set, please set LITELLM_MASTER_KEY"
+                    },
+                )
+
             jwt_token = jwt.encode(
                 {
                     "user_id": user_id,
@@ -7573,11 +7589,15 @@ async def login(request: Request):
                     "login_method": "username_password",
                     "premium_user": premium_user,
                 },
-                "secret",
+                litellm_master_key_hash,
                 algorithm="HS256",
             )
-            litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwt_token
+            litellm_dashboard_ui += "?userID=" + user_id
-            return RedirectResponse(url=litellm_dashboard_ui, status_code=303)
+            redirect_response = RedirectResponse(
+                url=litellm_dashboard_ui, status_code=303
+            )
+            redirect_response.set_cookie(key="token", value=jwt_token)
+            return redirect_response
         else:
             raise ProxyException(
                 message=f"Invalid credentials used to access UI. Passed in username: {username}, passed in password: {password}.\nNot valid credentials for {username}",
@@ -7688,6 +7708,12 @@ async def onboarding(invite_link: str):
         litellm_dashboard_ui += "/ui/onboarding"
     import jwt
 
+    if litellm_master_key_hash is None:
+        raise HTTPException(
+            status_code=500,
+            detail={"error": "No master key set, please set LITELLM_MASTER_KEY"},
+        )
+
     jwt_token = jwt.encode(
         {
             "user_id": user_obj.user_id,
@@ -7697,7 +7723,7 @@ async def onboarding(invite_link: str):
             "login_method": "username_password",
             "premium_user": premium_user,
         },
-        "secret",
+        litellm_master_key_hash,
         algorithm="HS256",
     )
 
@@ -8108,6 +8134,12 @@ async def auth_callback(request: Request):
 
     import jwt
 
+    if litellm_master_key_hash is None:
+        raise HTTPException(
+            status_code=500,
+            detail={"error": "No master key set, please set LITELLM_MASTER_KEY"},
+        )
+
     jwt_token = jwt.encode(
         {
             "user_id": user_id,
@@ -8117,11 +8149,13 @@ async def auth_callback(request: Request):
             "login_method": "sso",
             "premium_user": premium_user,
         },
-        "secret",
+        litellm_master_key_hash,
         algorithm="HS256",
     )
-    litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwt_token
+    litellm_dashboard_ui += "?userID=" + user_id
-    return RedirectResponse(url=litellm_dashboard_ui)
+    redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
+    redirect_response.set_cookie(key="token", value=jwt_token)
+    return redirect_response
 
 
 #### INVITATION MANAGEMENT ####

ui/litellm-dashboard/src/app/onboarding/page.tsx

@@ -20,10 +20,19 @@ import {
 } from "@/components/networking";
 import { jwtDecode } from "jwt-decode";
 import { Form, Button as Button2, message } from "antd";
+
+function getCookie(name: string) {
+  console.log("COOKIES", document.cookie)
+  const cookieValue = document.cookie
+      .split('; ')
+      .find(row => row.startsWith(name + '='));
+  return cookieValue ? cookieValue.split('=')[1] : null;
+}
+
 export default function Onboarding() {
   const [form] = Form.useForm();
   const searchParams = useSearchParams();
-  const token = searchParams.get("token");
+  const token = getCookie('token');
   const inviteID = searchParams.get("id");
   const [accessToken, setAccessToken] = useState<string | null>(null);
   const [defaultUserEmail, setDefaultUserEmail] = useState<string>("");

ui/litellm-dashboard/src/app/page.tsx

@@ -19,6 +19,15 @@ import CacheDashboard from "@/components/cache_dashboard";
 import { jwtDecode } from "jwt-decode";
 import { Typography } from "antd";
 
+function getCookie(name: string) {
+  console.log("COOKIES", document.cookie)
+  const cookieValue = document.cookie
+      .split('; ')
+      .find(row => row.startsWith(name + '='));
+  return cookieValue ? cookieValue.split('=')[1] : null;
+}
+
+
 function formatUserRole(userRole: string) {
   if (!userRole) {
     return "Undefined Role";
@@ -68,7 +77,7 @@ const CreateKeyPage = () => {
   const searchParams = useSearchParams();
   const [modelData, setModelData] = useState<any>({ data: [] });
   const userID = searchParams.get("userID");
-  const token = searchParams.get("token");
+  const token = getCookie('token');
 
   const [page, setPage] = useState("api-keys");
   const [accessToken, setAccessToken] = useState<string | null>(null);

ui/litellm-dashboard/src/components/user_dashboard.tsx

@@ -24,6 +24,14 @@ type UserSpendData = {
   max_budget?: number | null;
 };
 
+function getCookie(name: string) {
+  console.log("COOKIES", document.cookie)
+  const cookieValue = document.cookie
+      .split('; ')
+      .find(row => row.startsWith(name + '='));
+  return cookieValue ? cookieValue.split('=')[1] : null;
+}
+
 interface UserDashboardProps {
   userID: string | null;
   userRole: string | null;
@@ -66,7 +74,8 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
   const viewSpend = searchParams.get("viewSpend");
   const router = useRouter();
 
-  const token = searchParams.get("token");
+  const token = getCookie('token');
+
   const [accessToken, setAccessToken] = useState<string | null>(null);
   const [teamSpend, setTeamSpend] = useState<number | null>(null);
   const [userModels, setUserModels] = useState<string[]>([]);