diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 1293a05e4a6..cf0539ec18f 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -15424,14 +15424,31 @@ "title": "Jwt Issuer" }, "key": { - "title": "Key", - "type": "string" + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key" + }, + "token": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token" } }, "required": [ "jwt_claim_name", - "jwt_claim_value", - "key" + "jwt_claim_value" ], "title": "CreateJWTKeyMappingRequest", "type": "object" @@ -15595,6 +15612,17 @@ } ], "title": "Key" + }, + "token": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Token" } }, "required": [ diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3f2d2f57fd4..70e49bb18b9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4733,7 +4733,8 @@ class KeyHealthResponse(TypedDict, total=False): class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase): jwt_claim_name: str jwt_claim_value: str - key: str + key: str | None = None + token: str | None = None jwt_issuer: str | None = None description: str | None = None @@ -4741,6 +4742,7 @@ class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase): class UpdateJWTKeyMappingRequest(LiteLLMPydanticObjectBase): id: str key: str | None = None + token: str | None = None jwt_issuer: str | None = None description: str | None = None is_active: bool | None = None diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 07234883062..292cec1346d 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -1,3 +1,4 @@ +import re from collections.abc import Mapping, Sequence from datetime import datetime from typing import Final, Protocol @@ -21,6 +22,51 @@ from litellm.repositories.table_repositories import JWTKeyMappingRepository router: Final = APIRouter() +_TOKEN_HASH_PATTERN: Final = re.compile(r"[0-9a-f]{64}") + + +def _validated_token_hash(token: str) -> str: + """Guards a plaintext key from being stored as a hash of a hash, which would never match.""" + if _TOKEN_HASH_PATTERN.fullmatch(token) is None: + raise HTTPException( + status_code=400, + detail=( + "`token` must be the SHA-256 hash of a virtual key " + "(64 lowercase hex characters). Pass the plaintext as `key` instead." + ), + ) + return token + + +_EXACTLY_ONE_IDENTIFIER: Final = ( + "Provide exactly one of `key` (the plaintext virtual key) or `token` (its SHA-256 hash)." +) +_AT_MOST_ONE_IDENTIFIER: Final = ( + "Provide at most one of `key` (the plaintext virtual key) or `token` (its SHA-256 hash)." +) + + +def _token_hash_for_create(data: CreateJWTKeyMappingRequest) -> str: + """Resolve the token hash to store, from either the plaintext key or its hash.""" + if data.key is not None and data.token is not None: + raise HTTPException(status_code=400, detail=_EXACTLY_ONE_IDENTIFIER) + if data.token is not None: + return _validated_token_hash(data.token) + if data.key is not None: + return hash_token(data.key) + raise HTTPException(status_code=400, detail=_EXACTLY_ONE_IDENTIFIER) + + +def _token_hash_for_update(data: UpdateJWTKeyMappingRequest) -> str | None: + """Resolve the token hash to store, or None to leave the mapped key alone.""" + if data.key is not None and data.token is not None: + raise HTTPException(status_code=400, detail=_AT_MOST_ONE_IDENTIFIER) + if data.token is not None: + return _validated_token_hash(data.token) + if data.key is not None: + return hash_token(data.key) + return None + class _JWTKeyMappingRecord(Protocol): """A ``LiteLLM_JWTKeyMapping`` row, viewed through the columns these endpoints read.""" @@ -111,7 +157,7 @@ async def create_jwt_key_mapping( raise HTTPException(status_code=500, detail="Database not connected") try: - hashed_key: Final = hash_token(data.key) + hashed_key: Final = _token_hash_for_create(data) create_data: Final = { "jwt_issuer": data.jwt_issuer or "", "jwt_claim_name": data.jwt_claim_name, @@ -166,9 +212,10 @@ async def update_jwt_key_mapping( if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") - update_data: Final = data.model_dump(exclude_unset=True, exclude={"id", "key"}) - if data.key is not None: - update_data["token"] = hash_token(data.key) + update_data: Final = data.model_dump(exclude_unset=True, exclude={"id", "key", "token"}) + token_hash: Final = _token_hash_for_update(data) + if token_hash is not None: + update_data["token"] = token_hash if "jwt_issuer" in update_data: # DB column is NOT NULL (see schema.prisma); "" is the global/unscoped sentinel. update_data["jwt_issuer"] = update_data["jwt_issuer"] or "" diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index ee5b42fe0b7..8f40ed6dfb7 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -18,6 +18,15 @@ longer signal it. - **team_member_add**: `tpm_limit`, `rpm_limit`, `budget_duration`, and `allowed_models` attributes on `litellm_team_member_add`, applied to every member of the resource; `budget_duration` and `allowed_models` ride on `/team/member_add`, while the limits are sent through `/team/member_update`, which is where the proxy accepts them - **team**: Optional `team_id` argument on `litellm_team`, so teams can be created with a stable, human-readable ID instead of a provider-generated UUID; changing it forces replacement +- `litellm_jwt_key_mapping` accepts `token_id` as an alternative to `key`, so a + mapping can name its virtual key by the SHA-256 hash the proxy stores instead + of by the plaintext. Exactly one of the two is required. This is what lets a + mapping reference a key managed in the same configuration + (`token_id = litellm_key.foo.token_id`), which `key` cannot do, because + `litellm_key` marks its generated key write-only and referencing it fails at + plan time. `POST /jwt/key/mapping/new` and `/jwt/key/mapping/update` gained a + matching `token` field, validated as 64 lowercase hex characters so a + plaintext key sent by mistake is rejected instead of hashed twice - **jwt_key_mapping**: New `litellm_jwt_key_mapping` resource for the proxy's JWT to virtual key mappings, so JWT clients identified by a claim (`client_id`, `azp`, `sub`) map to virtual keys and inherit their models, budgets and rate limits. Supports `description` and `is_active`, rotating the mapped key in place, and forces replacement when the claim name or value changes - **team**: `soft_budget`, `tags`, and `soft_budget_alerting_emails` attributes on `litellm_team`, matching what `/team/new` and `/team/update` already accept; `soft_budget_alerting_emails` is sent under `metadata`, where the proxy reads it - **user**: New `litellm_user` resource and `litellm_user` / `litellm_users` data sources for managing internal users diff --git a/terraform/provider/docs/resources/jwt_key_mapping.md b/terraform/provider/docs/resources/jwt_key_mapping.md index fbc30947113..726c4b16021 100644 --- a/terraform/provider/docs/resources/jwt_key_mapping.md +++ b/terraform/provider/docs/resources/jwt_key_mapping.md @@ -65,7 +65,8 @@ resource "litellm_jwt_key_mapping" "developer" { - `jwt_claim_name` - (Required, ForceNew) Name of the JWT claim to match on, for example `client_id`, `azp` or `sub`. Must match `virtual_key_claim_field` in the proxy JWT config - `jwt_claim_value` - (Required, ForceNew) Value of the claim identifying the JWT client. Unique together with `jwt_claim_name`, so a second mapping for the same pair fails with a 409 -- `key` - (Required, Sensitive) The virtual key this claim value maps to. It has to exist already, otherwise the proxy rejects the mapping with `The provided key does not match an existing virtual key` +- `key` - (Optional, Sensitive) The virtual key this claim value maps to, as plaintext. It has to exist already, otherwise the proxy rejects the mapping with `The provided key does not match an existing virtual key`. Exactly one of `key` or `token_id` is required. `litellm_key` marks its generated `key` write-only, so this cannot reference a `litellm_key` resource -- use `token_id` for that, or supply the plaintext from a variable or a secret manager +- `token_id` - (Optional) The SHA-256 hash of the virtual key this claim value maps to, which is what the proxy stores. `litellm_key` exposes it as `token_id`, so unlike `key` it can be referenced directly from a `litellm_key` resource. Not a secret, so it is not marked sensitive. Exactly one of `key` or `token_id` is required - `description` - (Optional) Description of the mapping - `is_active` - (Optional) Whether the mapping is active. Inactive mappings are ignored during JWT auth. Defaults to `true` diff --git a/terraform/provider/litellm/resource_jwt_key_mapping.go b/terraform/provider/litellm/resource_jwt_key_mapping.go index e606e865737..ea968821527 100644 --- a/terraform/provider/litellm/resource_jwt_key_mapping.go +++ b/terraform/provider/litellm/resource_jwt_key_mapping.go @@ -29,10 +29,17 @@ func resourceLiteLLMJWTKeyMapping() *schema.Resource { Description: "Value of the claim identifying the JWT client. Unique together with jwt_claim_name", }, "key": { - Type: schema.TypeString, - Required: true, - Sensitive: true, - Description: "The virtual key this claim value maps to. The proxy stores only a hash of it and never returns it, so drift on this attribute cannot be detected and Terraform tracks the configured value", + Type: schema.TypeString, + Optional: true, + Sensitive: true, + ExactlyOneOf: []string{"key", "token_id"}, + Description: "The virtual key this claim value maps to, as plaintext. The proxy stores only a hash of it and never returns it, so drift on this attribute cannot be detected and Terraform tracks the configured value. litellm_key marks its generated key write-only, so this cannot reference a litellm_key resource; use token_id for that, or supply the plaintext from a variable or a secret manager", + }, + "token_id": { + Type: schema.TypeString, + Optional: true, + ExactlyOneOf: []string{"key", "token_id"}, + Description: "The SHA-256 hash of the virtual key this claim value maps to, which is what the proxy stores. litellm_key exposes it as token_id, so unlike key it can be referenced directly from a litellm_key resource. Not a secret, so it is not marked sensitive", }, "description": { Type: schema.TypeString, diff --git a/terraform/provider/litellm/resource_jwt_key_mapping_crud.go b/terraform/provider/litellm/resource_jwt_key_mapping_crud.go index 725235305f6..6c3883de2e1 100644 --- a/terraform/provider/litellm/resource_jwt_key_mapping_crud.go +++ b/terraform/provider/litellm/resource_jwt_key_mapping_crud.go @@ -19,6 +19,7 @@ func resourceLiteLLMJWTKeyMappingCreate(d *schema.ResourceData, m interface{}) e JWTClaimName: d.Get("jwt_claim_name").(string), JWTClaimValue: d.Get("jwt_claim_value").(string), Key: d.Get("key").(string), + Token: d.Get("token_id").(string), Description: d.Get("description").(string), } @@ -95,6 +96,7 @@ func resourceLiteLLMJWTKeyMappingUpdate(d *schema.ResourceData, m interface{}) e client := m.(*Client) oldKey, _ := d.GetChange("key") + oldTokenID, _ := d.GetChange("token_id") oldDescription, _ := d.GetChange("description") oldIsActive, _ := d.GetChange("is_active") @@ -104,6 +106,7 @@ func resourceLiteLLMJWTKeyMappingUpdate(d *schema.ResourceData, m interface{}) e // attempting to resync, so a failed refresh can't leave the rejected // values persisted into state. d.Set("key", oldKey) + d.Set("token_id", oldTokenID) d.Set("description", oldDescription) d.Set("is_active", oldIsActive) if readErr := resourceLiteLLMJWTKeyMappingRead(d, m); readErr != nil { @@ -146,6 +149,7 @@ func updateJWTKeyMapping(d *schema.ResourceData, client *Client) error { updateRequest := JWTKeyMappingUpdateRequest{ ID: d.Id(), Key: d.Get("key").(string), + Token: d.Get("token_id").(string), Description: d.Get("description").(string), IsActive: d.Get("is_active").(bool), } diff --git a/terraform/provider/litellm/resource_jwt_key_mapping_crud_test.go b/terraform/provider/litellm/resource_jwt_key_mapping_crud_test.go index 8007d1d4e08..27b75849fe6 100644 --- a/terraform/provider/litellm/resource_jwt_key_mapping_crud_test.go +++ b/terraform/provider/litellm/resource_jwt_key_mapping_crud_test.go @@ -628,3 +628,99 @@ func TestJWTKeyMappingCreateDoesNotLeakKeyInErrors(t *testing.T) { t.Fatalf("the virtual key must be redacted in errors, got %v", err) } } + +func TestJWTKeyMappingCreateSendsTokenIDAndOmitsKey(t *testing.T) { + srv, calls := jwtKeyMappingTestServer(t, jwtKeyMappingFixture()) + defer srv.Close() + + const tokenHash = "1923314ae0efc8b2523c7d421bac5a7cf88df291273b139948b526d396974a41" + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "token_id": tokenHash, + "is_active": true, + }) + + if err := resourceLiteLLMJWTKeyMappingCreate(d, client); err != nil { + t.Fatalf("create failed: %v", err) + } + + create := (*calls)[0] + if create.Body["token"] != tokenHash { + t.Fatalf("token hash not sent: %v", create.Body["token"]) + } + if _, sent := create.Body["key"]; sent { + t.Fatalf("key must be omitted when token_id is used, got: %v", create.Body) + } +} + +func TestJWTKeyMappingCreateOmitsTokenWhenKeyIsUsed(t *testing.T) { + srv, calls := jwtKeyMappingTestServer(t, jwtKeyMappingFixture()) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-abc123", + "is_active": true, + }) + + if err := resourceLiteLLMJWTKeyMappingCreate(d, client); err != nil { + t.Fatalf("create failed: %v", err) + } + + create := (*calls)[0] + if create.Body["key"] != "sk-abc123" { + t.Fatalf("virtual key not sent: %v", create.Body["key"]) + } + if _, sent := create.Body["token"]; sent { + t.Fatalf("token must be omitted when key is used, got: %v", create.Body) + } +} + +func TestJWTKeyMappingUpdateSendsTokenIDAndOmitsKey(t *testing.T) { + srv, calls := jwtKeyMappingTestServer(t, jwtKeyMappingFixture()) + defer srv.Close() + + const oldHash = "1111111111111111111111111111111111111111111111111111111111111111" + const newHash = "2222222222222222222222222222222222222222222222222222222222222222" + + client := NewClient(srv.URL, "test-key", true) + d := resourceDataWithChange(t, + map[string]string{ + "id": "map-abc-123", + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "token_id": oldHash, + "is_active": "true", + }, + map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "token_id": newHash, + "is_active": true, + }) + + if err := resourceLiteLLMJWTKeyMappingUpdate(d, client); err != nil { + t.Fatalf("update failed: %v", err) + } + + var update *jwtKeyMappingCall + for i := range *calls { + if (*calls)[i].Path == "/jwt/key/mapping/update" { + update = &(*calls)[i] + } + } + if update == nil { + t.Fatalf("expected an update call, got %v", *calls) + } + if update.Body["token"] != newHash { + t.Fatalf("new token hash not sent: %v", update.Body["token"]) + } + if _, sent := update.Body["key"]; sent { + t.Fatalf("key must be omitted when token_id is used, got: %v", update.Body) + } +} diff --git a/terraform/provider/litellm/types.go b/terraform/provider/litellm/types.go index 7bef44409fd..8bcf7dc4fe3 100644 --- a/terraform/provider/litellm/types.go +++ b/terraform/provider/litellm/types.go @@ -276,13 +276,15 @@ type VectorStoreInfoRequest struct { type JWTKeyMappingRequest struct { JWTClaimName string `json:"jwt_claim_name"` JWTClaimValue string `json:"jwt_claim_value"` - Key string `json:"key"` + Key string `json:"key,omitempty"` + Token string `json:"token,omitempty"` Description string `json:"description,omitempty"` } type JWTKeyMappingUpdateRequest struct { ID string `json:"id"` Key string `json:"key,omitempty"` + Token string `json:"token,omitempty"` Description string `json:"description"` IsActive bool `json:"is_active"` } diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index 3f2c04336a7..e95ed42013b 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -19,6 +19,8 @@ from litellm.proxy._types import ( ) from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( _to_response, + _token_hash_for_create, + _token_hash_for_update, create_jwt_key_mapping, delete_jwt_key_mapping, info_jwt_key_mapping, @@ -1515,6 +1517,132 @@ def test_jwt_client_id_field_does_not_raise_on_duplicate(): assert auth.virtual_key_claim_field == "new_field" +# ────────────────────────────────────────────── +# Tests: identifying the mapped key by hash instead of plaintext +# ────────────────────────────────────────────── + +_TOKEN_HASH = "1923314ae0efc8b2523c7d421bac5a7cf88df291273b139948b526d396974a41" + + +def test_create_stores_a_supplied_token_hash_verbatim(): + """A caller that holds only the hash gets it stored as given, not hashed again.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="email", jwt_claim_value="user@example.com", token=_TOKEN_HASH + ) + + assert _token_hash_for_create(data) == _TOKEN_HASH + + +def test_create_hashes_a_supplied_plaintext_key(): + """Supplying `key` keeps the original behaviour, so existing configs are unaffected.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest, hash_token + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="email", jwt_claim_value="user@example.com", key="sk-test-key" + ) + + assert _token_hash_for_create(data) == hash_token("sk-test-key") + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({}, id="neither"), + pytest.param({"key": "sk-test-key", "token": _TOKEN_HASH}, id="both"), + ], +) +def test_create_requires_exactly_one_identifier(kwargs): + """Neither or both is a 400, so a mapping can never be created ambiguously.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="email", jwt_claim_value="user@example.com", **kwargs + ) + + with pytest.raises(HTTPException) as exc_info: + _token_hash_for_create(data) + + assert exc_info.value.status_code == 400 + assert "exactly one" in exc_info.value.detail.lower() + + +@pytest.mark.parametrize( + "token", + [ + pytest.param("sk-not-a-hash", id="plaintext-key"), + pytest.param("abc123", id="too-short"), + pytest.param(_TOKEN_HASH.upper(), id="uppercase"), + pytest.param(_TOKEN_HASH + "0", id="too-long"), + pytest.param(_TOKEN_HASH[:-1] + "g", id="non-hex-character"), + ], +) +def test_create_rejects_a_token_that_is_not_a_sha256_hash(token): + """hash_token hashes unconditionally, so a bad `token` would be stored as a hash + of a hash and then silently match nothing at auth time.""" + from litellm.proxy._types import CreateJWTKeyMappingRequest + + data = CreateJWTKeyMappingRequest( + jwt_claim_name="email", jwt_claim_value="user@example.com", token=token + ) + + with pytest.raises(HTTPException) as exc_info: + _token_hash_for_create(data) + + assert exc_info.value.status_code == 400 + assert "SHA-256" in exc_info.value.detail + + +def test_update_leaves_the_mapped_key_alone_when_neither_is_given(): + """Updating only the description must not blank out the mapped key.""" + from litellm.proxy._types import UpdateJWTKeyMappingRequest + + data = UpdateJWTKeyMappingRequest(id="mapping-1", description="new text") + + assert _token_hash_for_update(data) is None + + +def test_update_stores_a_supplied_token_hash_verbatim(): + from litellm.proxy._types import UpdateJWTKeyMappingRequest + + data = UpdateJWTKeyMappingRequest(id="mapping-1", token=_TOKEN_HASH) + + assert _token_hash_for_update(data) == _TOKEN_HASH + + +def test_update_hashes_a_supplied_plaintext_key(): + from litellm.proxy._types import UpdateJWTKeyMappingRequest, hash_token + + data = UpdateJWTKeyMappingRequest(id="mapping-1", key="sk-rotated") + + assert _token_hash_for_update(data) == hash_token("sk-rotated") + + +def test_update_rejects_both_identifiers(): + from litellm.proxy._types import UpdateJWTKeyMappingRequest + + data = UpdateJWTKeyMappingRequest(id="mapping-1", key="sk-abc", token=_TOKEN_HASH) + + with pytest.raises(HTTPException) as exc_info: + _token_hash_for_update(data) + + assert exc_info.value.status_code == 400 + assert "at most one" in exc_info.value.detail.lower() + + +def test_update_rejects_a_token_that_is_not_a_sha256_hash(): + from litellm.proxy._types import UpdateJWTKeyMappingRequest + + data = UpdateJWTKeyMappingRequest(id="mapping-1", token="sk-not-a-hash") + + with pytest.raises(HTTPException) as exc_info: + _token_hash_for_update(data) + + assert exc_info.value.status_code == 400 + assert "SHA-256" in exc_info.value.detail + + # ────────────────────────────────────────────── # Tests: cache eviction must happen AFTER the DB write commits # ────────────────────────────────────────────── diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a0f6d2c8d8b..08b7ad5706a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27983,7 +27983,9 @@ export interface components { /** Jwt Issuer */ jwt_issuer?: string | null; /** Key */ - key: string; + key?: string | null; + /** Token */ + token?: string | null; }; /** CreateSearchToolRequest */ CreateSearchToolRequest: { @@ -39856,6 +39858,8 @@ export interface components { jwt_issuer?: string | null; /** Key */ key?: string | null; + /** Token */ + token?: string | null; }; /** UpdateKeyRequest */ UpdateKeyRequest: {