dify/api/services/knowledge_fs/query_images.py

"""Validation and bounded loading for KnowledgeFS query-image UploadFile references."""

from __future__ import annotations

import hashlib
from collections.abc import Sequence
from dataclasses import dataclass

from core.db.session_factory import session_factory
from extensions.ext_storage import storage
from libs.datetime_utils import naive_utc_now
from models.enums import CreatorUserRole
from models.model import UploadFile
from services.file_service import FileService

QUERY_IMAGE_MAX_COUNT = 4
QUERY_IMAGE_MAX_BYTES = 10 * 1024 * 1024
QUERY_IMAGE_MAX_TOTAL_BYTES = 32 * 1024 * 1024
QUERY_IMAGE_MIME_TYPES = frozenset({"image/gif", "image/jpeg", "image/png", "image/webp"})


class KnowledgeFSQueryImageError(ValueError):
    """A safe validation error raised before query model or retrieval work starts."""

    def __init__(self, code: str, message: str) -> None:
        super().__init__(message)
        self.code = code


@dataclass(frozen=True)
class KnowledgeFSQueryImageMetadata:
    upload_file_id: str
    byte_size: int
    mime_type: str


@dataclass(frozen=True)
class KnowledgeFSResolvedQueryImage(KnowledgeFSQueryImageMetadata):
    body: bytes
    sha256: str


def validate_query_image_references(
    *,
    tenant_id: str,
    account_id: str,
    upload_file_ids: Sequence[str],
    mark_used: bool,
) -> list[KnowledgeFSQueryImageMetadata]:
    """Validate tenant/actor ownership and static bounds in one database round trip."""

    normalized_ids = _validate_reference_ids(upload_file_ids)
    if not normalized_ids:
        return []

    with session_factory.create_session() as session:
        files_by_id = FileService.get_upload_files_by_ids(tenant_id, normalized_ids, session=session)
        result: list[KnowledgeFSQueryImageMetadata] = []
        total_bytes = 0
        for upload_file_id in normalized_ids:
            upload_file = files_by_id.get(upload_file_id)
            if upload_file is None:
                raise KnowledgeFSQueryImageError("QUERY_IMAGE_NOT_FOUND", "Query image was not found")
            _assert_actor_owned(upload_file, account_id=account_id)
            mime_type = _validate_metadata(upload_file)
            total_bytes += upload_file.size
            result.append(
                KnowledgeFSQueryImageMetadata(
                    upload_file_id=upload_file_id,
                    byte_size=upload_file.size,
                    mime_type=mime_type,
                )
            )
            if mark_used:
                upload_file.used = True
                upload_file.used_by = account_id
                upload_file.used_at = naive_utc_now()

        if total_bytes > QUERY_IMAGE_MAX_TOTAL_BYTES:
            raise KnowledgeFSQueryImageError(
                "QUERY_IMAGE_TOTAL_TOO_LARGE",
                f"Query images exceed aggregate max bytes {QUERY_IMAGE_MAX_TOTAL_BYTES}",
            )
        if mark_used:
            session.commit()
        return result


def load_query_image(*, tenant_id: str, account_id: str, upload_file_id: str) -> KnowledgeFSResolvedQueryImage:
    """Load and sniff one validated image from Dify's configured object storage."""

    normalized_ids = _validate_reference_ids([upload_file_id])
    with session_factory.create_session() as session:
        upload_file = FileService.get_upload_files_by_ids(tenant_id, normalized_ids, session=session).get(
            upload_file_id
        )
        if upload_file is None:
            raise KnowledgeFSQueryImageError("QUERY_IMAGE_NOT_FOUND", "Query image was not found")
        _assert_actor_owned(upload_file, account_id=account_id)
        declared_mime_type = _validate_metadata(upload_file)
        object_key = upload_file.key
        declared_size = upload_file.size

    body = _load_bounded_body(object_key=object_key, expected_size=declared_size)
    if not body:
        raise KnowledgeFSQueryImageError("QUERY_IMAGE_EMPTY", "Query image is empty")
    detected_mime_type = _detect_image_mime_type(body)
    if detected_mime_type != declared_mime_type:
        raise KnowledgeFSQueryImageError(
            "QUERY_IMAGE_MIME_MISMATCH", "Query image content does not match its MIME type"
        )

    return KnowledgeFSResolvedQueryImage(
        upload_file_id=upload_file_id,
        byte_size=len(body),
        mime_type=detected_mime_type,
        body=body,
        sha256=hashlib.sha256(body).hexdigest(),
    )


