From 37f3e0ac38b77c2ce0cf17ae18146b1dda1cd4bb Mon Sep 17 00:00:00 2001
From: Joe <1264204425@qq.com>
Date: Wed, 18 Dec 2024 10:33:34 +0800
Subject: [PATCH] fix: file upload auth

---
 .github/workflows/build-push.yml |  1 +
 api/controllers/console/files.py | 99 ++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+)
 create mode 100644 api/controllers/console/files.py

diff --git a/.github/workflows/build-push.yml b/.github/workflows/build-push.yml
index 407bd47d9b..355c797af7 100644
--- a/.github/workflows/build-push.yml
+++ b/.github/workflows/build-push.yml
@@ -5,6 +5,7 @@ on:
     branches:
       - "main"
       - "deploy/dev"
+      - "fix/0.8.3-upload-auth"
   release:
     types: [published]
 
diff --git a/api/controllers/console/files.py b/api/controllers/console/files.py
new file mode 100644
index 0000000000..fe48b0f221
--- /dev/null
+++ b/api/controllers/console/files.py
@@ -0,0 +1,99 @@
+from flask import request
+from flask_login import current_user
+from flask_restful import Resource, marshal_with
+from werkzeug.exceptions import Forbidden
+
+import services
+from configs import dify_config
+from constants import DOCUMENT_EXTENSIONS
+from controllers.common.errors import FilenameNotExistsError
+from controllers.console.wraps import (
+    account_initialization_required,
+    cloud_edition_billing_resource_check,
+    setup_required,
+)
+from fields.file_fields import file_fields, upload_config_fields
+from libs.login import login_required
+from services.file_service import FileService
+
+from .error import (
+    FileTooLargeError,
+    NoFileUploadedError,
+    TooManyFilesError,
+    UnsupportedFileTypeError,
+)
+
+PREVIEW_WORDS_LIMIT = 3000
+
+
+class FileApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(upload_config_fields)
+    def get(self):
+        return {
+            "file_size_limit": dify_config.UPLOAD_FILE_SIZE_LIMIT,
+            "batch_count_limit": dify_config.UPLOAD_FILE_BATCH_LIMIT,
+            "image_file_size_limit": dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT,
+            "video_file_size_limit": dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT,
+            "audio_file_size_limit": dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT,
+            "workflow_file_upload_limit": dify_config.WORKFLOW_FILE_UPLOAD_LIMIT,
+        }, 200
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(file_fields)
+    @cloud_edition_billing_resource_check("documents")
+    def post(self):
+        file = request.files["file"]
+        source = request.form.get("source")
+
+        if "file" not in request.files:
+            raise NoFileUploadedError()
+
+        if len(request.files) > 1:
+            raise TooManyFilesError()
+
+        if not file.filename:
+            raise FilenameNotExistsError
+
+        if source == "datasets" and not current_user.is_dataset_editor:
+            raise Forbidden()
+
+        if source not in {"datasets", None}:
+            source = None
+
+        try:
+            upload_file = FileService.upload_file(
+                filename=file.filename,
+                content=file.read(),
+                mimetype=file.mimetype,
+                user=current_user,
+                source=source,
+            )
+        except services.errors.file.FileTooLargeError as file_too_large_error:
+            raise FileTooLargeError(file_too_large_error.description)
+        except services.errors.file.UnsupportedFileTypeError:
+            raise UnsupportedFileTypeError()
+
+        return upload_file, 201
+
+
+class FilePreviewApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, file_id):
+        file_id = str(file_id)
+        text = FileService.get_file_preview(file_id)
+        return {"content": text}
+
+
+class FileSupportTypeApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self):
+        return {"allowed_extensions": DOCUMENT_EXTENSIONS}