dify/api/services/knowledge_fs/object_storage.py

"""KnowledgeFS object operations backed by Dify's configured storage provider.

KnowledgeFS is a separate Node process, so it reaches this service through a
trusted inner API rather than importing Python storage implementations or
holding a second set of provider credentials. Logical KnowledgeFS keys are
confined below a dedicated physical namespace. Metadata is stored as a small
sidecar because Dify's common ``BaseStorage`` contract intentionally exposes
portable byte operations rather than provider-specific object metadata.

Object keys are immutable/versioned in KnowledgeFS product flows. A backend
must also support ``scan`` for bounded logical pagination and cleanup; health
fails closed when that capability is unavailable.
"""

from __future__ import annotations

import json
from base64 import b64decode, b64encode
from collections.abc import Generator, Mapping
from dataclasses import dataclass
from hashlib import sha256
from typing import NotRequired, Protocol, TypedDict

from extensions.ext_storage import storage

KNOWLEDGE_FS_OBJECT_MAX_BYTES = 64 * 1024 * 1024
_MAX_KEY_BYTES = 1024
_MAX_LIST_LIMIT = 100
_MAX_METADATA_ENTRIES = 64
_MAX_METADATA_KEY_BYTES = 64
_MAX_METADATA_VALUE_BYTES = 1024
_DATA_ROOT = "knowledge-fs/objects"
_METADATA_ROOT = "knowledge-fs/object-metadata"
_HEALTH_SCAN_ROOT = "knowledge-fs/health-capability"


class KnowledgeFSObjectStorageError(Exception):
    """Base error for the trusted KnowledgeFS storage gateway."""


class KnowledgeFSObjectStorageInvalidInputError(KnowledgeFSObjectStorageError):
    """The caller supplied an invalid key, cursor, metadata value, or limit."""


class KnowledgeFSObjectStorageTooLargeError(KnowledgeFSObjectStorageInvalidInputError):
    """The caller supplied an object body above the portable gateway limit."""


class KnowledgeFSObjectStorageChecksumError(KnowledgeFSObjectStorageError):
    """The supplied checksum does not match the uploaded object body."""


class KnowledgeFSObjectStorageCorruptError(KnowledgeFSObjectStorageError):
    """The physical object and its portable metadata sidecar disagree."""


class KnowledgeFSObjectStorageUnavailableError(KnowledgeFSObjectStorageError):
    """The selected Dify storage backend lacks a required portable capability."""


class _StorageBackend(Protocol):
    def save(self, filename: str, data: bytes) -> None: ...

    def load_once(self, filename: str) -> bytes: ...

    def load_stream(self, filename: str) -> Generator[bytes, None, None]: ...

    def exists(self, filename: str) -> bool: ...

    def delete(self, filename: str) -> None: ...

    def scan(self, path: str, files: bool = True, directories: bool = False) -> list[str]: ...


class _StoredMetadataPayload(TypedDict):
    checksum_sha256_base64: str
    content_type: str | None
    key: str
    metadata: dict[str, str]
    size_bytes: int
    version: int
    source_path: NotRequired[str]


@dataclass(frozen=True, slots=True)
class _DecodedStoredMetadata:
    metadata: KnowledgeFSObjectMetadata
    source_path: str | None = None


@dataclass(frozen=True, slots=True)
class KnowledgeFSObjectMetadata:
    """Portable metadata returned to the KnowledgeFS adapter."""

    checksum_sha256_base64: str
    content_type: str | None
    key: str
    metadata: Mapping[str, str]
    size_bytes: int


@dataclass(frozen=True, slots=True)
class KnowledgeFSObjectList:
    """One keyset-paginated page of logical KnowledgeFS objects."""

    objects: tuple[KnowledgeFSObjectMetadata, ...]
    next_cursor: str | None = None