def _load_bounded_body(*, object_key: str, expected_size: int) -> bytes:
    chunks: list[bytes] = []
    total_bytes = 0
    stream = storage.load(object_key, stream=True)
    try:
        for chunk in stream:
            if not isinstance(chunk, (bytes, bytearray, memoryview)):
                raise KnowledgeFSQueryImageError("QUERY_IMAGE_SIZE_INVALID", "Query image stream is invalid")
            body_chunk = bytes(chunk)
            total_bytes += len(body_chunk)
            if total_bytes > QUERY_IMAGE_MAX_BYTES or total_bytes > expected_size:
                raise KnowledgeFSQueryImageError(
                    "QUERY_IMAGE_SIZE_INVALID", "Query image size does not match its metadata"
                )
            chunks.append(body_chunk)
    finally:
        close = getattr(stream, "close", None)
        if callable(close):
            close()

    if total_bytes != expected_size:
        raise KnowledgeFSQueryImageError("QUERY_IMAGE_SIZE_INVALID", "Query image size does not match its metadata")
    return b"".join(chunks)


def _validate_reference_ids(upload_file_ids: Sequence[str]) -> list[str]:
    normalized = [str(value).strip() for value in upload_file_ids]
    if len(normalized) > QUERY_IMAGE_MAX_COUNT:
        raise KnowledgeFSQueryImageError(
            "QUERY_IMAGE_COUNT_EXCEEDED", f"queryImages exceeds max count {QUERY_IMAGE_MAX_COUNT}"
        )
    if any(not value for value in normalized):
        raise KnowledgeFSQueryImageError("QUERY_IMAGE_REFERENCE_INVALID", "Query image reference is invalid")
    if len(set(normalized)) != len(normalized):
        raise KnowledgeFSQueryImageError(
            "QUERY_IMAGE_REFERENCE_DUPLICATE", "queryImages must not contain duplicate uploadFileId values"
        )
    return normalized


def _assert_actor_owned(upload_file: UploadFile, *, account_id: str) -> None:
    if upload_file.created_by_role != CreatorUserRole.ACCOUNT or upload_file.created_by != account_id:
        # Deliberately use the same not-found result as an unknown id to avoid disclosing foreign files.
        raise KnowledgeFSQueryImageError("QUERY_IMAGE_NOT_FOUND", "Query image was not found")


def _validate_metadata(upload_file: UploadFile) -> str:
    mime_type = (upload_file.mime_type or "").split(";", 1)[0].strip().lower()
    if mime_type not in QUERY_IMAGE_MIME_TYPES:
        raise KnowledgeFSQueryImageError("QUERY_IMAGE_MIME_UNSUPPORTED", "Query image MIME type is not supported")
    if upload_file.size < 1:
        raise KnowledgeFSQueryImageError("QUERY_IMAGE_EMPTY", "Query image is empty")
    if upload_file.size > QUERY_IMAGE_MAX_BYTES:
        raise KnowledgeFSQueryImageError(
            "QUERY_IMAGE_TOO_LARGE", f"Query image exceeds max bytes {QUERY_IMAGE_MAX_BYTES}"
        )
    return mime_type


def _detect_image_mime_type(body: bytes) -> str:
    if body.startswith(b"\x89PNG\r\n\x1a\n"):
        return "image/png"
    if body.startswith(b"\xff\xd8\xff"):
        return "image/jpeg"
    if body.startswith((b"GIF87a", b"GIF89a")):
        return "image/gif"
    if len(body) >= 12 and body.startswith(b"RIFF") and body[8:12] == b"WEBP":
        return "image/webp"
    raise KnowledgeFSQueryImageError("QUERY_IMAGE_CONTENT_UNSUPPORTED", "Query image content is not supported")