251 lines
10 KiB
Python
251 lines
10 KiB
Python
"""Transactional KnowledgeFS authorization effects of Dify Workspace membership changes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from enum import StrEnum
|
|
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.orm import Session
|
|
|
|
from libs.datetime_utils import naive_utc_now
|
|
from models.knowledge_fs import (
|
|
KnowledgeFSAuthorizationRevision,
|
|
KnowledgeFSControlSpace,
|
|
KnowledgeFSControlSpacePermission,
|
|
KnowledgeFSControlSpacePermissionRole,
|
|
KnowledgeFSControlSpacePermissionStatus,
|
|
)
|
|
from services.knowledge_fs.revocation_commands import KnowledgeFSRevocationCommandProducer
|
|
|
|
_MAX_SAFE_INTEGER = 2**53 - 1
|
|
|
|
|
|
class KnowledgeFSWorkspaceMembershipChange(StrEnum):
|
|
MEMBER_ADDED = "member_added"
|
|
MEMBER_REMOVED = "member_removed"
|
|
ROLE_CHANGED = "role_changed"
|
|
|
|
|
|
class KnowledgeFSMembershipChangeInvariantError(RuntimeError):
|
|
"""Workspace membership state cannot be applied safely to KnowledgeFS authorization."""
|
|
|
|
|
|
def knowledge_fs_membership_schema_installed(session: Session) -> bool:
|
|
"""Support additive rollout and lightweight tests before the control-plane table exists."""
|
|
return sa.inspect(session.connection()).has_table(KnowledgeFSControlSpace.__tablename__)
|
|
|
|
|
|
def apply_workspace_membership_change(
|
|
*,
|
|
session: Session,
|
|
tenant_id: str,
|
|
actor_account_id: str,
|
|
account_ids: Sequence[str],
|
|
change: KnowledgeFSWorkspaceMembershipChange,
|
|
removed_account_id: str | None = None,
|
|
replacement_owner_account_id: str | None = None,
|
|
) -> None:
|
|
"""Advance every Space epoch and append revokes in the caller's membership transaction."""
|
|
normalized_tenant_id = _identifier(tenant_id, "tenant_id")
|
|
normalized_actor_id = _identifier(actor_account_id, "actor_account_id")
|
|
normalized_accounts = tuple(dict.fromkeys(_identifier(value, "account_id") for value in account_ids))
|
|
if not normalized_accounts:
|
|
raise KnowledgeFSMembershipChangeInvariantError("Membership change requires an affected account")
|
|
normalized_removed_id = _optional_identifier(removed_account_id, "removed_account_id")
|
|
normalized_replacement_id = _optional_identifier(
|
|
replacement_owner_account_id,
|
|
"replacement_owner_account_id",
|
|
)
|
|
if change is KnowledgeFSWorkspaceMembershipChange.MEMBER_REMOVED:
|
|
if normalized_removed_id not in normalized_accounts or normalized_replacement_id is None:
|
|
raise KnowledgeFSMembershipChangeInvariantError(
|
|
"Member removal requires the removed account and replacement Workspace owner"
|
|
)
|
|
if normalized_removed_id == normalized_replacement_id:
|
|
raise KnowledgeFSMembershipChangeInvariantError("Removed member cannot remain the replacement owner")
|
|
elif normalized_removed_id is not None or normalized_replacement_id is not None:
|
|
raise KnowledgeFSMembershipChangeInvariantError("Owner replacement is valid only for member removal")
|
|
|
|
if not knowledge_fs_membership_schema_installed(session):
|
|
return
|
|
spaces, revisions = _lock_workspace_authorization(session, normalized_tenant_id)
|
|
if not spaces:
|
|
return
|
|
|
|
producer = KnowledgeFSRevocationCommandProducer()
|
|
for space in spaces:
|
|
revision = revisions[space.id]
|
|
if revision.membership_epoch >= _MAX_SAFE_INTEGER:
|
|
raise KnowledgeFSMembershipChangeInvariantError("KnowledgeFS membership epoch is exhausted")
|
|
revision.membership_epoch += 1
|
|
|
|
acl_changed = False
|
|
if normalized_removed_id is not None and normalized_replacement_id is not None:
|
|
acl_changed = _remove_space_member(
|
|
session=session,
|
|
space=space,
|
|
removed_account_id=normalized_removed_id,
|
|
replacement_owner_account_id=normalized_replacement_id,
|
|
actor_account_id=normalized_actor_id,
|
|
)
|
|
if acl_changed:
|
|
if revision.space_acl_epoch >= _MAX_SAFE_INTEGER:
|
|
raise KnowledgeFSMembershipChangeInvariantError("KnowledgeFS Space ACL epoch is exhausted")
|
|
revision.space_acl_epoch += 1
|
|
|
|
if change is KnowledgeFSWorkspaceMembershipChange.MEMBER_ADDED:
|
|
continue
|
|
reason_code = (
|
|
"workspace_membership_removed"
|
|
if change is KnowledgeFSWorkspaceMembershipChange.MEMBER_REMOVED
|
|
else "workspace_role_changed"
|
|
)
|
|
for account_id in normalized_accounts:
|
|
producer.enqueue_principal_grants(
|
|
session=session,
|
|
tenant_id=normalized_tenant_id,
|
|
control_space_id=space.id,
|
|
subject=f"dify-account:{account_id}",
|
|
reason_code=reason_code,
|
|
caller_kinds=("interactive",),
|
|
)
|
|
|
|
|
|
def apply_workspace_rbac_role_change(*, session: Session, tenant_id: str) -> None:
|
|
"""Advance Workspace membership epochs and revoke every interactive grant.
|
|
|
|
Custom Enterprise RBAC role edits can affect any member bound to the role,
|
|
while Dify deliberately does not mirror those bindings locally. The safe
|
|
invalidation therefore scans durable grant provenance per control-space
|
|
instead of relying on a potentially incomplete member list.
|
|
"""
|
|
normalized_tenant_id = _identifier(tenant_id, "tenant_id")
|
|
if not knowledge_fs_membership_schema_installed(session):
|
|
return
|
|
spaces, revisions = _lock_workspace_authorization(session, normalized_tenant_id)
|
|
if not spaces:
|
|
return
|
|
|
|
producer = KnowledgeFSRevocationCommandProducer()
|
|
for space in spaces:
|
|
revision = revisions[space.id]
|
|
if revision.membership_epoch >= _MAX_SAFE_INTEGER:
|
|
raise KnowledgeFSMembershipChangeInvariantError("KnowledgeFS membership epoch is exhausted")
|
|
revision.membership_epoch += 1
|
|
producer.enqueue_control_space_grants(
|
|
session=session,
|
|
tenant_id=normalized_tenant_id,
|
|
control_space_id=space.id,
|
|
reason_code="workspace_rbac_role_changed",
|
|
caller_kinds=("interactive",),
|
|
)
|
|
|
|
|
|
def _lock_workspace_authorization(
|
|
session: Session,
|
|
tenant_id: str,
|
|
) -> tuple[tuple[KnowledgeFSControlSpace, ...], dict[str, KnowledgeFSAuthorizationRevision]]:
|
|
spaces = tuple(
|
|
session.scalars(
|
|
sa.select(KnowledgeFSControlSpace)
|
|
.where(KnowledgeFSControlSpace.tenant_id == tenant_id)
|
|
.order_by(KnowledgeFSControlSpace.id)
|
|
.with_for_update()
|
|
)
|
|
)
|
|
if not spaces:
|
|
return (), {}
|
|
revisions = {
|
|
revision.control_space_id: revision
|
|
for revision in session.scalars(
|
|
sa.select(KnowledgeFSAuthorizationRevision)
|
|
.where(
|
|
KnowledgeFSAuthorizationRevision.tenant_id == tenant_id,
|
|
KnowledgeFSAuthorizationRevision.control_space_id.in_(space.id for space in spaces),
|
|
)
|
|
.order_by(KnowledgeFSAuthorizationRevision.control_space_id)
|
|
.with_for_update()
|
|
)
|
|
}
|
|
if len(revisions) != len(spaces):
|
|
raise KnowledgeFSMembershipChangeInvariantError("KnowledgeFS authorization revision is missing")
|
|
return spaces, revisions
|
|
|
|
|
|
def _remove_space_member(
|
|
*,
|
|
session: Session,
|
|
space: KnowledgeFSControlSpace,
|
|
removed_account_id: str,
|
|
replacement_owner_account_id: str,
|
|
actor_account_id: str,
|
|
) -> bool:
|
|
permissions = {
|
|
permission.account_id: permission
|
|
for permission in session.scalars(
|
|
sa.select(KnowledgeFSControlSpacePermission)
|
|
.where(
|
|
KnowledgeFSControlSpacePermission.tenant_id == space.tenant_id,
|
|
KnowledgeFSControlSpacePermission.control_space_id == space.id,
|
|
KnowledgeFSControlSpacePermission.account_id.in_((removed_account_id, replacement_owner_account_id)),
|
|
)
|
|
.with_for_update()
|
|
)
|
|
}
|
|
changed = False
|
|
now = naive_utc_now()
|
|
removed_permission = permissions.get(removed_account_id)
|
|
if removed_permission is not None and (removed_permission.status is KnowledgeFSControlSpacePermissionStatus.ACTIVE):
|
|
removed_permission.status = KnowledgeFSControlSpacePermissionStatus.REVOKED
|
|
removed_permission.revision += 1
|
|
removed_permission.revoked_at = now
|
|
removed_permission.revoked_by_account_id = actor_account_id
|
|
changed = True
|
|
|
|
if space.owner_account_id != removed_account_id:
|
|
return changed
|
|
space.owner_account_id = replacement_owner_account_id
|
|
space.resource_version += 1
|
|
replacement_permission = permissions.get(replacement_owner_account_id)
|
|
if replacement_permission is None:
|
|
session.add(
|
|
KnowledgeFSControlSpacePermission(
|
|
tenant_id=space.tenant_id,
|
|
control_space_id=space.id,
|
|
account_id=replacement_owner_account_id,
|
|
role=KnowledgeFSControlSpacePermissionRole.OWNER,
|
|
granted_by_account_id=actor_account_id,
|
|
)
|
|
)
|
|
elif (
|
|
replacement_permission.role is not KnowledgeFSControlSpacePermissionRole.OWNER
|
|
or replacement_permission.status is not KnowledgeFSControlSpacePermissionStatus.ACTIVE
|
|
):
|
|
replacement_permission.role = KnowledgeFSControlSpacePermissionRole.OWNER
|
|
replacement_permission.status = KnowledgeFSControlSpacePermissionStatus.ACTIVE
|
|
replacement_permission.revision += 1
|
|
replacement_permission.revoked_at = None
|
|
replacement_permission.revoked_by_account_id = None
|
|
return True
|
|
|
|
|
|
def _identifier(value: str, field: str) -> str:
|
|
normalized = value.strip()
|
|
if not normalized or normalized != value or len(normalized) > 255:
|
|
raise KnowledgeFSMembershipChangeInvariantError(f"{field} must be a normalized identifier")
|
|
return normalized
|
|
|
|
|
|
def _optional_identifier(value: str | None, field: str) -> str | None:
|
|
return None if value is None else _identifier(value, field)
|
|
|
|
|
|
__all__ = [
|
|
"KnowledgeFSMembershipChangeInvariantError",
|
|
"KnowledgeFSWorkspaceMembershipChange",
|
|
"apply_workspace_membership_change",
|
|
"apply_workspace_rbac_role_change",
|
|
"knowledge_fs_membership_schema_installed",
|
|
]
|