feat: support azure keyvault (#39933)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: -LAN- <laipz8200@outlook.com>

.gitignore

@@ -140,6 +140,9 @@ dmypy.json
 pyrightconfig.json
 !api/pyrightconfig.json
 
+# import-linter
+.import_linter_cache/
+
 # Pyre type checker
 .pyre/
 .idea/'

api/.env.example

@@ -117,6 +117,19 @@ SQLALCHEMY_POOL_RESET_ON_RETURN=rollback
 # storage type: opendal, s3, aliyun-oss, azure-blob, baidu-obs, google-storage, huawei-obs, oci-storage, tencent-cos, volcengine-tos, supabase
 STORAGE_TYPE=opendal
 
+# Key provider configuration, used to encrypt/decrypt tenant credentials (LLM/tool provider secrets)
+# key provider type: local, azure-keyvault
+KEY_PROVIDER_TYPE=local
+
+# Azure Key Vault configuration, required when KEY_PROVIDER_TYPE=azure-keyvault
+# authentication uses DefaultAzureCredential (managed identity, environment variables, or Azure CLI login)
+AZURE_KEYVAULT_VAULT_URL=https://<your-vault-name>.vault.azure.net
+AZURE_KEYVAULT_KEY_SIZE=2048
+# optional: auto-rotate each tenant's key every N days. Leave empty to manage rotation manually.
+# rotation is safe because old key versions are never given an expiry and stay usable forever --
+# do NOT set a rotation policy on this key in the Azure portal that expires old versions.
+AZURE_KEYVAULT_ROTATION_INTERVAL_DAYS=
+
 # Apache OpenDAL storage configuration, refer to https://github.com/apache/opendal
 OPENDAL_SCHEME=fs
 OPENDAL_FS_ROOT=storage

api/.importlinter

@@ -15,6 +15,29 @@ root_packages =
     services
 include_external_packages = True
 
+[importlinter:contract:no-direct-rsa-imports]
+# Note: `libs` itself is deliberately excluded from source_modules -- import-linter's
+# `forbidden` contract cannot check a package against its own descendant (libs.rsa lives
+# inside libs). Sibling modules under libs/ importing libs.rsa are not covered by this
+# contract; the realistic risk this guards against is application/service code reaching
+# past the key provider abstraction, which is fully covered below.
+name = Only the key provider abstraction may import libs.rsa directly
+type = forbidden
+source_modules =
+    core
+    constants
+    context
+    configs
+    controllers
+    extensions
+    factories
+    models
+    tasks
+    services
+forbidden_modules =
+    libs.rsa
+allow_indirect_imports = True
+
 [importlinter:contract:machinery-framework-boundary]
 name = API machinery is framework neutral
 type = forbidden

api/app_factory.py

@@ -184,6 +184,7 @@ def initialize_extensions(app: DifyApp):
         ext_forward_refs,
         ext_hosting_provider,
         ext_import_modules,
+        ext_key_provider,
         ext_logging,
         ext_login,
         ext_logstore,
@@ -219,6 +220,7 @@ def initialize_extensions(app: DifyApp):
         ext_migrate,
         ext_redis,
         ext_storage,
+        ext_key_provider,  # Initialize after storage, since RSAKeyProvider reads private keys from it
         ext_set_secretkey,
         ext_logstore,  # Initialize logstore after storage, before celery
         ext_celery,

api/commands/system.py

@@ -11,7 +11,7 @@ from events.app_event import app_was_created
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from libs.db_migration_lock import DbMigrationAutoRenewLock
-from libs.rsa import generate_key_pair
+from libs.key_providers import generate_key_pair
 from models import Tenant
 from models.model import App, AppMode, Conversation
 from models.provider import Provider, ProviderModel

api/configs/middleware/__init__.py

@@ -7,6 +7,7 @@ from pydantic_settings import BaseSettings
 
 from .cache.redis_config import RedisConfig
 from .cache.redis_pubsub_config import RedisPubSubConfig
+from .key_provider.azure_keyvault_config import AzureKeyVaultConfig
 from .storage.aliyun_oss_storage_config import AliyunOSSStorageConfig
 from .storage.amazon_s3_storage_config import S3StorageConfig
 from .storage.azure_blob_storage_config import AzureBlobStorageConfig
@@ -83,6 +84,22 @@ class StorageConfig(BaseSettings):
     )
 
 
+_VALID_KEY_PROVIDER_TYPE = Literal[
+    "local",
+    "azure-keyvault",
+]
+
+
+class KeyProviderConfig(BaseSettings):
+    KEY_PROVIDER_TYPE: _VALID_KEY_PROVIDER_TYPE = Field(
+        description="Key provider used to encrypt/decrypt tenant credentials (LLM/tool provider secrets)."
+        " Options: 'local' (per-tenant RSA key pair, private key kept in the STORAGE_TYPE backend),"
+        " 'azure-keyvault' (per-tenant RSA key kept in Azure Key Vault, private key never leaves the vault)."
+        " Default is 'local'.",
+        default=cast(_VALID_KEY_PROVIDER_TYPE, "local"),
+    )
+
+
 class VectorStoreConfig(BaseSettings):
     VECTOR_STORE: str | None = Field(
         description="Type of vector store to use for efficient similarity search."
@@ -359,6 +376,9 @@ class MiddlewareConfig(
     KeywordStoreConfig,
     RedisConfig,
     RedisPubSubConfig,
+    # configs of the tenant credential encryption key provider
+    KeyProviderConfig,
+    AzureKeyVaultConfig,
     # configs of storage and storage providers
     StorageConfig,
     AliyunOSSStorageConfig,

api/configs/middleware/key_provider/azure_keyvault_config.py

@@ -0,0 +1,38 @@
+from pydantic import Field, field_validator
+from pydantic_settings import BaseSettings
+
+
+class AzureKeyVaultConfig(BaseSettings):
+    """
+    Configuration settings for Azure Key Vault, used as a tenant credential encryption key provider
+    """
+
+    AZURE_KEYVAULT_VAULT_URL: str | None = Field(
+        description="URL of the Azure Key Vault instance (e.g., 'https://<your-vault-name>.vault.azure.net')."
+        " Required when KEY_PROVIDER_TYPE is set to 'azure-keyvault'. Authentication uses"
+        " DefaultAzureCredential (managed identity, environment variables, or Azure CLI login).",
+        default=None,
+    )
+
+    AZURE_KEYVAULT_KEY_SIZE: int = Field(
+        description="RSA key size (in bits) used when Dify provisions a new per-tenant key in Azure Key Vault.",
+        default=2048,
+    )
+
+    AZURE_KEYVAULT_ROTATION_INTERVAL_DAYS: int | None = Field(
+        description="If set, Dify configures each newly-created per-tenant Key Vault key to auto-rotate every"
+        " N days (using a 'time after create' trigger, with no expiry set on generated versions)."
+        " Old key versions are kept forever and remain usable for decrypting credentials encrypted"
+        " before the rotation, so rotation requires no manual re-encryption. Leave unset (default) to"
+        " not configure a rotation policy and manage rotation manually in Azure. Must be at least 7 days.",
+        default=None,
+        ge=7,
+    )
+
+    @field_validator("AZURE_KEYVAULT_ROTATION_INTERVAL_DAYS", mode="before")
+    @classmethod
+    def _empty_string_to_none_for_rotation_interval(cls, v):
+        """Allow empty string in env/.env (e.g. an unfilled template value) to mean 'unset'."""
+        if isinstance(v, str) and v.strip() == "":
+            return None
+        return v

api/core/helper/encrypter.py

@@ -1,8 +1,5 @@
 import base64
-
-from Crypto.PublicKey import RSA
-
-from libs import rsa
+from typing import Any
 
 
 def obfuscated_token(token: str) -> str:
@@ -18,29 +15,35 @@ def full_mask_token(token_length: int = 20) -> str:
 
 
 def encrypt_token(tenant_id: str, token: str) -> str:
-    from models.account import Tenant
-    from models.engine import db
+    from extensions.ext_key_provider import key_provider_manager
 
-    if not (tenant := db.session.get(Tenant, tenant_id)):
-        raise ValueError(f"Tenant with id {tenant_id} not found")
-    assert tenant.encrypt_public_key is not None
-    encrypted_token = rsa.encrypt(token, tenant.encrypt_public_key)
+    encrypted_token = key_provider_manager.provider.encrypt(tenant_id, token)
     return base64.b64encode(encrypted_token).decode()
 
 
 def decrypt_token(tenant_id: str, token: str) -> str:
-    return rsa.decrypt(base64.b64decode(token), tenant_id)
+    from extensions.ext_key_provider import key_provider_manager
+
+    return key_provider_manager.provider.decrypt(tenant_id, base64.b64decode(token))
 
 
 def batch_decrypt_token(tenant_id: str, tokens: list[str]) -> list[str]:
-    rsa_key, cipher_rsa = rsa.get_decrypt_decoding(tenant_id)
-
-    return [rsa.decrypt_token_with_decoding(base64.b64decode(token), rsa_key, cipher_rsa) for token in tokens]
+    decoding = get_decrypt_decoding(tenant_id)
+    return [decrypt_token_with_decoding(token, decoding) for token in tokens]
 
 
-def get_decrypt_decoding(tenant_id: str) -> tuple[RSA.RsaKey, object]:
-    return rsa.get_decrypt_decoding(tenant_id)
+def get_decrypt_decoding(tenant_id: str) -> Any:
+    """
+    Return a reusable decoding context for batch/repeated decryption of a tenant's credentials
+    (e.g. across many provider/model configs in the same request). The returned object is opaque
+    and must only be passed back into decrypt_token_with_decoding.
+    """
+    from extensions.ext_key_provider import key_provider_manager
+
+    return key_provider_manager.provider.get_decrypt_decoding(tenant_id)
 
 
-def decrypt_token_with_decoding(token: str, rsa_key: RSA.RsaKey, cipher_rsa: object) -> str:
-    return rsa.decrypt_token_with_decoding(base64.b64decode(token), rsa_key, cipher_rsa)
+def decrypt_token_with_decoding(token: str, decoding: Any) -> str:
+    from extensions.ext_key_provider import key_provider_manager
+
+    return key_provider_manager.provider.decrypt_with_decoding(base64.b64decode(token), decoding)

api/core/provider_manager.py

@@ -574,17 +574,24 @@ class ProviderManager:
     instance scope.
     """
 
-    decoding_rsa_key: Any | None
-    decoding_cipher_rsa: Any | None
+    # Keyed by tenant_id -- a single ProviderManager instance may be asked to decrypt
+    # credentials belonging to different tenants (e.g. load balancing configs each carry
+    # their own tenant_id), so this cache must not collapse to a single shared value.
+    _decoding_contexts: dict[str, Any]
     _model_runtime: ModelRuntime
     _configurations_cache: dict[str, ProviderConfigurations]
 
     def __init__(self, model_runtime: ModelRuntime):
-        self.decoding_rsa_key = None
-        self.decoding_cipher_rsa = None
+        self._decoding_contexts = {}
         self._model_runtime = model_runtime
         self._configurations_cache = {}
 
+    def _get_decoding_context(self, tenant_id: str) -> Any:
+        """Return this manager's cached decoding context for `tenant_id`, fetching it once if absent."""
+        if tenant_id not in self._decoding_contexts:
+            self._decoding_contexts[tenant_id] = encrypter.get_decrypt_decoding(tenant_id)
+        return self._decoding_contexts[tenant_id]
+
     def clear_configurations_cache(self, tenant_id: str | None = None) -> None:
         """Drop assembled provider configurations cached on this manager instance."""
         if tenant_id is None:
@@ -1501,16 +1508,14 @@ class ProviderManager:
             return {}
 
         # Decrypt secret variables
-        if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None:
-            self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding(tenant_id)
+        decoding_context = self._get_decoding_context(tenant_id)
 
         for variable in secret_variables:
             if variable in credentials:
                 with contextlib.suppress(ValueError):
                     credentials[variable] = encrypter.decrypt_token_with_decoding(
                         credentials.get(variable) or "",
-                        self.decoding_rsa_key,
-                        self.decoding_cipher_rsa,
+                        decoding_context,
                     )
 
         # Cache the decrypted credentials
@@ -1662,17 +1667,15 @@ class ProviderManager:
                         else []
                     )
 
-                    # Get decoding rsa key and cipher for decrypting credentials
-                    if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None:
-                        self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding(tenant_id)
+                    # Get decoding context for decrypting credentials
+                    decoding_context = self._get_decoding_context(tenant_id)
 
                     for variable in provider_credential_secret_variables:
                         if variable in provider_credentials:
                             try:
                                 provider_credentials[variable] = encrypter.decrypt_token_with_decoding(
                                     provider_credentials.get(variable, ""),
-                                    self.decoding_rsa_key,
-                                    self.decoding_cipher_rsa,
+                                    decoding_context,
                                 )
                             except ValueError:
                                 pass
@@ -1806,19 +1809,15 @@ class ProviderManager:
                             except (ValueError, JSONDecodeError):
                                 continue
 
-                            # Get decoding rsa key and cipher for decrypting credentials
-                            if self.decoding_rsa_key is None or self.decoding_cipher_rsa is None:
-                                self.decoding_rsa_key, self.decoding_cipher_rsa = encrypter.get_decrypt_decoding(
-                                    load_balancing_model_config.tenant_id
-                                )
+                            # Get decoding context for decrypting credentials
+                            decoding_context = self._get_decoding_context(load_balancing_model_config.tenant_id)
 
                             for variable in model_credential_secret_variables:
                                 if variable in provider_model_credentials:
                                     try:
                                         provider_model_credentials[variable] = encrypter.decrypt_token_with_decoding(
                                             provider_model_credentials.get(variable) or "",
-                                            self.decoding_rsa_key,
-                                            self.decoding_cipher_rsa,
+                                            decoding_context,
                                         )
                                     except ValueError:
                                         pass

api/extensions/ext_key_provider.py

@@ -0,0 +1,50 @@
+import logging
+from collections.abc import Callable
+
+from flask import Flask
+
+from configs import dify_config
+from dify_app import DifyApp
+from libs.key_providers.base import BaseKeyProvider
+from libs.key_providers.key_provider_type import KeyProviderType
+
+logger = logging.getLogger(__name__)
+
+
+class KeyProviderManager:
+    _provider: BaseKeyProvider | None = None
+
+    def init_app(self, app: Flask):
+        with app.app_context():
+            self._provider = self._build_provider()
+
+    @property
+    def provider(self) -> BaseKeyProvider:
+        if self._provider is None:
+            self._provider = self._build_provider()
+        return self._provider
+
+    def _build_provider(self) -> BaseKeyProvider:
+        provider_factory = self.get_provider_factory(dify_config.KEY_PROVIDER_TYPE)
+        return provider_factory()
+
+    @staticmethod
+    def get_provider_factory(provider_type: str) -> Callable[[], BaseKeyProvider]:
+        match provider_type:
+            case KeyProviderType.LOCAL:
+                from libs.key_providers.rsa_key_provider import RSAKeyProvider
+
+                return RSAKeyProvider
+            case KeyProviderType.AZURE_KEYVAULT:
+                from libs.key_providers.azure_keyvault_key_provider import AzureKeyVaultKeyProvider
+
+                return AzureKeyVaultKeyProvider
+            case _:
+                raise ValueError(f"unsupported key provider type {provider_type}")
+
+
+key_provider_manager = KeyProviderManager()
+
+
+def init_app(app: DifyApp):
+    key_provider_manager.init_app(app)

api/libs/key_providers/__init__.py

@@ -0,0 +1,14 @@
+from libs.key_providers.base import BaseKeyProvider
+
+__all__ = ["BaseKeyProvider", "generate_key_pair"]
+
+
+def generate_key_pair(tenant_id: str) -> str:
+    """
+    Provision the tenant credential encryption key using the configured KEY_PROVIDER_TYPE.
+
+    Returns the opaque reference to be stored in Tenant.encrypt_public_key.
+    """
+    from extensions.ext_key_provider import key_provider_manager
+
+    return key_provider_manager.provider.generate_key_pair(tenant_id)

api/libs/key_providers/azure_keyvault_key_provider.py

@@ -0,0 +1,203 @@
+import json
+from datetime import UTC, datetime
+from typing import Any, override
+
+from azure.core.exceptions import AzureError
+from azure.identity import DefaultAzureCredential
+from azure.keyvault.keys import (
+    KeyClient,
+    KeyRotationLifetimeAction,
+    KeyRotationPolicy,
+    KeyRotationPolicyAction,
+)
+from azure.keyvault.keys.crypto import CryptographyClient, KeyWrapAlgorithm
+from Crypto.Cipher import AES
+from Crypto.Random import get_random_bytes
+
+from configs import dify_config
+from libs.key_providers.base import BaseKeyProvider
+
+# Marker kept identical to libs/rsa.py so ciphertext produced by either provider is
+# self-describing, even though the two providers never decode each other's payloads.
+_PREFIX = b"HYBRID:"
+
+# Bump this if the *binary layout* below ever changes. Adding new metadata fields does NOT
+# require a bump (metadata is a JSON object -- old readers just ignore unknown keys via .get()).
+_ENVELOPE_VERSION = 1
+
+_DEFAULT_WRAP_ALGORITHM = KeyWrapAlgorithm.rsa_oaep_256
+
+
+class AzureKeyVaultKeyProvider(BaseKeyProvider):
+    """
+    Envelope-encryption key provider backed by Azure Key Vault.
+
+    Ciphertext envelope (self-describing, forward-compatible):
+        PREFIX (7 bytes)
+      + envelope_version (1 byte)
+      + metadata_len (2 bytes, big-endian) + metadata (UTF-8 JSON object)
+      + wrapped_key_len (2 bytes, big-endian) + wrapped_key
+      + nonce (16 bytes) + tag (16 bytes) + ciphertext
+
+    `metadata` currently carries {"key_version": ..., "wrap_alg": ...}. It's a JSON object
+    rather than fixed-width fields so new attributes can be added later without touching the
+    binary layout or breaking old ciphertext; `envelope_version` exists separately to guard
+    the binary layout itself, in case that ever needs to change.
+
+    Recording the key_version that wrapped each token (rather than always resolving "the
+    current version" at decrypt time) is what makes Key Vault's native automatic key rotation
+    safe to use here: old tokens keep decrypting against the version that encrypted them,
+    while new tokens pick up whatever version is current. This only holds as long as old
+    versions are never allowed to *expire* -- see generate_key_pair().
+    """
+
+    def __init__(self):
+        vault_url = dify_config.AZURE_KEYVAULT_VAULT_URL
+        if not vault_url:
+            raise ValueError("AZURE_KEYVAULT_VAULT_URL must be configured when KEY_PROVIDER_TYPE=azure-keyvault")
+
+        self._vault_url = vault_url
+        self._credential = DefaultAzureCredential()
+        self._key_client = KeyClient(vault_url=vault_url, credential=self._credential)
+
+    @staticmethod
+    def _key_name(tenant_id: str) -> str:
+        return f"dify-tenant-{tenant_id}"
+
+    def _get_crypto_client(self, tenant_id: str, version: str | None = None) -> tuple[CryptographyClient, str]:
+        """
+        Return a CryptographyClient bound to `version`, along with the resolved version string
+        that was actually used.
+        """
+        key_name = self._key_name(tenant_id)
+        if version is not None:
+            resolved_version = version
+        else:
+            versions = list(self._key_client.list_properties_of_key_versions(key_name))
+            if not versions:
+                raise ValueError(f"No key versions found for key {key_name}")
+            current = max(
+                versions,
+                key=lambda properties: properties.created_on or datetime.min.replace(tzinfo=UTC),
+            )
+            resolved_version = current.version or ""
+        return self._key_client.get_cryptography_client(key_name, key_version=resolved_version), resolved_version
+
+    @override
+    def generate_key_pair(self, tenant_id: str) -> str:
+        key_name = self._key_name(tenant_id)
+        self._key_client.create_rsa_key(key_name, size=dify_config.AZURE_KEYVAULT_KEY_SIZE)
+
+        rotation_interval_days = dify_config.AZURE_KEYVAULT_ROTATION_INTERVAL_DAYS
+        if rotation_interval_days:
+            self._key_client.update_key_rotation_policy(
+                key_name,
+                policy=KeyRotationPolicy(
+                    lifetime_actions=[
+                        KeyRotationLifetimeAction(
+                            KeyRotationPolicyAction.rotate,
+                            time_after_create=f"P{rotation_interval_days}D",
+                        )
+                    ],
+                    # Deliberately no `expires_in` here: this provider pins each ciphertext to
+                    # the key_version that encrypted it and relies on old versions staying
+                    # usable forever. If versions were also given an expiry (time_before_expiry
+                    # trigger / expires_in), old ciphertext would become permanently
+                    # undecryptable once its version expired, unless a separate re-wrap/
+                    # migration job proactively moves it to the new version first.
+                ),
+            )
+        return key_name
+
+    @override
+    def encrypt(self, tenant_id: str, text: str) -> bytes:
+        aes_key = get_random_bytes(16)
+        cipher_aes = AES.new(aes_key, AES.MODE_EAX)
+        ciphertext, tag = cipher_aes.encrypt_and_digest(text.encode())
+
+        crypto_client, key_version = self._get_crypto_client(tenant_id)
+        wrapped_key = crypto_client.wrap_key(_DEFAULT_WRAP_ALGORITHM, aes_key).encrypted_key
+
+        metadata = json.dumps({"key_version": key_version, "wrap_alg": _DEFAULT_WRAP_ALGORITHM.value}).encode()
+
+        return (
+            _PREFIX
+            + _ENVELOPE_VERSION.to_bytes(1, "big")
+            + len(metadata).to_bytes(2, "big")
+            + metadata
+            + len(wrapped_key).to_bytes(2, "big")
+            + wrapped_key
+            + cipher_aes.nonce
+            + tag
+            + ciphertext
+        )
+
+    @override
+    def get_decrypt_decoding(self, tenant_id: str) -> str:
+        return tenant_id
+
+    @override
+    def decrypt_with_decoding(self, encrypted_text: bytes, decoding: str) -> str:
+        tenant_id = decoding
+        if not encrypted_text.startswith(_PREFIX):
+            raise ValueError("Unsupported ciphertext format for Azure Key Vault key provider")
+
+        # Bytes slicing never raises on out-of-range indices in Python (it just returns a
+        # shorter/empty slice), so a truncated envelope wouldn't otherwise surface as an error
+        # until (maybe) AES decryption fails much later, or not at all. Validate lengths
+        # explicitly and turn any parsing failure into ValueError, matching what callers
+        # (e.g. core/provider_manager.py) already expect and suppress for malformed credentials.
+        try:
+            body = encrypted_text[len(_PREFIX) :]
+            if len(body) < 1:
+                raise ValueError("Malformed Azure Key Vault envelope: missing envelope version")
+            envelope_version = body[0]
+            if envelope_version != _ENVELOPE_VERSION:
+                raise ValueError(f"Unsupported Azure Key Vault envelope version: {envelope_version}")
+            offset = 1
+
+            if len(body) < offset + 2:
+                raise ValueError("Malformed Azure Key Vault envelope: truncated metadata length")
+            metadata_len = int.from_bytes(body[offset : offset + 2], "big")
+            offset += 2
+            if len(body) < offset + metadata_len:
+                raise ValueError("Malformed Azure Key Vault envelope: truncated metadata")
+            metadata: Any = json.loads(body[offset : offset + metadata_len])
+            if not isinstance(metadata, dict):
+                raise ValueError("Malformed Azure Key Vault envelope: metadata is not a JSON object")
+            offset += metadata_len
+
+            if len(body) < offset + 2:
+                raise ValueError("Malformed Azure Key Vault envelope: truncated wrapped key length")
+            key_len = int.from_bytes(body[offset : offset + 2], "big")
+            offset += 2
+            if len(body) < offset + key_len + 16 + 16:
+                raise ValueError("Malformed Azure Key Vault envelope: truncated wrapped key/nonce/tag")
+            wrapped_key = body[offset : offset + key_len]
+            offset += key_len
+            nonce = body[offset : offset + 16]
+            offset += 16
+            tag = body[offset : offset + 16]
+            offset += 16
+            ciphertext = body[offset:]
+        except (IndexError, TypeError) as exc:
+            raise ValueError("Malformed Azure Key Vault envelope") from exc
+
+        wrap_alg = KeyWrapAlgorithm(metadata["wrap_alg"]) if metadata.get("wrap_alg") else _DEFAULT_WRAP_ALGORITHM
+        try:
+            # A specific key_version can legitimately become unusable after this ciphertext was
+            # created -- disabled, deleted, or (if a rotation policy with an expiry was
+            # misconfigured despite generate_key_pair()'s warning against it) expired. Every
+            # caller of decrypt_token_with_decoding (core/provider_manager.py,
+            # services/model_load_balancing_service.py) already only expects/suppresses
+            # ValueError for "this particular credential can't be decrypted right now", so Azure
+            # SDK errors must be translated here rather than left to escape as a different type
+            # and crash the whole call chain (e.g. building a tenant's full provider
+            # configuration just to create an unrelated new credential).
+            crypto_client, _ = self._get_crypto_client(tenant_id, version=metadata.get("key_version"))
+            aes_key = crypto_client.unwrap_key(wrap_alg, wrapped_key).key
+        except AzureError as exc:
+            raise ValueError(f"Failed to unwrap credential via Azure Key Vault: {exc}") from exc
+
+        cipher_aes = AES.new(aes_key, AES.MODE_EAX, nonce=nonce)
+        return cipher_aes.decrypt_and_verify(ciphertext, tag).decode()

api/libs/key_providers/base.py

@@ -0,0 +1,34 @@
+"""Abstract interface for tenant credential encryption key providers."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class BaseKeyProvider(ABC):
+    """Interface for providers that manage the keys used to encrypt/decrypt tenant credentials."""
+
+    @abstractmethod
+    def generate_key_pair(self, tenant_id: str) -> str:
+        """
+        Provision the encryption key for a tenant.
+
+        Returns an opaque reference to be stored in Tenant.encrypt_public_key
+        (e.g. a PEM public key, or a key vault key name/identifier).
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def encrypt(self, tenant_id: str, text: str) -> bytes:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_decrypt_decoding(self, tenant_id: str) -> Any:
+        """Return a reusable decoding context, so batch decryption can avoid repeated key lookups."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def decrypt_with_decoding(self, encrypted_text: bytes, decoding: Any) -> str:
+        raise NotImplementedError
+
+    def decrypt(self, tenant_id: str, encrypted_text: bytes) -> str:
+        return self.decrypt_with_decoding(encrypted_text, self.get_decrypt_decoding(tenant_id))

api/libs/key_providers/key_provider_type.py

@@ -0,0 +1,6 @@
+from enum import StrEnum
+
+
+class KeyProviderType(StrEnum):
+    LOCAL = "local"
+    AZURE_KEYVAULT = "azure-keyvault"

api/libs/key_providers/rsa_key_provider.py

@@ -0,0 +1,47 @@
+from typing import override
+
+from Crypto.PublicKey import RSA
+
+from libs import rsa
+from libs.key_providers.base import BaseKeyProvider
+
+
+class RSAKeyProvider(BaseKeyProvider):
+    """
+    Default key provider: per-tenant RSA key pair.
+
+    The private key is kept in the configured STORAGE_TYPE backend (see libs/rsa.py).
+    This provider only composes the existing libs.rsa implementation; the underlying
+    crypto logic is intentionally left untouched.
+    """
+
+    @override
+    def generate_key_pair(self, tenant_id: str) -> str:
+        return rsa.generate_key_pair(tenant_id)
+
+    @override
+    def encrypt(self, tenant_id: str, text: str) -> bytes:
+        from models.account import Tenant
+        from models.engine import db
+
+        if not (tenant := db.session.get(Tenant, tenant_id)):
+            raise ValueError(f"Tenant with id {tenant_id} not found")
+        if tenant.encrypt_public_key is None:
+            raise ValueError(f"Tenant with id {tenant_id} has no encrypt_public_key")
+        return rsa.encrypt(text, tenant.encrypt_public_key)
+
+    @override
+    def get_decrypt_decoding(self, tenant_id: str) -> tuple[RSA.RsaKey, object]:
+        return rsa.get_decrypt_decoding(tenant_id)
+
+    @override
+    def decrypt_with_decoding(self, encrypted_text: bytes, decoding: tuple[RSA.RsaKey, object]) -> str:
+        rsa_key, cipher_rsa = decoding
+        return rsa.decrypt_token_with_decoding(encrypted_text, rsa_key, cipher_rsa)
+
+    @override
+    def decrypt(self, tenant_id: str, encrypted_text: bytes) -> str:
+        # Overrides BaseKeyProvider's generic get_decrypt_decoding()+decrypt_with_decoding()
+        # composition to call libs.rsa.decrypt() directly (a single-shot equivalent), matching
+        # this provider's one supported decrypt path in libs/rsa.py.
+        return rsa.decrypt(encrypted_text, tenant_id)

api/libs/rsa.py

@@ -1,3 +1,20 @@
+"""
+Low-level implementation of the default ("local") tenant credential encryption key provider.
+
+Do NOT import this module directly to encrypt/decrypt tenant credentials. It only implements
+one specific key provider (per-tenant RSA key pair, private key kept in STORAGE_TYPE). Other
+KEY_PROVIDER_TYPE options (e.g. 'azure-keyvault') are not implemented here.
+
+Instead use:
+  - core.helper.encrypter (encrypt_token / decrypt_token / batch_decrypt_token / ...) for
+    application code that needs to encrypt or decrypt tenant credentials.
+  - libs.key_providers (generate_key_pair) when provisioning the key for a new tenant.
+
+This module is only meant to be imported by libs.key_providers.rsa_key_provider.RSAKeyProvider,
+which the rest of the codebase should reach through extensions.ext_key_provider.key_provider_manager.
+This is enforced by the "no-direct-rsa-imports" contract in .importlinter (run via `make lint`).
+"""
+
 import hashlib
 from typing import Union
 

api/pyproject.toml

@@ -107,7 +107,7 @@ dify-trace-tencent = { workspace = true }
 dify-trace-weave = { workspace = true }
 
 [tool.uv]
-default-groups = ["storage", "tools", "vdb-all", "trace-all"]
+default-groups = ["kms", "storage", "tools", "vdb-all", "trace-all"]
 package = false
 override-dependencies = [
     "litellm>=1.83.10,<2.0.0",
@@ -190,6 +190,14 @@ dev = [
     "types-croniter>=6.2.4.20260711",
 ]
 
+############################################################
+# [ KMS ] dependency group
+# Required for key management / credential encryption key providers
+############################################################
+kms = [
+    "azure-keyvault-keys>=4.10.0,<5.0.0",
+]
+
 ############################################################
 # [ Storage ] dependency group
 # Required for storage clients

api/services/account_service.py

@@ -28,9 +28,9 @@ from extensions.ext_redis import redis_client, redis_fallback
 from libs.datetime_utils import naive_utc_now
 from libs.helper import RateLimiter, TokenManager
 from libs.helper import timezone as validate_timezone
+from libs.key_providers import generate_key_pair
 from libs.passport import PassportService
 from libs.password import compare_password, hash_password, valid_password
-from libs.rsa import generate_key_pair
 from libs.token import generate_csrf_token
 from models.account import (
     Account,

api/services/model_load_balancing_service.py

@@ -178,8 +178,8 @@ class ModelLoadBalancingService:
         # Get credential form schemas from model credential schema or provider credential schema
         credential_schemas = self._get_credential_schema(provider_configuration)
 
-        # Get decoding rsa key and cipher for decrypting credentials
-        decoding_rsa_key, decoding_cipher_rsa = encrypter.get_decrypt_decoding(tenant_id)
+        # Get decoding context for decrypting credentials
+        decoding_context = encrypter.get_decrypt_decoding(tenant_id)
 
         # fetch status and ttl for each config
         datas: list[LoadBalancingConfigSummaryDict] = []
@@ -213,8 +213,7 @@ class ModelLoadBalancingService:
                         if isinstance(token_value, str):
                             credentials[variable] = encrypter.decrypt_token_with_decoding(
                                 token_value,
-                                decoding_rsa_key,
-                                decoding_cipher_rsa,
+                                decoding_context,
                             )
                     except ValueError:
                         pass

api/tests/unit_tests/libs/key_providers/test_azure_keyvault_key_provider.py

@@ -0,0 +1,265 @@
+"""
+Unit tests for AzureKeyVaultKeyProvider, with the Azure SDK mocked out (no network calls).
+
+The main thing under test is that ciphertext is pinned to the key *version* that wrapped it,
+so that Key Vault's native key rotation (a new "current" version appearing) does not break
+decryption of tokens encrypted before the rotation.
+"""
+
+import datetime
+import json
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+from azure.core.exceptions import HttpResponseError
+from azure.keyvault.keys import KeyRotationPolicy
+
+from configs import dify_config
+from libs.key_providers.azure_keyvault_key_provider import AzureKeyVaultKeyProvider
+
+
+class FakeCryptographyClient:
+    """
+    Stands in for azure.keyvault.keys.crypto.CryptographyClient. wrap_key/unwrap_key here just
+    tag the payload with the bound key version, so unwrapping with the "wrong" version's client
+    can be detected -- mirroring how a real RSA key from a different version can't unwrap data
+    wrapped under another version's key.
+
+    Constructed from a key *id string* (e.g. ".../keys/<name>/<version>"), mirroring how the
+    real provider now avoids ever handing CryptographyClient an already-materialized key --
+    see AzureKeyVaultKeyProvider._get_crypto_client().
+    """
+
+    # Class-level so a test can mark a version "disabled" (mirroring Key Vault's real behavior
+    # for a disabled/deleted key version) without threading state through every fixture.
+    disabled_versions: set[str] = set()
+
+    def __init__(self, key: str, credential: object = None) -> None:
+        self.version = key.rsplit("/", 1)[-1]
+        self.credential = credential
+        self.closed = False
+
+    def close(self) -> None:
+        self.closed = True
+
+    def wrap_key(self, algorithm: object, key_bytes: bytes) -> SimpleNamespace:
+        self.last_wrap_algorithm = algorithm
+        return SimpleNamespace(encrypted_key=f"{self.version}:".encode() + key_bytes)
+
+    def unwrap_key(self, algorithm: object, wrapped: bytes) -> SimpleNamespace:
+        self.last_unwrap_algorithm = algorithm
+        if self.version in self.disabled_versions:
+            raise HttpResponseError(message="Operation unwrapKey is not allowed on a disabled key.")
+        prefix = f"{self.version}:".encode()
+        if not wrapped.startswith(prefix):
+            raise ValueError(f"key version {self.version} cannot unwrap data wrapped by another version")
+        return SimpleNamespace(key=wrapped[len(prefix) :])
+
+
+class FakeKeyClient:
+    """Stands in for azure.keyvault.keys.KeyClient."""
+
+    def __init__(self, vault_url: str | None = None, credential: object = None) -> None:
+        self.vault_url = vault_url
+        self.credential = credential
+        self.current_version = "v1"
+        self.rotation_policies: dict[str, KeyRotationPolicy] = {}
+        self.created_keys: dict[str, int] = {}
+        # version -> created_on, in creation order; mirrors what list_properties_of_key_versions
+        # reports in real Key Vault, which the provider now relies on (instead of GetKey) to
+        # resolve "the current version" without ever fetching key material.
+        self._version_created_on: dict[str, datetime.datetime] = {}
+
+    def _register_current_version(self) -> None:
+        if self.current_version not in self._version_created_on:
+            self._version_created_on[self.current_version] = datetime.datetime(2024, 1, 1) + datetime.timedelta(
+                seconds=len(self._version_created_on)
+            )
+
+    def create_rsa_key(self, name: str, size: int = 2048) -> SimpleNamespace:
+        self.created_keys[name] = size
+        self._register_current_version()
+        return SimpleNamespace(properties=SimpleNamespace(version=self.current_version))
+
+    def list_properties_of_key_versions(self, name: str) -> list[SimpleNamespace]:
+        assert name in self.created_keys, f"key {name} was never created"
+        self._register_current_version()
+        return [
+            SimpleNamespace(version=version, created_on=created_on, id=f"{self.vault_url}/keys/{name}/{version}")
+            for version, created_on in self._version_created_on.items()
+        ]
+
+    def update_key_rotation_policy(self, name: str, policy: KeyRotationPolicy) -> KeyRotationPolicy:
+        self.rotation_policies[name] = policy
+        return policy
+
+    def get_cryptography_client(self, name: str, *, key_version: str | None = None) -> "FakeCryptographyClient":
+        assert name in self.created_keys, f"key {name} was never created"
+        key_id = f"{self.vault_url}/keys/{name}/{key_version}"
+        return FakeCryptographyClient(key_id, credential=self.credential)
+
+
+@pytest.fixture
+def fake_key_client(monkeypatch: pytest.MonkeyPatch) -> FakeKeyClient:
+    client = FakeKeyClient()
+    monkeypatch.setattr(
+        "libs.key_providers.azure_keyvault_key_provider.KeyClient",
+        MagicMock(return_value=client),
+    )
+    monkeypatch.setattr(
+        "libs.key_providers.azure_keyvault_key_provider.CryptographyClient",
+        FakeCryptographyClient,
+    )
+    monkeypatch.setattr(
+        "libs.key_providers.azure_keyvault_key_provider.DefaultAzureCredential",
+        MagicMock(),
+    )
+    FakeCryptographyClient.disabled_versions = set()
+    return client
+
+
+@pytest.fixture(autouse=True)
+def azure_keyvault_config(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(dify_config, "AZURE_KEYVAULT_VAULT_URL", "https://fake-vault.vault.azure.net")
+    monkeypatch.setattr(dify_config, "AZURE_KEYVAULT_KEY_SIZE", 2048)
+    monkeypatch.setattr(dify_config, "AZURE_KEYVAULT_ROTATION_INTERVAL_DAYS", None)
+
+
+@pytest.mark.usefixtures("fake_key_client")
+def test_missing_vault_url_raises(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(dify_config, "AZURE_KEYVAULT_VAULT_URL", None)
+    with pytest.raises(ValueError, match="AZURE_KEYVAULT_VAULT_URL"):
+        AzureKeyVaultKeyProvider()
+
+
+@pytest.mark.usefixtures("fake_key_client")
+def test_encrypt_decrypt_roundtrip() -> None:
+    provider = AzureKeyVaultKeyProvider()
+    provider.generate_key_pair("tenant-1")
+
+    encrypted = provider.encrypt("tenant-1", "super-secret")
+    decoding = provider.get_decrypt_decoding("tenant-1")
+    assert provider.decrypt_with_decoding(encrypted, decoding) == "super-secret"
+
+
+def _embedded_key_version(envelope: bytes) -> str:
+    """Peel out the {"key_version": ...} metadata Dify embeds in each ciphertext, for assertions."""
+    body = envelope[len(b"HYBRID:") :]
+    metadata_len = int.from_bytes(body[1:3], "big")
+    metadata = json.loads(body[3 : 3 + metadata_len])
+    return metadata["key_version"]
+
+
+def test_rotation_does_not_break_decryption_of_old_ciphertext(fake_key_client: FakeKeyClient) -> None:
+    """
+    The core guarantee: a token encrypted before rotation must still decrypt correctly after
+    Key Vault promotes a new "current" version, and new tokens must promptly pick up the new
+    version rather than keep using a stale cached "current" client.
+    """
+    provider = AzureKeyVaultKeyProvider()
+    provider.generate_key_pair("tenant-1")
+
+    encrypted_v1 = provider.encrypt("tenant-1", "secret-before-rotation")
+    assert _embedded_key_version(encrypted_v1) == "v1"
+
+    # Simulate Key Vault's rotation policy promoting a new version in the background.
+    fake_key_client.current_version = "v2"
+
+    encrypted_v2 = provider.encrypt("tenant-1", "secret-after-rotation")
+    # The new version must be picked up immediately, not after some cache TTL elapses.
+    assert _embedded_key_version(encrypted_v2) == "v2"
+
+    decoding = provider.get_decrypt_decoding("tenant-1")
+    assert provider.decrypt_with_decoding(encrypted_v1, decoding) == "secret-before-rotation"
+    assert provider.decrypt_with_decoding(encrypted_v2, decoding) == "secret-after-rotation"
+
+
+def test_generate_key_pair_without_rotation_interval_does_not_set_policy(fake_key_client: FakeKeyClient) -> None:
+    provider = AzureKeyVaultKeyProvider()
+    provider.generate_key_pair("tenant-1")
+    assert fake_key_client.rotation_policies == {}
+
+
+def test_generate_key_pair_with_rotation_interval_sets_time_after_create_only(
+    monkeypatch: pytest.MonkeyPatch, fake_key_client: FakeKeyClient
+) -> None:
+    monkeypatch.setattr(dify_config, "AZURE_KEYVAULT_ROTATION_INTERVAL_DAYS", 30)
+
+    provider = AzureKeyVaultKeyProvider()
+    provider.generate_key_pair("tenant-1")
+
+    policy = fake_key_client.rotation_policies["dify-tenant-tenant-1"]
+    assert policy.expires_in is None
+    (action,) = policy.lifetime_actions
+    assert action.time_after_create == "P30D"
+    assert action.time_before_expiry is None
+
+
+@pytest.mark.usefixtures("fake_key_client")
+def test_decrypt_rejects_unknown_envelope_version() -> None:
+    provider = AzureKeyVaultKeyProvider()
+    provider.generate_key_pair("tenant-1")
+
+    encrypted = bytearray(provider.encrypt("tenant-1", "secret"))
+    # Byte right after the "HYBRID:" prefix is the envelope version.
+    encrypted[len(b"HYBRID:")] = 99
+
+    with pytest.raises(ValueError, match="envelope version"):
+        provider.decrypt_with_decoding(bytes(encrypted), "tenant-1")
+
+
+@pytest.mark.usefixtures("fake_key_client")
+def test_decrypt_rejects_unrecognized_prefix() -> None:
+    provider = AzureKeyVaultKeyProvider()
+    with pytest.raises(ValueError, match="Unsupported ciphertext format"):
+        provider.decrypt_with_decoding(b"not-a-valid-envelope", "tenant-1")
+
+
+@pytest.mark.usefixtures("fake_key_client")
+def test_decrypt_rejects_truncated_envelope_missing_version_byte() -> None:
+    provider = AzureKeyVaultKeyProvider()
+    with pytest.raises(ValueError, match="Malformed Azure Key Vault envelope"):
+        provider.decrypt_with_decoding(b"HYBRID:", "tenant-1")
+
+
+@pytest.mark.usefixtures("fake_key_client")
+@pytest.mark.parametrize(
+    "truncate_at",
+    [
+        len(b"HYBRID:") + 2,  # cut inside the metadata-length prefix
+        len(b"HYBRID:") + 3,  # cut inside the metadata JSON blob
+    ],
+)
+def test_decrypt_rejects_truncated_envelope_raises_value_error(truncate_at: int) -> None:
+    provider = AzureKeyVaultKeyProvider()
+    provider.generate_key_pair("tenant-1")
+
+    encrypted = provider.encrypt("tenant-1", "secret")
+    truncated = encrypted[:truncate_at]
+
+    with pytest.raises(ValueError):
+        provider.decrypt_with_decoding(truncated, "tenant-1")
+
+
+@pytest.mark.usefixtures("fake_key_client")
+def test_decrypt_translates_disabled_key_version_into_value_error() -> None:
+    """
+    A specific key_version can become unusable after a credential was encrypted (disabled,
+    deleted, or expired in Key Vault). Callers of decrypt_token_with_decoding
+    (core/provider_manager.py, services/model_load_balancing_service.py) only catch ValueError
+    for "this credential can't be decrypted right now" -- if the underlying Azure SDK error
+    leaked through untranslated, it would crash the whole call chain (e.g. building a tenant's
+    full provider configuration just to create an unrelated new credential).
+    """
+    provider = AzureKeyVaultKeyProvider()
+    provider.generate_key_pair("tenant-1")
+
+    encrypted = provider.encrypt("tenant-1", "secret")
+    assert _embedded_key_version(encrypted) == "v1"
+
+    # Simulate the vault operator disabling/deleting the version that encrypted this credential.
+    FakeCryptographyClient.disabled_versions.add("v1")
+
+    with pytest.raises(ValueError, match="Failed to unwrap credential"):
+        provider.decrypt_with_decoding(encrypted, "tenant-1")

api/tests/unit_tests/services/test_model_load_balancing_service.py

@@ -219,7 +219,7 @@ def test_get_configs_inserts_inherit_and_filters_tenant_provider_and_source(
     )
     monkeypatch.setattr(
         "services.model_load_balancing_service.encrypter.decrypt_token_with_decoding",
-        lambda _value, _key, _cipher: "plain",
+        lambda _value, _decoding: "plain",
     )
     monkeypatch.setattr(
         "services.model_load_balancing_service.LBModelManager.get_config_in_cooldown_and_ttl",

api/uv.lock

@@ -424,6 +424,21 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/49/9a/417b3a533e01953a7c618884df2cb05a71e7b68bdbce4fbdb62349d2a2e8/azure_identity-1.25.3-py3-none-any.whl", hash = "sha256:f4d0b956a8146f30333e071374171f3cfa7bdb8073adb8c3814b65567aa7447c", size = 192138, upload-time = "2026-03-13T01:12:22.951Z" },
 ]
 
+[[package]]
+name = "azure-keyvault-keys"
+version = "4.11.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "azure-core" },
+    { name = "cryptography" },
+    { name = "isodate" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/f4/03/5ce6db28b545427d4ab572f6a4ef2a727b6b4e7bf6941cedddf98822535b/azure_keyvault_keys-4.11.1.tar.gz", hash = "sha256:90caa3a7b2c8f6b53c247ec115cf1c1dad7f107cc3aa9f35aff4838bbce7e562", size = 260915, upload-time = "2026-05-19T20:01:08.041Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/0c/3d/7bed91ae9268cf48124cf6990d8cd2c3daff7a2bb1e91a439d26ee90d705/azure_keyvault_keys-4.11.1-py3-none-any.whl", hash = "sha256:f46cdf6ee7a9baf27f70e6838327032886c8a087041dd56397773b1639da8fe2", size = 200651, upload-time = "2026-05-19T20:01:09.809Z" },
+]
+
 [[package]]
 name = "azure-storage-blob"
 version = "12.30.0"
@@ -1445,6 +1460,9 @@ dev = [
     { name = "types-ujson" },
     { name = "xinference-client" },
 ]
+kms = [
+    { name = "azure-keyvault-keys" },
+]
 storage = [
     { name = "azure-storage-blob" },
     { name = "bce-python-sdk" },
@@ -1732,6 +1750,7 @@ dev = [
     { name = "types-ujson", specifier = ">=5.10.0" },
     { name = "xinference-client", specifier = ">=2.7.0" },
 ]
+kms = [{ name = "azure-keyvault-keys", specifier = ">=4.10.0,<5.0.0" }]
 storage = [
     { name = "azure-storage-blob", specifier = ">=12.30.0,<13.0.0" },
     { name = "bce-python-sdk", specifier = "==0.9.76" },

docker/envs/core-services/shared.env.example

@@ -319,6 +319,9 @@ ARCHIVE_STORAGE_EXPORT_BUCKET=
 ARCHIVE_STORAGE_REGION=auto
 AZURE_BLOB_ACCOUNT_NAME=difyai
 AZURE_BLOB_CONTAINER_NAME=difyai-container
+AZURE_KEYVAULT_VAULT_URL=
+AZURE_KEYVAULT_KEY_SIZE=2048
+AZURE_KEYVAULT_ROTATION_INTERVAL_DAYS=
 GOOGLE_STORAGE_BUCKET_NAME=your-bucket-name
 GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64=
 ALIYUN_OSS_BUCKET_NAME=your-bucket-name
@@ -452,6 +455,7 @@ MAX_SUBMIT_COUNT=100
 
 # Vector Store Configuration
 STORAGE_TYPE=opendal
+KEY_PROVIDER_TYPE=local
 VECTOR_STORE=weaviate
 VECTOR_INDEX_NAME_PREFIX=Vector_index
 WEAVIATE_ENDPOINT=http://weaviate:8080