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

fix(proxy_server.py): secure /team/info endpoint

Browse source 
make sure user requesting team info is part of team or admin
This commit is contained in:
Krrish Dholakia 2024-04-16 11:39:52 -07:00
parent 77df5d7658
commit 3f8abe2754
2 changed files with 39 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

9
litellm/proxy/proxy_server.py
View file

@ -1040,6 +1040,15 @@ async def user_api_key_auth(
elif route == "/model/info":
# /model/info just shows models user has access to
pass
elif route == "/team/info":
# check if key can access this team's info
query_params = request.query_params
team_id = query_params.get("team_id")
if team_id != valid_token.team_id:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="key not allowed to access this team's info",
)
else:
raise Exception(
f"Only master key can be used to generate, delete, update info for new keys/users."

30
tests/test_team.py
View file

@ -260,7 +260,18 @@ async def get_team_info(session, get_team, call_key):
@pytest.mark.asyncio
async def test_team_info():
"""
Scenario 1:
- test with admin key -> expect to work
Scenario 2:
- test with team key -> expect to work
Scenario 3:
- test with non-team key -> expect to fail
"""
async with aiohttp.ClientSession() as session:
"""
Scenario 1 - as admin
"""
new_team_data = await new_team(
session,
0,
@ -268,6 +279,25 @@ async def test_team_info():
team_id = new_team_data["team_id"]
## as admin ##
await get_team_info(session=session, get_team=team_id, call_key="sk-1234")
"""
Scenario 2 - as team key
"""
key_gen = await generate_key(session=session, i=0, team_id=team_id)
key = key_gen["key"]
await get_team_info(session=session, get_team=team_id, call_key=key)
"""
Scenario 3 - as non-team key
"""
key_gen = await generate_key(session=session, i=0)
key = key_gen["key"]
try:
await get_team_info(session=session, get_team=team_id, call_key=key)
pytest.fail(f"Expected call to fail")
except Exception as e:
pass
"""