mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
Merge pull request #42064 from BerriAI/litellm_fix_jwt_deactivated_users_8217
fix(auth): reject deactivated JWT users and refresh cached status
This commit is contained in:
joshua-berri
GitHub
commit
c25601fb5f
8 changed files with 250 additions and 10 deletions
|
|
@ -2468,7 +2468,22 @@ class JWTAuthManager:
|
|||
jwt_valid_token, handler, prisma_client, user_api_key_cache, parent_otel_span, proxy_logging_obj
|
||||
)
|
||||
return {**admin_result, "user_object": identity.user_object}
|
||||
return admin_result
|
||||
if prisma_client is None:
|
||||
return admin_result
|
||||
try:
|
||||
admin_user: Final = await get_user_object(
|
||||
user_id=user_id,
|
||||
user_email=user_email,
|
||||
sso_user_id=user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except UserNotFoundError:
|
||||
return admin_result
|
||||
return {**admin_result, "user_object": admin_user}
|
||||
|
||||
# Get team with model access
|
||||
## Check if team_id is specified via x-litellm-team-id header
|
||||
|
|
|
|||
|
|
@ -1714,6 +1714,16 @@ async def _user_api_key_auth_builder(
|
|||
jwt_claims = result.get("jwt_claims", None)
|
||||
agent_id: Final[str | None] = result.get("agent_id")
|
||||
|
||||
if (
|
||||
user_object is not None
|
||||
and isinstance(user_object.metadata, dict)
|
||||
and user_object.metadata.get("scim_active") is False
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=f"User={user_id} has been deactivated via SCIM. Keys owned by this user cannot be used.",
|
||||
)
|
||||
|
||||
if is_proxy_admin:
|
||||
# Proxy admins authenticate via auth_builder (full
|
||||
# access), not via a mapped virtual key. If
|
||||
|
|
|
|||
|
|
@ -1571,7 +1571,7 @@ async def _update_single_user_helper(
|
|||
|
||||
await _invalidate_user_spend_counter_if_changed(non_default_values)
|
||||
|
||||
if "model_max_budget" in non_default_values:
|
||||
if "model_max_budget" in non_default_values or "metadata" in data_json:
|
||||
await evict_and_broadcast(
|
||||
cache_keys=(non_default_values["user_id"],),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.proxy.auth.auth_checks import _delete_cache_key_object
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
|
||||
from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
|
||||
from litellm.proxy.management_endpoints.scim.scim_transformations import (
|
||||
|
|
@ -1804,6 +1805,9 @@ async def update_user(
|
|||
where={"user_id": user_id},
|
||||
data=update_data,
|
||||
)
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
await evict_and_broadcast(cache_keys=(user_id,), user_api_key_cache=user_api_key_cache)
|
||||
|
||||
if client_set_active:
|
||||
new_active: Final = _scim_active_value(metadata)
|
||||
|
|
@ -2375,6 +2379,9 @@ async def patch_user(
|
|||
where={"user_id": user_id},
|
||||
data=update_data,
|
||||
)
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
await evict_and_broadcast(cache_keys=(user_id,), user_api_key_cache=user_api_key_cache)
|
||||
|
||||
if new_active is not None and new_active != (True if prev_active is None else prev_active):
|
||||
await _set_user_keys_blocked(user_id=user_id, blocked=not new_active)
|
||||
|
|
|
|||
|
|
@ -7144,3 +7144,52 @@ async def test_admin_jwt_team_header_only_provisions_during_admission(monkeypatc
|
|||
else:
|
||||
create_team.assert_not_awaited()
|
||||
assert result["team_id"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("existing_user", [False, True])
|
||||
@pytest.mark.parametrize("warm_cache", [False, True])
|
||||
@pytest.mark.parametrize("email", [None, "admin@external.example", "admin@allowed.example"])
|
||||
async def test_scope_admin_admission_resolves_existing_user_without_provisioning(
|
||||
monkeypatch: pytest.MonkeyPatch, existing_user: bool, warm_cache: bool, email: str | None
|
||||
) -> None:
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
private_key, jwk = _get_rsa_key_and_jwk("admin-status")
|
||||
cache: Final = UserApiKeyCache()
|
||||
cache.set_cache("litellm_jwt_auth_keys_https://admin.example/jwks", [jwk])
|
||||
user_id: Final = f"admin-status-{existing_user}-{warm_cache}-{email}"
|
||||
user: Final = LiteLLM_UserTable(user_id=user_id, user_email="admin@allowed.example", metadata={"scim_active": False}, organization_memberships=[])
|
||||
if existing_user and warm_cache:
|
||||
cache.set_cache(user_id, user)
|
||||
database: Final = MagicMock()
|
||||
users: Final = database.db.litellm_usertable
|
||||
users.find_unique = AsyncMock(return_value=user if existing_user else None)
|
||||
users.find_first = AsyncMock(return_value=None)
|
||||
users.create = AsyncMock()
|
||||
handler: Final = JWTHandler()
|
||||
handler.update_environment(
|
||||
prisma_client=database,
|
||||
user_api_key_cache=cache,
|
||||
litellm_jwtauth=LiteLLM_JWTAuth(
|
||||
user_id_jwt_field="sub", user_id_upsert=True, user_email_jwt_field="email",
|
||||
user_allowed_email_domain="allowed.example",
|
||||
),
|
||||
)
|
||||
monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "https://admin.example/jwks")
|
||||
monkeypatch.setenv("JWT_ISSUER", "https://admin.example")
|
||||
monkeypatch.setenv("JWT_AUDIENCE", "gateway")
|
||||
token: Final = _encode_rsa_jwt(
|
||||
private_key, "https://admin.example", "gateway", "admin-status",
|
||||
{"sub": user_id, "scope": "litellm_proxy_admin", **({"email": email} if email else {})},
|
||||
)
|
||||
result: Final = await JWTAuthManager.auth_builder(
|
||||
api_key=token, jwt_handler=handler, prisma_client=database, user_api_key_cache=cache,
|
||||
parent_otel_span=None, proxy_logging_obj=MagicMock(), request_data={}, general_settings={}, route="/user/info",
|
||||
)
|
||||
assert result["is_proxy_admin"] is True
|
||||
assert result["user_id"] == user_id
|
||||
assert result["user_object"] == (user if existing_user else None)
|
||||
users.create.assert_not_awaited()
|
||||
if existing_user:
|
||||
assert users.find_unique.await_count == (0 if warm_cache else 1)
|
||||
|
|
|
|||
|
|
@ -2091,7 +2091,8 @@ async def test_auto_register_binds_api_key_to_token_hash():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_register_first_request_propagates_user_email():
|
||||
@pytest.mark.parametrize("active", [True, False])
|
||||
async def test_auto_register_first_request_propagates_user_email(active: bool) -> None:
|
||||
"""
|
||||
The first auto-registered JWT request must also carry user_email (resolved
|
||||
from the validated LiteLLM_UserTable), so attribution is consistent with the
|
||||
|
|
@ -2120,6 +2121,7 @@ async def test_auto_register_first_request_propagates_user_email():
|
|||
user_id="validated-user",
|
||||
user_email="validated@example.com",
|
||||
user_role="internal_user",
|
||||
metadata={"scim_active": active},
|
||||
)
|
||||
mock_jwt_result = {
|
||||
"is_proxy_admin": False,
|
||||
|
|
@ -2150,7 +2152,7 @@ async def test_auto_register_first_request_propagates_user_email():
|
|||
patch("litellm.proxy.proxy_server.master_key", "sk-master"),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache),
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock(return_value=None))),
|
||||
patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler),
|
||||
patch(
|
||||
"litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key",
|
||||
|
|
@ -2170,8 +2172,22 @@ async def test_auto_register_first_request_propagates_user_email():
|
|||
"litellm.proxy.auth.user_api_key_auth._auto_register_jwt_mapping",
|
||||
new_callable=AsyncMock,
|
||||
return_value=auto_registered_key,
|
||||
),
|
||||
) as auto_register,
|
||||
):
|
||||
if not active:
|
||||
with pytest.raises(ProxyException, match="deactivated via SCIM") as exc:
|
||||
await _user_api_key_auth_builder(
|
||||
request=mock_request,
|
||||
api_key=jwt_token,
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
request_data={},
|
||||
)
|
||||
assert int(exc.value.code) == 401
|
||||
auto_register.assert_not_awaited()
|
||||
return
|
||||
result = await _user_api_key_auth_builder(
|
||||
request=mock_request,
|
||||
api_key=jwt_token,
|
||||
|
|
@ -7315,15 +7331,15 @@ class TestJWTAuthUserEmail:
|
|||
the Prometheus `user_email` label and `user_api_key_user_email` in
|
||||
StandardLogging/SpendLogs metadata, which were always None for JWT traffic."""
|
||||
|
||||
def _jwt_request(self, jwt_token):
|
||||
def _jwt_request(self, jwt_token, route="/v1/chat/completions"):
|
||||
mock_request = MagicMock()
|
||||
mock_request.url.path = "/v1/chat/completions"
|
||||
mock_request.method = "POST"
|
||||
mock_request.url.path = route
|
||||
mock_request.method = "GET" if route.endswith("/list") else "POST"
|
||||
mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
|
||||
mock_request.query_params = {}
|
||||
return mock_request
|
||||
|
||||
async def _run_jwt_auth(self, mock_jwt_result, jwt_token):
|
||||
async def _run_jwt_auth(self, mock_jwt_result, jwt_token, route="/v1/chat/completions"):
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
|
|
@ -7344,7 +7360,7 @@ class TestJWTAuthUserEmail:
|
|||
litellm_jwtauth=LiteLLM_JWTAuth(),
|
||||
)
|
||||
return await user_api_key_auth(
|
||||
request=self._jwt_request(jwt_token),
|
||||
request=self._jwt_request(jwt_token, route),
|
||||
api_key=f"Bearer {jwt_token}",
|
||||
)
|
||||
|
||||
|
|
@ -7376,6 +7392,44 @@ class TestJWTAuthUserEmail:
|
|||
assert result.user_id == "jwt-human-user"
|
||||
assert result.user_email == "resolved@example.com"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", ["/mcp-rest/tools/list", "/mcp-rest/tools/call", "/v1/chat/completions", "/user/info"])
|
||||
@pytest.mark.parametrize("active", [False, True, None, "false", 0])
|
||||
@pytest.mark.parametrize("is_admin", [False, True])
|
||||
async def test_jwt_auth_rejects_deactivated_user(
|
||||
self, route: str, active: bool | str | int | None, is_admin: bool
|
||||
) -> None:
|
||||
from typing import Final
|
||||
|
||||
jwt_token: Final = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature"
|
||||
result: Final = {
|
||||
"is_proxy_admin": is_admin,
|
||||
"team_object": None,
|
||||
"user_object": LiteLLM_UserTable(
|
||||
user_id="jwt-human-user",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN.value if is_admin else LitellmUserRoles.INTERNAL_USER.value,
|
||||
metadata={} if active is None else {"scim_active": active},
|
||||
),
|
||||
"end_user_object": None,
|
||||
"org_object": None,
|
||||
"token": jwt_token,
|
||||
"team_id": None,
|
||||
"user_id": "jwt-human-user",
|
||||
"user_email": None,
|
||||
"end_user_id": None,
|
||||
"org_id": None,
|
||||
"team_membership": None,
|
||||
"jwt_claims": {"sub": "user1"},
|
||||
}
|
||||
|
||||
if active is False:
|
||||
with pytest.raises(ProxyException, match="deactivated via SCIM") as exc:
|
||||
await self._run_jwt_auth(result, jwt_token, route)
|
||||
assert int(exc.value.code) == 401
|
||||
else:
|
||||
token: Final = await self._run_jwt_auth(result, jwt_token, route)
|
||||
assert token.user_id == "jwt-human-user"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_jwt_auth_populates_user_email_on_proxy_admin(self):
|
||||
jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature"
|
||||
|
|
|
|||
|
|
@ -541,3 +541,63 @@ async def test_scim_put_user_explicit_active_false_blocks_keys():
|
|||
assert update_kwargs["where"] == {"token": "hash-block-me"}
|
||||
assert update_kwargs["data"]["blocked"] is True
|
||||
assert '"scim_blocked": true' in update_kwargs["data"]["metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("method", ["PUT", "PATCH"])
|
||||
@pytest.mark.parametrize("active", [False, True])
|
||||
@pytest.mark.parametrize("failure", [None, "write", "keys"])
|
||||
@pytest.mark.parametrize("status_change", [False, True])
|
||||
async def test_scim_status_write_refreshes_user_cache(
|
||||
method: str, active: bool, failure: str | None, status_change: bool
|
||||
) -> None:
|
||||
import json
|
||||
from typing import Final
|
||||
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
user_id: Final = "scim-cache-user"
|
||||
saved: Final = LiteLLM_UserTable(
|
||||
user_id=user_id, user_email="x@example.com", teams=[], metadata={"scim_active": not active if status_change else active},
|
||||
)
|
||||
updated: Final = LiteLLM_UserTable(
|
||||
user_id=user_id, user_email="x@example.com", teams=[], metadata={"scim_active": active},
|
||||
)
|
||||
client, db = _build_prisma_with_keys([], mock_user=saved.model_copy(deep=True), updated_user=updated)
|
||||
if failure == "write":
|
||||
db.litellm_usertable.update.side_effect = RuntimeError("status write failed")
|
||||
if failure == "keys":
|
||||
db.litellm_verificationtoken.find_many.side_effect = RuntimeError("key update failed")
|
||||
cache: Final = UserApiKeyCache()
|
||||
await cache.async_set_cache(key=user_id, value=saved, model_type=LiteLLM_UserTable)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", client), # test-quality-ok: substitute the database dependency
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: exercise a real isolated cache
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), # test-quality-ok: isolate the logging dependency
|
||||
patch( # test-quality-ok: observe the Redis publication boundary
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
|
||||
new_callable=AsyncMock,
|
||||
) as broadcast,
|
||||
):
|
||||
request: Final = (
|
||||
update_user(user_id=user_id, user=SCIMUser.model_validate(_build_put_user_payload(user_id, active=active)))
|
||||
if method == "PUT" else
|
||||
patch_user(user_id=user_id, patch_ops=SCIMPatchOp(
|
||||
Operations=[SCIMPatchOperation(op="replace", path="active", value=active)]
|
||||
))
|
||||
)
|
||||
if failure == "write" or (failure == "keys" and status_change):
|
||||
with pytest.raises(ProxyException, match="status write failed" if failure == "write" else "key update failed"):
|
||||
await request
|
||||
else:
|
||||
response: Final = await request
|
||||
assert response.active is active
|
||||
assert json.loads(db.litellm_usertable.update.await_args.kwargs["data"]["metadata"])["scim_active"] is active
|
||||
cached: Final = await cache.async_get_cache(key=user_id, model_type=LiteLLM_UserTable)
|
||||
if failure == "write":
|
||||
assert cached == saved
|
||||
broadcast.assert_not_awaited()
|
||||
else:
|
||||
assert cached is None
|
||||
broadcast.assert_awaited_once_with(cache_key=user_id)
|
||||
|
|
|
|||
|
|
@ -2228,6 +2228,51 @@ async def test_user_model_budget_update_by_email_refreshes_cached_user(mocker: M
|
|||
broadcast.assert_awaited_once_with(cache_key=saved_user.user_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("by_email", [False, True])
|
||||
@pytest.mark.parametrize("active", [False, True, None])
|
||||
async def test_user_status_update_refreshes_cached_user(
|
||||
mocker: MockerFixture, by_email: bool, active: bool | None
|
||||
) -> None:
|
||||
from litellm.proxy._types import LiteLLM_UserTable
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper
|
||||
|
||||
saved_user: Final = LiteLLM_UserTable(
|
||||
user_id="user-spruce",
|
||||
user_email="spruce@example.test",
|
||||
metadata={"scim_active": False if active is None else not active, "department": "engineering"},
|
||||
)
|
||||
prisma_client: Final = mocker.MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user)
|
||||
prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user])
|
||||
prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user})
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
|
||||
cache: Final = UserApiKeyCache()
|
||||
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
|
||||
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
|
||||
new_callable=mocker.AsyncMock,
|
||||
)
|
||||
|
||||
await _update_single_user_helper(
|
||||
user_request=UpdateUserRequest(
|
||||
user_id=None if by_email else saved_user.user_id,
|
||||
user_email=saved_user.user_email if by_email else None,
|
||||
metadata={"department": "engineering"} if active is None else {"scim_active": active},
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
assert prisma_client.update_data.call_args.kwargs["user_id"] == saved_user.user_id
|
||||
assert prisma_client.update_data.call_args.kwargs["data"]["metadata"] == (
|
||||
{"department": "engineering"} if active is None else {"scim_active": active}
|
||||
)
|
||||
assert await cache.async_get_cache(key=saved_user.user_id, model_type=LiteLLM_UserTable) is None
|
||||
broadcast.assert_awaited_once_with(cache_key=saved_user.user_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocker: MockerFixture) -> None:
|
||||
from litellm.proxy._types import LiteLLM_UserTable
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue