diff --git a/api/commands/system.py b/api/commands/system.py index 39b2e991ed..7755d3b5bc 100644 --- a/api/commands/system.py +++ b/api/commands/system.py @@ -14,6 +14,7 @@ from libs.rsa import generate_key_pair from models import Tenant from models.model import App, AppMode, Conversation from models.provider import Provider, ProviderModel +from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider logger = logging.getLogger(__name__) @@ -23,13 +24,16 @@ DB_UPGRADE_LOCK_TTL_SECONDS = 60 @click.command( "reset-encrypt-key-pair", help="Reset the asymmetric key pair of workspace for encrypt LLM credentials. " - "After the reset, all LLM credentials will become invalid, " - "requiring re-entry." + "After the reset, all LLM credentials and tool provider credentials " + "(builtin / API / MCP) will be purged, requiring re-entry. " "Only support SELF_HOSTED mode.", ) @click.confirmation_option( prompt=click.style( - "Are you sure you want to reset encrypt key pair? This operation cannot be rolled back!", fg="red" + "Are you sure you want to reset encrypt key pair? " + "This will also purge builtin / API / MCP tool provider records for every tenant. " + "This operation cannot be rolled back!", + fg="red", ) ) def reset_encrypt_key_pair(): @@ -53,6 +57,13 @@ def reset_encrypt_key_pair(): session.execute(delete(Provider).where(Provider.provider_type == "custom", Provider.tenant_id == tenant.id)) session.execute(delete(ProviderModel).where(ProviderModel.tenant_id == tenant.id)) + # Purge tool provider records that hold credentials encrypted under the + # tenant key. Leaving them in place causes /console/api/workspaces/current/ + # tool-providers to 500 because decryption fails on stale ciphertext (#35396). + session.execute(delete(BuiltinToolProvider).where(BuiltinToolProvider.tenant_id == tenant.id)) + session.execute(delete(ApiToolProvider).where(ApiToolProvider.tenant_id == tenant.id)) + session.execute(delete(MCPToolProvider).where(MCPToolProvider.tenant_id == tenant.id)) + click.echo( click.style( f"Congratulations! The asymmetric key pair of workspace {tenant.id} has been reset.", diff --git a/api/tests/unit_tests/commands/test_reset_encrypt_key_pair.py b/api/tests/unit_tests/commands/test_reset_encrypt_key_pair.py new file mode 100644 index 0000000000..31b4d71d0f --- /dev/null +++ b/api/tests/unit_tests/commands/test_reset_encrypt_key_pair.py @@ -0,0 +1,108 @@ +"""Unit tests for the reset-encrypt-key-pair CLI command (#35396). + +The command must purge every table that stores ciphertext encrypted with the +tenant's asymmetric key, otherwise stale rows cause downstream API failures +such as `/console/api/workspaces/current/tool-providers` returning 500. +""" + +from unittest.mock import MagicMock, patch + +import commands +from commands import system as system_commands +from models.provider import Provider, ProviderModel +from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider + + +def _invoke_reset() -> int: + try: + commands.reset_encrypt_key_pair.callback() + except SystemExit as e: + return int(e.code or 0) + return 0 + + +def _delete_targets(session_mock: MagicMock) -> list: + """Extract the model class targeted by each `delete(...)` call on the session.""" + targets = [] + for call in session_mock.execute.call_args_list: + stmt = call.args[0] + # `delete(Foo)` constructs a `Delete` statement whose entity is `Foo`. + try: + targets.append(stmt.table.name) + except AttributeError: + targets.append(repr(stmt)) + return targets + + +def test_reset_aborts_when_not_self_hosted(monkeypatch, capsys): + monkeypatch.setattr(system_commands.dify_config, "EDITION", "CLOUD") + + exit_code = _invoke_reset() + captured = capsys.readouterr() + + assert exit_code == 0 + assert "only for SELF_HOSTED" in captured.out + + +def test_reset_purges_provider_and_tool_tables_for_each_tenant(monkeypatch, capsys): + """The command must purge LLM provider rows AND every tool provider table + that stores ciphertext encrypted under the tenant key (#35396).""" + monkeypatch.setattr(system_commands.dify_config, "EDITION", "SELF_HOSTED") + monkeypatch.setattr(system_commands, "generate_key_pair", lambda tenant_id: f"new-key-{tenant_id}") + + fake_tenant = MagicMock(id="tenant-abc", encrypt_public_key="old-key") + session = MagicMock() + session.scalars.return_value.all.return_value = [fake_tenant] + + fake_sessionmaker = MagicMock() + fake_sessionmaker.begin.return_value.__enter__.return_value = session + fake_sessionmaker.begin.return_value.__exit__.return_value = False + + with ( + patch.object(system_commands, "db", MagicMock()), + patch.object(system_commands, "sessionmaker", return_value=fake_sessionmaker), + ): + exit_code = _invoke_reset() + + captured = capsys.readouterr() + assert exit_code == 0 + assert "tenant-abc" in captured.out + + # New key pair generated and assigned. + assert fake_tenant.encrypt_public_key == "new-key-tenant-abc" + + # Every encrypted-credential table should have been purged for this tenant. + table_names = _delete_targets(session) + expected = { + Provider.__tablename__, + ProviderModel.__tablename__, + BuiltinToolProvider.__tablename__, + ApiToolProvider.__tablename__, + MCPToolProvider.__tablename__, + } + assert expected.issubset(set(table_names)), f"missing purges: expected {expected}, got {table_names}" + + +def test_reset_iterates_all_tenants(monkeypatch, capsys): + """Multi-tenant deployments must purge every tenant, not just the first.""" + monkeypatch.setattr(system_commands.dify_config, "EDITION", "SELF_HOSTED") + monkeypatch.setattr(system_commands, "generate_key_pair", lambda tenant_id: f"new-key-{tenant_id}") + + tenants = [MagicMock(id=f"tenant-{i}", encrypt_public_key="old") for i in range(3)] + session = MagicMock() + session.scalars.return_value.all.return_value = tenants + + fake_sessionmaker = MagicMock() + fake_sessionmaker.begin.return_value.__enter__.return_value = session + fake_sessionmaker.begin.return_value.__exit__.return_value = False + + with ( + patch.object(system_commands, "db", MagicMock()), + patch.object(system_commands, "sessionmaker", return_value=fake_sessionmaker), + ): + _invoke_reset() + + # Five purges per tenant × 3 tenants = 15 execute calls. + assert session.execute.call_count == 15 + for tenant in tenants: + assert tenant.encrypt_public_key == f"new-key-{tenant.id}"