diff --git a/api/controllers/console/app/app_asset.py b/api/controllers/console/app/app_asset.py index 34c8f99db5..80d7673995 100644 --- a/api/controllers/console/app/app_asset.py +++ b/api/controllers/console/app/app_asset.py @@ -1,7 +1,6 @@ -from flask import Response, request +from flask import request from flask_restx import Resource from pydantic import BaseModel, Field, field_validator -from werkzeug.exceptions import Forbidden from controllers.console import console_ns from controllers.console.app.error import ( @@ -273,33 +272,3 @@ class AppAssetFileDownloadUrlResource(Resource): return {"download_url": download_url} except ServiceNodeNotFoundError: raise AppAssetNodeNotFoundError() - - -@console_ns.route("/apps//assets/files//download") -class AppAssetFileDownloadResource(Resource): - @setup_required - @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW]) - def get(self, app_model: App, node_id: str): - timestamp = request.args.get("timestamp", "") - nonce = request.args.get("nonce", "") - sign = request.args.get("sign", "") - - if not AppAssetService.verify_download_signature( - app_id=app_model.id, - node_id=node_id, - timestamp=timestamp, - nonce=nonce, - sign=sign, - ): - raise Forbidden("Invalid or expired download link") - - try: - content, filename = AppAssetService.get_file_for_download(app_model, node_id) - except ServiceNodeNotFoundError: - raise AppAssetNodeNotFoundError() - - return Response( - content, - mimetype="application/octet-stream", - headers={"Content-Disposition": f'attachment; filename="{filename}"'}, - ) diff --git a/api/controllers/files/__init__.py b/api/controllers/files/__init__.py index f8976b86b9..5a0fbf7b1d 100644 --- a/api/controllers/files/__init__.py +++ b/api/controllers/files/__init__.py @@ -14,7 +14,7 @@ api = ExternalApi( files_ns = Namespace("files", description="File operations", path="/") -from . import image_preview, tool_files, upload +from . import image_preview, storage_download, tool_files, upload api.add_namespace(files_ns) @@ -23,6 +23,7 @@ __all__ = [ "bp", "files_ns", "image_preview", + "storage_download", "tool_files", "upload", ] diff --git a/api/controllers/files/storage_download.py b/api/controllers/files/storage_download.py new file mode 100644 index 0000000000..dfa6193a80 --- /dev/null +++ b/api/controllers/files/storage_download.py @@ -0,0 +1,56 @@ +from urllib.parse import quote, unquote + +from flask import Response, request +from flask_restx import Resource +from pydantic import BaseModel, Field +from werkzeug.exceptions import Forbidden, NotFound + +from controllers.files import files_ns +from extensions.ext_storage import storage +from extensions.storage.file_presign_storage import FilePresignStorage + +DEFAULT_REF_TEMPLATE_SWAGGER_2_0 = "#/definitions/{model}" + + +class StorageDownloadQuery(BaseModel): + timestamp: str = Field(..., description="Unix timestamp used in the signature") + nonce: str = Field(..., description="Random string for signature") + sign: str = Field(..., description="HMAC signature") + + +files_ns.schema_model( + StorageDownloadQuery.__name__, + StorageDownloadQuery.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0), +) + + +@files_ns.route("/storage//download") +class StorageFileDownloadApi(Resource): + def get(self, filename: str): + filename = unquote(filename) + + args = StorageDownloadQuery.model_validate(request.args.to_dict(flat=True)) + + if not FilePresignStorage.verify_signature( + filename=filename, + timestamp=args.timestamp, + nonce=args.nonce, + sign=args.sign, + ): + raise Forbidden("Invalid or expired download link") + + try: + generator = storage.load_stream(filename) + except FileNotFoundError: + raise NotFound("File not found") + + encoded_filename = quote(filename.split("/")[-1]) + + return Response( + generator, + mimetype="application/octet-stream", + direct_passthrough=True, + headers={ + "Content-Disposition": f"attachment; filename*=UTF-8''{encoded_filename}", + }, + ) diff --git a/api/extensions/storage/file_presign_storage.py b/api/extensions/storage/file_presign_storage.py new file mode 100644 index 0000000000..7ec0775398 --- /dev/null +++ b/api/extensions/storage/file_presign_storage.py @@ -0,0 +1,73 @@ +import base64 +import hashlib +import hmac +import os +import time +import urllib.parse +from collections.abc import Generator + +from configs import dify_config +from extensions.storage.base_storage import BaseStorage + + +class FilePresignStorage(BaseStorage): + SIGNATURE_PREFIX = "storage-download" + + def __init__(self, storage: BaseStorage): + super().__init__() + self._storage = storage + + def save(self, filename: str, data: bytes): + self._storage.save(filename, data) + + def load_once(self, filename: str) -> bytes: + return self._storage.load_once(filename) + + def load_stream(self, filename: str) -> Generator: + return self._storage.load_stream(filename) + + def download(self, filename: str, target_filepath: str): + self._storage.download(filename, target_filepath) + + def exists(self, filename: str) -> bool: + return self._storage.exists(filename) + + def delete(self, filename: str): + self._storage.delete(filename) + + def scan(self, path: str, files: bool = True, directories: bool = False) -> list[str]: + return self._storage.scan(path, files=files, directories=directories) + + def get_download_url(self, filename: str, expires_in: int = 3600) -> str: + try: + return self._storage.get_download_url(filename, expires_in) + except NotImplementedError: + return self._generate_signed_proxy_url(filename) + + def _generate_signed_proxy_url(self, filename: str) -> str: + base_url = dify_config.FILES_URL + encoded_filename = urllib.parse.quote(filename, safe="") + url = f"{base_url}/files/storage/{encoded_filename}/download" + + timestamp = str(int(time.time())) + nonce = os.urandom(16).hex() + sign = self._create_signature(filename, timestamp, nonce) + + query = urllib.parse.urlencode({"timestamp": timestamp, "nonce": nonce, "sign": sign}) + return f"{url}?{query}" + + @classmethod + def _create_signature(cls, filename: str, timestamp: str, nonce: str) -> str: + key = dify_config.SECRET_KEY.encode() + msg = f"{cls.SIGNATURE_PREFIX}|{filename}|{timestamp}|{nonce}" + sign = hmac.new(key, msg.encode(), hashlib.sha256).digest() + return base64.urlsafe_b64encode(sign).decode() + + @classmethod + def verify_signature(cls, *, filename: str, timestamp: str, nonce: str, sign: str) -> bool: + expected_sign = cls._create_signature(filename, timestamp, nonce) + if sign != expected_sign: + return False + + current_time = int(time.time()) + return current_time - int(timestamp) <= dify_config.FILES_ACCESS_TIMEOUT diff --git a/api/services/app_asset_service.py b/api/services/app_asset_service.py index d2bd3ca7b7..7d736ba080 100644 --- a/api/services/app_asset_service.py +++ b/api/services/app_asset_service.py @@ -1,15 +1,11 @@ -import base64 import hashlib -import hmac import io import logging -import time import zipfile from uuid import uuid4 from sqlalchemy.orm import Session -from configs import dify_config from core.app.entities.app_asset_entities import ( AppAssetFileTree, AppAssetNode, @@ -20,6 +16,7 @@ from core.app.entities.app_asset_entities import ( ) from extensions.ext_database import db from extensions.ext_storage import storage +from extensions.storage.file_presign_storage import FilePresignStorage from libs.datetime_utils import naive_utc_now from models.app_asset import AppAssets from models.model import App @@ -339,53 +336,5 @@ class AppAssetService: raise AppAssetNodeNotFoundError(f"File node {node_id} not found") storage_key = AppAssets.get_storage_key(app_model.tenant_id, app_model.id, node_id) - - try: - return storage.get_download_url(storage_key, expires_in) - except NotImplementedError: - raise NotImplementedError("Download URL not implemented for storage, please contact administrator") - - @staticmethod - def verify_download_signature( - *, - app_id: str, - node_id: str, - timestamp: str, - nonce: str, - sign: str, - ) -> bool: - data_to_sign = f"app-asset-download|{app_id}|{node_id}|{timestamp}|{nonce}" - secret_key = dify_config.SECRET_KEY.encode() - recalculated_sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest() - recalculated_encoded_sign = base64.urlsafe_b64encode(recalculated_sign).decode() - - if sign != recalculated_encoded_sign: - return False - - current_time = int(time.time()) - return current_time - int(timestamp) <= dify_config.FILES_ACCESS_TIMEOUT - - @staticmethod - def get_file_for_download(app_model: App, node_id: str) -> tuple[bytes, str]: - with Session(db.engine) as session: - assets = ( - session.query(AppAssets) - .filter( - AppAssets.tenant_id == app_model.tenant_id, - AppAssets.app_id == app_model.id, - AppAssets.version == AppAssets.VERSION_DRAFT, - ) - .first() - ) - if not assets: - raise AppAssetNodeNotFoundError(f"Assets not found for app {app_model.id}") - - tree = assets.asset_tree - node = tree.get(node_id) - if not node or node.node_type != AssetNodeType.FILE: - raise AppAssetNodeNotFoundError(f"File node {node_id} not found") - - storage_key = AppAssets.get_storage_key(app_model.tenant_id, app_model.id, node_id) - content = storage.load_once(storage_key) - - return content, node.name + presign_storage = FilePresignStorage(storage.storage_runner) + return presign_storage.get_download_url(storage_key, expires_in)