class KnowledgeFSObjectStorageService:
    """Confine KnowledgeFS object I/O to Dify's unified storage namespace."""

    _backend: _StorageBackend

    def __init__(self, *, backend: _StorageBackend = storage) -> None:
        self._backend = backend

    def put_object(
        self,
        *,
        key: str,
        body: bytes,
        metadata: Mapping[str, str],
        checksum_sha256_base64: str | None = None,
        content_type: str | None = None,
    ) -> KnowledgeFSObjectMetadata:
        """Validate and persist one immutable logical object and its metadata sidecar."""
        normalized_key = _normalize_key(key)
        normalized_metadata = _normalize_metadata(metadata)
        normalized_content_type = _normalize_content_type(content_type)
        if len(body) > KNOWLEDGE_FS_OBJECT_MAX_BYTES:
            raise KnowledgeFSObjectStorageTooLargeError(f"object exceeds max bytes {KNOWLEDGE_FS_OBJECT_MAX_BYTES}")

        actual_checksum = b64encode(sha256(body).digest()).decode()
        if checksum_sha256_base64 is not None:
            expected_checksum = _normalize_checksum(checksum_sha256_base64)
            if expected_checksum != actual_checksum:
                raise KnowledgeFSObjectStorageChecksumError("object checksum does not match body")

        object_metadata = KnowledgeFSObjectMetadata(
            checksum_sha256_base64=actual_checksum,
            content_type=normalized_content_type,
            key=normalized_key,
            metadata=normalized_metadata,
            size_bytes=len(body),
        )
        data_path = _data_path(normalized_key)
        self._backend.save(data_path, body)
        try:
            self._backend.save(_metadata_path(normalized_key), _encode_metadata(object_metadata))
        except Exception:
            self._backend.delete(data_path)
            raise
        return object_metadata

    def adopt_upload_file(
        self,
        *,
        key: str,
        source_path: str,
        tenant_id: str,
        size_bytes: int,
        checksum_sha256_base64: str,
        metadata: Mapping[str, str],
        content_type: str | None = None,
    ) -> KnowledgeFSObjectMetadata:
        """Transfer one trusted Dify UploadFile into the KnowledgeFS logical namespace.

        The document bytes stay at their existing unified-storage path. A small logical
        marker and a v2 sidecar make the object visible to KnowledgeFS scans while reads
        stream the original object. Deletion follows the reference and removes both the
        logical marker and the transferred source object.
        """
        normalized_key = _normalize_key(key)
        normalized_source_path = _normalize_adopted_source_path(source_path, tenant_id=tenant_id)
        normalized_metadata = _normalize_metadata(metadata)
        normalized_content_type = _normalize_content_type(content_type)
        normalized_checksum = _normalize_checksum(checksum_sha256_base64)
        if not isinstance(size_bytes, int) or isinstance(size_bytes, bool) or size_bytes <= 0:
            raise KnowledgeFSObjectStorageInvalidInputError("adopted object size is invalid")
        if size_bytes > KNOWLEDGE_FS_OBJECT_MAX_BYTES:
            raise KnowledgeFSObjectStorageTooLargeError(f"object exceeds max bytes {KNOWLEDGE_FS_OBJECT_MAX_BYTES}")
        if not self._backend.exists(normalized_source_path):
            raise KnowledgeFSObjectStorageInvalidInputError("adopted upload file is missing")

        object_metadata = KnowledgeFSObjectMetadata(
            checksum_sha256_base64=normalized_checksum,
            content_type=normalized_content_type,
            key=normalized_key,
            metadata=normalized_metadata,
            size_bytes=size_bytes,
        )
        existing = self._read_stored_metadata(normalized_key)
        if existing is not None:
            if existing != _DecodedStoredMetadata(metadata=object_metadata, source_path=normalized_source_path):
                raise KnowledgeFSObjectStorageInvalidInputError("logical object already references different bytes")
            return object_metadata

        data_path = _data_path(normalized_key)
        self._backend.save(data_path, b"")
        try:
            self._backend.save(
                _metadata_path(normalized_key),
                _encode_metadata(object_metadata, source_path=normalized_source_path),
            )
        except Exception:
            self._backend.delete(data_path)
            raise
        return object_metadata

    def head_object(self, *, key: str) -> KnowledgeFSObjectMetadata | None:
        """Return portable metadata without loading the object body."""
        normalized_key = _normalize_key(key)
        if not self._backend.exists(_data_path(normalized_key)):
            return None
        stored = self._read_stored_metadata(normalized_key)
        if stored is None:
            return None
        if stored.source_path is not None and not self._backend.exists(stored.source_path):
            raise KnowledgeFSObjectStorageCorruptError("adopted object source is missing")
        return stored.metadata

    def load_stream(self, *, key: str) -> Generator[bytes, None, None] | None:
        """Return the configured backend's stream for one logical object."""
        normalized_key = _normalize_key(key)
        stored = self._read_stored_metadata(normalized_key)
        if stored is None:
            return None
        source_path = stored.source_path or _data_path(normalized_key)
        if not self._backend.exists(source_path):
            raise KnowledgeFSObjectStorageCorruptError("object body is missing")
        return self._backend.load_stream(source_path)

    def delete_object(self, *, key: str) -> None:
        """Idempotently delete an object and its portable metadata sidecar."""
        normalized_key = _normalize_key(key)
        stored = self._read_stored_metadata(normalized_key)
        if stored is not None and stored.source_path is not None:
            self._backend.delete(stored.source_path)
        self._backend.delete(_data_path(normalized_key))
        self._backend.delete(_metadata_path(normalized_key))

    def _read_stored_metadata(self, normalized_key: str) -> _DecodedStoredMetadata | None:
        if not self._backend.exists(_data_path(normalized_key)):
            return None
        try:
            encoded = self._backend.load_once(_metadata_path(normalized_key))
        except FileNotFoundError as exc:
            raise KnowledgeFSObjectStorageCorruptError("object metadata sidecar is missing") from exc
        stored = _decode_metadata(encoded)
        if stored.metadata.key != normalized_key:
            raise KnowledgeFSObjectStorageCorruptError("object metadata key does not match")
        return stored

    def list_objects(
        self,
        *,
        prefix: str,
        limit: int,
        cursor: str | None = None,
    ) -> KnowledgeFSObjectList:
        """List logical keys using a stable lexical cursor and portable metadata."""
        normalized_prefix = _normalize_prefix(prefix)
        if not isinstance(limit, int) or isinstance(limit, bool) or not 1 <= limit <= _MAX_LIST_LIMIT:
            raise KnowledgeFSObjectStorageInvalidInputError("list limit must be between 1 and 100")
        normalized_cursor = _normalize_key(cursor) if cursor is not None else None
        if normalized_cursor is not None and not normalized_cursor.startswith(normalized_prefix):
            raise KnowledgeFSObjectStorageInvalidInputError("list cursor must be within prefix")

        try:
            physical_paths = self._backend.scan(
                _physical_scan_prefix(normalized_prefix),
                files=True,
                directories=False,
            )
        except FileNotFoundError:
            physical_paths = []
        except NotImplementedError as exc:
            raise KnowledgeFSObjectStorageUnavailableError(
                "configured Dify storage backend does not support object listing"
            ) from exc
        except Exception as exc:
            raise KnowledgeFSObjectStorageUnavailableError("Dify storage listing failed") from exc

        data_prefix = f"{_DATA_ROOT}/"
        logical_keys = sorted(
            {
                physical_path[len(data_prefix) :]
                for physical_path in physical_paths
                if physical_path.startswith(data_prefix)
                and physical_path[len(data_prefix) :].startswith(normalized_prefix)
                and (normalized_cursor is None or physical_path[len(data_prefix) :] > normalized_cursor)
            }
        )
        page_keys = logical_keys[: limit + 1]
        has_more = len(page_keys) > limit
        objects: list[KnowledgeFSObjectMetadata] = []
        for logical_key in page_keys[:limit]:
            metadata = self.head_object(key=logical_key)
            if metadata is not None:
                objects.append(metadata)
        next_cursor = objects[-1].key if has_more and objects else None
        return KnowledgeFSObjectList(objects=tuple(objects), next_cursor=next_cursor)

    def health(self) -> bool:
        """Check the portable list capability without scanning the object namespace."""
        try:
            self._backend.scan(_HEALTH_SCAN_ROOT, files=True, directories=False)
        except FileNotFoundError:
            return True
        except Exception:
            return False
        return True


