From 204f7725ee75b1749a9d526ce9a81e95f1702b39 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 22 Jun 2024 08:47:43 -0700 Subject: [PATCH 1/4] ui - read jwts from cookie --- .../src/components/user_dashboard.tsx | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index 5aa0c92dd53..ec66bcea0fa 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/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 = ({ const viewSpend = searchParams.get("viewSpend"); const router = useRouter(); - const token = searchParams.get("token"); + const token = getCookie('token'); + const [accessToken, setAccessToken] = useState(null); const [teamSpend, setTeamSpend] = useState(null); const [userModels, setUserModels] = useState([]); From 03d1d9229d1bca36a410ded26e338c540e0daa73 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 22 Jun 2024 08:50:26 -0700 Subject: [PATCH 2/4] ui - use cookies to return JWTs --- litellm/proxy/proxy_server.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 021b59e295f..f3aeaddf983 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7514,8 +7514,10 @@ async def login(request: Request): "secret", algorithm="HS256", ) - litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwt_token - return RedirectResponse(url=litellm_dashboard_ui, status_code=303) + litellm_dashboard_ui += "?userID=" + user_id + 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 @@ -7576,8 +7578,12 @@ async def login(request: Request): "secret", algorithm="HS256", ) - litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwt_token - return RedirectResponse(url=litellm_dashboard_ui, status_code=303) + litellm_dashboard_ui += "?userID=" + user_id + 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}", @@ -8120,8 +8126,10 @@ async def auth_callback(request: Request): "secret", algorithm="HS256", ) - litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwt_token - return RedirectResponse(url=litellm_dashboard_ui) + litellm_dashboard_ui += "?userID=" + user_id + redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) + redirect_response.set_cookie(key="token", value=jwt_token) + return redirect_response #### INVITATION MANAGEMENT #### From c86efa55a5620cef8fbc243919db596a338c662d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 22 Jun 2024 08:58:31 -0700 Subject: [PATCH 3/4] use hash of master key encode the jwt --- litellm/proxy/proxy_server.py | 34 ++++++++++++++++++++++++++++++---- 1 file changed, 30 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f3aeaddf983..630aa3f3e15 100644 --- a/litellm/proxy/proxy_server.py +++ b/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,7 +7517,7 @@ 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 @@ -7566,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, @@ -7575,7 +7589,7 @@ 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 @@ -7694,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, @@ -7703,7 +7723,7 @@ async def onboarding(invite_link: str): "login_method": "username_password", "premium_user": premium_user, }, - "secret", + litellm_master_key_hash, algorithm="HS256", ) @@ -8114,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, @@ -8123,7 +8149,7 @@ 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 From 066ed01d0f1894f0d171334331568306edf7a7a4 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 22 Jun 2024 09:05:06 -0700 Subject: [PATCH 4/4] read token from cookie --- ui/litellm-dashboard/src/app/onboarding/page.tsx | 11 ++++++++++- ui/litellm-dashboard/src/app/page.tsx | 11 ++++++++++- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/app/onboarding/page.tsx b/ui/litellm-dashboard/src/app/onboarding/page.tsx index 44c37610cec..a31afb0f9e5 100644 --- a/ui/litellm-dashboard/src/app/onboarding/page.tsx +++ b/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(null); const [defaultUserEmail, setDefaultUserEmail] = useState(""); diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index ed9c98df7c6..41f0e8529b3 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/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({ 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(null);