diff --git a/litellm/__init__.py b/litellm/__init__.py index a044676a843..471b273f00d 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -381,6 +381,7 @@ enable_model_config_credential_overrides: bool = False enable_key_alias_format_validation: bool = ( False # opt-in validation of key_alias format on /key/generate and /key/update ) +key_alias_pattern: str | None = None enable_gemini_default_thinking_level_low: bool = ( False # opt-in: force thinkingLevel low/minimal for Gemini 3 thinking param mapping ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 801c45501ce..1845f11c0fa 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -7590,25 +7590,47 @@ async def test_key_logging( _KEY_ALIAS_PATTERN: Final = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9_\-/\.@]{0,253}[a-zA-Z0-9]$") +_KEY_ALIAS_PATTERN_MESSAGE: Final = ( + "Invalid key_alias format. Must be 2-255 characters, start/end with alphanumeric, and only contain a-zA-Z0-9_-/.@." +) +_KEY_ALIAS_MAX_LENGTH: Final = 255 + + +def parse_key_alias_pattern(value: object) -> str | None: + if value is None: + return None + if not isinstance(value, str): + raise ValueError( + f"Invalid regex set for litellm_settings.key_alias_pattern - value={value!r}: must be a string" + ) + try: + re.compile(value) + except re.error as e: + raise ValueError(f"Invalid regex set for litellm_settings.key_alias_pattern - value={value}: {e}") from e + return value + + +def _key_alias_rule() -> tuple[re.Pattern[str], str] | None: + if litellm.key_alias_pattern is not None: + return ( + re.compile(litellm.key_alias_pattern), + f"Invalid key_alias format. Must be at most {_KEY_ALIAS_MAX_LENGTH} characters and match the configured" + f" key_alias_pattern: {litellm.key_alias_pattern}", + ) + if litellm.enable_key_alias_format_validation: + return (_KEY_ALIAS_PATTERN, _KEY_ALIAS_PATTERN_MESSAGE) + return None def _validate_key_alias_format(key_alias: str | None) -> None: """ Validate the format of the key_alias. - A baseline validation always runs, regardless of - ``litellm.enable_key_alias_format_validation``. - - The remaining charset/length rules are gated behind - ``litellm.enable_key_alias_format_validation`` (default **False**). When disabled, - only the baseline validation above is performed, so existing workflows are not - broken. - - Rules (when enabled): - - None is OK (no alias). - - Otherwise must be 2–255 chars - - start/end with alphanumeric - - only allow a-zA-Z0-9_-/.@ + Path traversal and control characters are always rejected. The alias then has to + stay within ``_KEY_ALIAS_MAX_LENGTH`` and fully match ``litellm.key_alias_pattern`` + when one is configured, else the built-in pattern when + ``litellm.enable_key_alias_format_validation`` is on, else nothing more is checked + so existing workflows are not broken. """ if key_alias is None: return @@ -7623,12 +7645,14 @@ def _validate_key_alias_format(key_alias: str | None) -> None: code=400, ) - if not litellm.enable_key_alias_format_validation: + rule: Final = _key_alias_rule() + if rule is None: return - if not _KEY_ALIAS_PATTERN.match(key_alias): + pattern, message = rule + if len(key_alias) > _KEY_ALIAS_MAX_LENGTH or pattern.fullmatch(key_alias) is None: raise ProxyException( - message="Invalid key_alias format. Must be 2-255 characters, start/end with alphanumeric, and only contain a-zA-Z0-9_-/.@.", + message=message, type=ProxyErrorTypes.bad_request_error, param="key_alias", code=400, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 10110dd4f76..acb36fc4a59 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -596,6 +596,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_verification_tokens, duration_in_seconds, generate_key_helper_fn, + parse_key_alias_pattern, ) from litellm.proxy.management_endpoints.key_management_endpoints import ( router as key_management_router, @@ -6266,6 +6267,8 @@ class ProxyConfig: litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(**value) else: raise Exception(f"Invalid value set for upperbound_key_generate_params - value={value}") + elif key == "key_alias_pattern": + litellm.key_alias_pattern = parse_key_alias_pattern(value) elif key == "json_logs" and value is True: litellm.json_logs = True litellm._turn_on_json() diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index c539a10e1e2..cd929dfeeb3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -3089,6 +3089,50 @@ async def test_update_key_by_alias_only(monkeypatch): assert result["key"] == hashed_token +@pytest.mark.asyncio +async def test_update_key_changed_alias_must_match_key_alias_pattern(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + monkeypatch.setattr(litellm, "key_alias_pattern", r"^[a-z0-9]+(-[a-z0-9]+)*$") + hashed_token = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b" + key_in_db = LiteLLM_VerificationToken(token=hashed_token, key_alias="Legacy Alias", user_id="test-user") + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=key_in_db) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[key_in_db]) + mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(return_value=None) + mock_prisma_client.update_data = AsyncMock(return_value={"data": {"max_budget": 50.0}}) + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ) + + with pytest.raises(ProxyException) as exc_info: + await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key=hashed_token, key_alias="Prod Key"), + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + assert str(exc_info.value.code) == "400" + assert "key_alias_pattern" in str(exc_info.value.message) + mock_prisma_client.update_data.assert_not_awaited() + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + return_value=None, + ): + await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key=hashed_token, key_alias="Legacy Alias", max_budget=50.0), + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + mock_prisma_client.update_data.assert_awaited_once() + + @pytest.mark.asyncio async def test_update_key_by_alias_not_found_returns_404(monkeypatch): """ @@ -10562,6 +10606,10 @@ class TestValidateKeyAliasFormat: def reset_key_alias_flag(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "enable_key_alias_format_validation", False) + @pytest.fixture(autouse=True) + def reset_key_alias_pattern(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "key_alias_pattern", None) + def test_validation_skipped_when_flag_disabled(self): """When enable_key_alias_format_validation is False (default), no charset/length validation occurs.""" from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -10644,6 +10692,67 @@ class TestValidateKeyAliasFormat: assert str(exc.value.code) == "400" assert "Invalid key_alias format" in str(exc.value.message) + def test_configured_pattern_applies_with_flag_off(self, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_key_alias_format, + ) + + monkeypatch.setattr(litellm, "key_alias_pattern", r"^[a-z0-9]+(-[a-z0-9]+)*$") + with pytest.raises(ProxyException) as exc: + _validate_key_alias_format("Prod Key") + assert str(exc.value.code) == "400" + assert exc.value.param == "key_alias" + assert "key_alias_pattern" in str(exc.value.message) + assert r"^[a-z0-9]+(-[a-z0-9]+)*$" in str(exc.value.message) + assert _validate_key_alias_format("prod-key-001") is None + assert _validate_key_alias_format(None) is None + + def test_configured_pattern_must_match_the_whole_alias(self, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_key_alias_format, + ) + + monkeypatch.setattr(litellm, "key_alias_pattern", r"team-[a-z]+") + _validate_key_alias_format("team-search") + for partial_match in ("team-search-2", "xteam-search"): + with pytest.raises(ProxyException): + _validate_key_alias_format(partial_match) + + def test_configured_pattern_replaces_the_builtin_rule(self, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_key_alias_format, + ) + + monkeypatch.setattr(litellm, "enable_key_alias_format_validation", True) + monkeypatch.setattr(litellm, "key_alias_pattern", r"^[a-z ]+$") + _validate_key_alias_format("alias with spaces") + with pytest.raises(ProxyException) as exc: + _validate_key_alias_format("Uppercase") + assert "key_alias_pattern" in str(exc.value.message) + + def test_configured_pattern_keeps_the_baseline_safety_check(self, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_key_alias_format, + ) + + monkeypatch.setattr(litellm, "key_alias_pattern", r".*") + with pytest.raises(ProxyException) as exc: + _validate_key_alias_format("../../../other-app/creds") + assert str(exc.value.code) == "400" + assert "key_alias_pattern" not in str(exc.value.message) + + def test_configured_pattern_bounds_the_alias_length(self, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_key_alias_format, + ) + + monkeypatch.setattr(litellm, "key_alias_pattern", r"^[a-z]+$") + _validate_key_alias_format("a" * 255) + with pytest.raises(ProxyException) as exc: + _validate_key_alias_format("a" * 256) + assert str(exc.value.code) == "400" + assert "at most 255 characters" in str(exc.value.message) + @pytest.mark.asyncio async def test_check_org_key_limits_on_update_within_bounds(): @@ -12734,6 +12843,55 @@ async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monk assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_changed_alias_must_match_key_alias_pattern( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + monkeypatch.setattr(litellm, "key_alias_pattern", r"^[a-z0-9]+(-[a-z0-9]+)*$") + mock_prisma_client = _make_regenerate_mock_prisma() + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=_make_regenerate_existing_key(), + hashed_api_key="abc123", + key="abc123", + data=RegenerateKeyRequest(key_alias="Regenerated Key"), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + assert str(exc_info.value.code) == "400" + assert exc_info.value.param == "key_alias" + assert r"^[a-z0-9]+(-[a-z0-9]+)*$" in str(exc_info.value.message) + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + + @pytest.mark.asyncio async def test_execute_virtual_key_regeneration_allows_within_limit_duration(monkeypatch): """Regenerate must accept durations within upperbound_key_generate_params.duration.""" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 32b89885ac2..89fd9c5c9d4 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3555,6 +3555,23 @@ async def test_load_config_rejects_malformed_role_permissions(tmp_path): await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) +@pytest.mark.asyncio +async def test_load_config_compiles_key_alias_pattern_at_startup(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setattr(litellm, "key_alias_pattern", None) + config_file: Final = tmp_path / "config.yaml" + + config_file.write_text(yaml.dump({"model_list": [], "litellm_settings": {"key_alias_pattern": "^team-("}})) + with pytest.raises(Exception, match=r"litellm_settings\.key_alias_pattern"): + await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) + assert litellm.key_alias_pattern is None + + config_file.write_text(yaml.dump({"model_list": [], "litellm_settings": {"key_alias_pattern": "^team-[a-z]+$"}})) + await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) + assert litellm.key_alias_pattern == "^team-[a-z]+$" + + def test_os_environ_resolution_leaves_the_config_layer_holding_the_reference(monkeypatch): from litellm.proxy.proxy_server import ProxyConfig