def _normalize_key(value: str) -> str:
    if not isinstance(value, str):
        raise KnowledgeFSObjectStorageInvalidInputError("object key must be a string")
    normalized = value.strip()
    segments = normalized.split("/")
    if (
        not normalized
        or normalized.startswith("/")
        or "\\" in normalized
        or "\x00" in normalized
        or any(segment in {"", ".", ".."} for segment in segments)
        or len(normalized.encode()) > _MAX_KEY_BYTES
    ):
        raise KnowledgeFSObjectStorageInvalidInputError("object key is invalid")
    return normalized


def _normalize_prefix(value: str) -> str:
    if value == "":
        return value
    normalized = value.strip()
    candidate = normalized.removesuffix("/")
    _normalize_key(candidate)
    return normalized


def _normalize_metadata(metadata: Mapping[str, str]) -> dict[str, str]:
    if len(metadata) > _MAX_METADATA_ENTRIES:
        raise KnowledgeFSObjectStorageInvalidInputError("object metadata has too many entries")
    normalized: dict[str, str] = {}
    for key, value in metadata.items():
        if (
            not isinstance(key, str)
            or not isinstance(value, str)
            or not key
            or len(key.encode()) > _MAX_METADATA_KEY_BYTES
            or len(value.encode()) > _MAX_METADATA_VALUE_BYTES
        ):
            raise KnowledgeFSObjectStorageInvalidInputError("object metadata is invalid")
        normalized[key] = value
    return normalized


