mirror of
https://github.com/langgenius/dify.git
synced 2026-05-09 21:28:25 +08:00
fix(commands): purge tenant tool credentials on reset-encrypt-key-pair (#35396)
The 'reset-encrypt-key-pair' CLI command rotates every tenant's asymmetric key and purges Provider / ProviderModel rows that hold LLM credentials encrypted under the old key. However it does not touch the tool provider tables that encrypt credentials under the same key: - tool_builtin_providers (BuiltinToolProvider.encrypted_credentials) - tool_api_providers (ApiToolProvider.credentials_str) - tool_mcp_providers (MCPToolProvider.encrypted_credentials/headers/server_url) Stale ciphertext in these tables surfaces as a 500 from /console/api/workspaces/current/tool-providers because the new key cannot decrypt it (#35396). This change extends the existing per-tenant purge to also delete rows from the three tool provider tables. The help text and confirmation prompt are updated to reflect the broader scope. The existing UX guarantee ('all LLM credentials will become invalid, requiring re-entry') already implies tool credentials, so this is a behavior fix rather than a contract change. Tests added under tests/unit_tests/commands/ following the existing test_upgrade_db.py pattern (mocked sessionmaker, no DB required): - aborts when EDITION is not SELF_HOSTED - purges every encrypted-credential table for a single tenant - iterates correctly across multiple tenants
This commit is contained in:
parent
f3c3534e33
commit
44562aa281
@ -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.",
|
||||
|
||||
108
api/tests/unit_tests/commands/test_reset_encrypt_key_pair.py
Normal file
108
api/tests/unit_tests/commands/test_reset_encrypt_key_pair.py
Normal 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}"
|
||||
Loading…
Reference in New Issue
Block a user