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

Merge pull request #4357 from BerriAI/litellm_use_jwt_in_cookies

Browse source 
[Security Fix - Proxy Server ADMIN UI] - Store credentials in cookies + use strong JWT signing secret
This commit is contained in:
Ishaan Jaff 2024-06-22 09:12:17 -07:00 • committed by GitHub
commit d32b8d589b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 74 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

54
litellm/proxy/proxy_server.py
View file

@ -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
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
@ -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
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}",
@ -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
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 ####

11
ui/litellm-dashboard/src/app/onboarding/page.tsx
View file

@ -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>("");

11
ui/litellm-dashboard/src/app/page.tsx
View file

@ -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);

11
ui/litellm-dashboard/src/components/user_dashboard.tsx
View file

@ -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[]>([]);