mirror of
https://github.com/langgenius/dify.git
synced 2026-05-21 01:07:03 +08:00
93 lines
2.8 KiB
Python
93 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Generator, Iterable
|
|
from contextlib import contextmanager
|
|
from contextvars import ContextVar
|
|
from dataclasses import dataclass, field, replace
|
|
|
|
from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom
|
|
|
|
_current_file_access_scope: ContextVar[FileAccessScope | None] = ContextVar(
|
|
"current_file_access_scope",
|
|
default=None,
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class FileAccessScope:
|
|
"""Request-scoped ownership context used by workflow-layer file lookups.
|
|
|
|
``granted_upload_file_ids`` is execution-local: callers may add upload files
|
|
that were returned by trusted retrieval paths without changing persistent
|
|
ownership markers.
|
|
|
|
``granted_retriever_segment_ids`` gates lazy attachment loading by segment
|
|
ID, so user-provided context cannot make a later LLM node load arbitrary
|
|
same-tenant knowledge attachments.
|
|
"""
|
|
|
|
tenant_id: str
|
|
user_id: str
|
|
user_from: UserFrom
|
|
invoke_from: InvokeFrom
|
|
granted_upload_file_ids: frozenset[str] = field(default_factory=frozenset)
|
|
granted_retriever_segment_ids: frozenset[str] = field(default_factory=frozenset)
|
|
|
|
@property
|
|
def requires_user_ownership(self) -> bool:
|
|
return self.user_from == UserFrom.END_USER
|
|
|
|
|
|
def get_current_file_access_scope() -> FileAccessScope | None:
|
|
return _current_file_access_scope.get()
|
|
|
|
|
|
def grant_upload_file_access(upload_file_ids: Iterable[str]) -> None:
|
|
scope = _current_file_access_scope.get()
|
|
if scope is None:
|
|
return
|
|
|
|
granted_upload_file_ids = frozenset(str(file_id) for file_id in upload_file_ids if file_id)
|
|
if not granted_upload_file_ids:
|
|
return
|
|
|
|
_current_file_access_scope.set(
|
|
replace(
|
|
scope,
|
|
granted_upload_file_ids=scope.granted_upload_file_ids | granted_upload_file_ids,
|
|
)
|
|
)
|
|
|
|
|
|
def grant_retriever_segment_access(segment_ids: Iterable[str]) -> None:
|
|
scope = _current_file_access_scope.get()
|
|
if scope is None:
|
|
return
|
|
|
|
granted_segment_ids = frozenset(str(segment_id) for segment_id in segment_ids if segment_id)
|
|
if not granted_segment_ids:
|
|
return
|
|
|
|
_current_file_access_scope.set(
|
|
replace(
|
|
scope,
|
|
granted_retriever_segment_ids=scope.granted_retriever_segment_ids | granted_segment_ids,
|
|
)
|
|
)
|
|
|
|
|
|
def is_retriever_segment_access_granted(segment_id: str) -> bool:
|
|
scope = _current_file_access_scope.get()
|
|
if scope is None or not scope.requires_user_ownership:
|
|
return True
|
|
return str(segment_id) in scope.granted_retriever_segment_ids
|
|
|
|
|
|
@contextmanager
|
|
def bind_file_access_scope(scope: FileAccessScope) -> Generator[None, None, None]:
|
|
token = _current_file_access_scope.set(scope)
|
|
try:
|
|
yield
|
|
finally:
|
|
_current_file_access_scope.reset(token)
|