def _normalize_content_type(value: str | None) -> str | None:
    if value is None:
        return None
    normalized = value.strip()
    if not normalized or len(normalized) > 255 or "\r" in normalized or "\n" in normalized:
        raise KnowledgeFSObjectStorageInvalidInputError("object content type is invalid")
    return normalized


def _normalize_checksum(value: str) -> str:
    try:
        decoded = b64decode(value, validate=True)
    except ValueError as exc:
        raise KnowledgeFSObjectStorageInvalidInputError("object checksum is invalid") from exc
    if len(decoded) != 32 or b64encode(decoded).decode() != value:
        raise KnowledgeFSObjectStorageInvalidInputError("object checksum is invalid")
    return value


def _normalize_adopted_source_path(value: str, *, tenant_id: str) -> str:
    if not isinstance(value, str) or not isinstance(tenant_id, str):
        raise KnowledgeFSObjectStorageInvalidInputError("adopted upload path is invalid")
    normalized_tenant_id = tenant_id.strip()
    normalized = value.strip()
    expected_prefix = f"upload_files/{normalized_tenant_id}/"
    relative = normalized.removeprefix(expected_prefix)
    if (
        not normalized_tenant_id
        or not normalized.startswith(expected_prefix)
        or not relative
        or "/" in relative
        or "\\" in normalized
        or "\x00" in normalized
        or relative in {".", ".."}
    ):
        raise KnowledgeFSObjectStorageInvalidInputError("adopted upload path is invalid")
    return normalized


def _data_path(key: str) -> str:
    return f"{_DATA_ROOT}/{key}"


def _metadata_path(key: str) -> str:
    digest = sha256(key.encode()).hexdigest()
    return f"{_METADATA_ROOT}/{digest[:2]}/{digest[2:4]}/{digest}.json"


def _physical_scan_prefix(prefix: str) -> str:
    parent = prefix.rstrip("/").rpartition("/")[0]
    return f"{_DATA_ROOT}/{parent}" if parent else _DATA_ROOT


def _encode_metadata(metadata: KnowledgeFSObjectMetadata, *, source_path: str | None = None) -> bytes:
    payload: _StoredMetadataPayload = {
        "checksum_sha256_base64": metadata.checksum_sha256_base64,
        "content_type": metadata.content_type,
        "key": metadata.key,
        "metadata": dict(metadata.metadata),
        "size_bytes": metadata.size_bytes,
        "version": 2 if source_path is not None else 1,
    }
    if source_path is not None:
        payload["source_path"] = source_path
    return json.dumps(payload, separators=(",", ":"), sort_keys=True).encode()


def _decode_metadata(encoded: bytes) -> _DecodedStoredMetadata:
    try:
        payload = json.loads(encoded)
        if not isinstance(payload, dict) or payload.get("version") not in {1, 2}:
            raise ValueError
        key = _normalize_key(payload["key"])
        checksum = _normalize_checksum(payload["checksum_sha256_base64"])
        content_type = _normalize_content_type(payload["content_type"])
        metadata = _normalize_metadata(payload["metadata"])
        size_bytes = payload["size_bytes"]
        if not isinstance(size_bytes, int) or isinstance(size_bytes, bool) or size_bytes < 0:
            raise ValueError
        source_path = payload.get("source_path")
        if payload["version"] == 1:
            if source_path is not None:
                raise ValueError
        else:
            tenant_id = metadata.get("tenant_id")
            if not isinstance(source_path, str) or not isinstance(tenant_id, str):
                raise ValueError
            source_path = _normalize_adopted_source_path(source_path, tenant_id=tenant_id)
    except (KeyError, TypeError, ValueError, json.JSONDecodeError, KnowledgeFSObjectStorageInvalidInputError) as exc:
        raise KnowledgeFSObjectStorageCorruptError("object metadata sidecar is invalid") from exc
    return _DecodedStoredMetadata(
        metadata=KnowledgeFSObjectMetadata(
            checksum_sha256_base64=checksum,
            content_type=content_type,
            key=key,
            metadata=metadata,
            size_bytes=size_bytes,
        ),
        source_path=source_path,
    )