feat(proxy): configurable key_alias_pattern for key generate, update, and regenerate (#42553)

* feat(proxy): configurable key_alias_pattern for key generate, update, and regenerate Adds litellm_settings.key_alias_pattern, a regex every key_alias sent to /key/generate, /key/service-account/generate, /key/update, and /key/{key}/regenerate has to fully match. A non-matching alias gets a 400 that names the setting and the pattern. When set, it replaces the built-in rule enable_key_alias_format_validation turns on, and the baseline unsafe-name check still runs first. An invalid regex fails config load. * fix(proxy): cap key_alias length under key_alias_pattern and type the test fixtures * style(proxy): declare key_alias_pattern with a PEP 604 union --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
2026-09-24 00:52:24 +00:00 · 2026-09-22 14:52:01 -07:00 · 2026-09-22 14:52:01 -07:00 · d7c27cdc08
commit d7c27cdc08
parent d31e8aac6d
5 changed files with 219 additions and 16 deletions
--- 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
 )
--- 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,
--- 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()
--- 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."""
--- 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