This commit is contained in:
Tim Ren 2026-05-09 04:35:52 +00:00 committed by GitHub
commit ebafc32e32
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 122 additions and 3 deletions

View File

@ -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.",

View File

@ -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}"