fix(remote-files): fallback to get when remote server not support head method (#10370)

api/controllers/console/error.py

@@ -62,3 +62,27 @@ class EmailSendIpLimitError(BaseHTTPException):
     error_code = "email_send_ip_limit"
     description = "Too many emails have been sent from this IP address recently. Please try again later."
     code = 429
+
+
+class FileTooLargeError(BaseHTTPException):
+    error_code = "file_too_large"
+    description = "File size exceeded. {message}"
+    code = 413
+
+
+class UnsupportedFileTypeError(BaseHTTPException):
+    error_code = "unsupported_file_type"
+    description = "File type not allowed."
+    code = 415
+
+
+class TooManyFilesError(BaseHTTPException):
+    error_code = "too_many_files"
+    description = "Only one file is allowed."
+    code = 400
+
+
+class NoFileUploadedError(BaseHTTPException):
+    error_code = "no_file_uploaded"
+    description = "Please upload your file."
+    code = 400

api/controllers/console/files/__init__.py → api/controllers/console/files.py

@@ -15,7 +15,7 @@ from fields.file_fields import file_fields, upload_config_fields
 from libs.login import login_required
 from services.file_service import FileService
 
-from .errors import (
+from .error import (
     FileTooLargeError,
     NoFileUploadedError,
     TooManyFilesError,

api/controllers/console/files/errors.py

@@ -1,25 +0,0 @@
-from libs.exception import BaseHTTPException
-
-
-class FileTooLargeError(BaseHTTPException):
-    error_code = "file_too_large"
-    description = "File size exceeded. {message}"
-    code = 413
-
-
-class UnsupportedFileTypeError(BaseHTTPException):
-    error_code = "unsupported_file_type"
-    description = "File type not allowed."
-    code = 415
-
-
-class TooManyFilesError(BaseHTTPException):
-    error_code = "too_many_files"
-    description = "Only one file is allowed."
-    code = 400
-
-
-class NoFileUploadedError(BaseHTTPException):
-    error_code = "no_file_uploaded"
-    description = "Please upload your file."
-    code = 400

api/controllers/console/remote_files.py

@@ -1,9 +1,11 @@
 import urllib.parse
 from typing import cast
 
+import httpx
 from flask_login import current_user
 from flask_restful import Resource, marshal_with, reqparse
 
+import services
 from controllers.common import helpers
 from core.file import helpers as file_helpers
 from core.helper import ssrf_proxy
@@ -11,19 +13,25 @@ from fields.file_fields import file_fields_with_signed_url, remote_file_info_fie
 from models.account import Account
 from services.file_service import FileService
 
+from .error import (
+    FileTooLargeError,
+    UnsupportedFileTypeError,
+)
+
 
 class RemoteFileInfoApi(Resource):
     @marshal_with(remote_file_info_fields)
     def get(self, url):
         decoded_url = urllib.parse.unquote(url)
-        try:
-            response = ssrf_proxy.head(decoded_url)
-            return {
-                "file_type": response.headers.get("Content-Type", "application/octet-stream"),
-                "file_length": int(response.headers.get("Content-Length", 0)),
-            }
-        except Exception as e:
-            return {"error": str(e)}, 400
+        resp = ssrf_proxy.head(decoded_url)
+        if resp.status_code != httpx.codes.OK:
+            # failed back to get method
+            resp = ssrf_proxy.get(decoded_url, timeout=3)
+        resp.raise_for_status()
+        return {
+            "file_type": resp.headers.get("Content-Type", "application/octet-stream"),
+            "file_length": int(resp.headers.get("Content-Length", 0)),
+        }
 
 
 class RemoteFileUploadApi(Resource):
@@ -35,17 +43,17 @@ class RemoteFileUploadApi(Resource):
 
         url = args["url"]
 
-        response = ssrf_proxy.head(url)
-        response.raise_for_status()
+        resp = ssrf_proxy.head(url=url)
+        if resp.status_code != httpx.codes.OK:
+            resp = ssrf_proxy.get(url=url, timeout=3)
+        resp.raise_for_status()
 
-        file_info = helpers.guess_file_info_from_response(response)
+        file_info = helpers.guess_file_info_from_response(resp)
 
         if not FileService.is_file_size_within_limit(extension=file_info.extension, file_size=file_info.size):
-            return {"error": "File size exceeded"}, 400
+            raise FileTooLargeError
 
-        response = ssrf_proxy.get(url)
-        response.raise_for_status()
-        content = response.content
+        content = resp.content if resp.request.method == "GET" else ssrf_proxy.get(url).content
 
         try:
             user = cast(Account, current_user)
@@ -56,8 +64,10 @@ class RemoteFileUploadApi(Resource):
                 user=user,
                 source_url=url,
             )
-        except Exception as e:
-            return {"error": str(e)}, 400
+        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 {
             "id": upload_file.id,

api/controllers/web/remote_files.py

@@ -1,7 +1,9 @@
 import urllib.parse
 
+import httpx
 from flask_restful import marshal_with, reqparse
 
+import services
 from controllers.common import helpers
 from controllers.web.wraps import WebApiResource
 from core.file import helpers as file_helpers
@@ -9,19 +11,22 @@ from core.helper import ssrf_proxy
 from fields.file_fields import file_fields_with_signed_url, remote_file_info_fields
 from services.file_service import FileService
 
+from .error import FileTooLargeError, UnsupportedFileTypeError
+
 
 class RemoteFileInfoApi(WebApiResource):
     @marshal_with(remote_file_info_fields)
     def get(self, app_model, end_user, url):
         decoded_url = urllib.parse.unquote(url)
-        try:
-            response = ssrf_proxy.head(decoded_url)
-            return {
-                "file_type": response.headers.get("Content-Type", "application/octet-stream"),
-                "file_length": int(response.headers.get("Content-Length", -1)),
-            }
-        except Exception as e:
-            return {"error": str(e)}, 400
+        resp = ssrf_proxy.head(decoded_url)
+        if resp.status_code != httpx.codes.OK:
+            # failed back to get method
+            resp = ssrf_proxy.get(decoded_url, timeout=3)
+        resp.raise_for_status()
+        return {
+            "file_type": resp.headers.get("Content-Type", "application/octet-stream"),
+            "file_length": int(resp.headers.get("Content-Length", -1)),
+        }
 
 
 class RemoteFileUploadApi(WebApiResource):
@@ -33,28 +38,30 @@ class RemoteFileUploadApi(WebApiResource):
 
         url = args["url"]
 
-        response = ssrf_proxy.head(url)
-        response.raise_for_status()
+        resp = ssrf_proxy.head(url=url)
+        if resp.status_code != httpx.codes.OK:
+            resp = ssrf_proxy.get(url=url, timeout=3)
+        resp.raise_for_status()
 
-        file_info = helpers.guess_file_info_from_response(response)
+        file_info = helpers.guess_file_info_from_response(resp)
 
         if not FileService.is_file_size_within_limit(extension=file_info.extension, file_size=file_info.size):
-            return {"error": "File size exceeded"}, 400
+            raise FileTooLargeError
 
-        response = ssrf_proxy.get(url)
-        response.raise_for_status()
-        content = response.content
+        content = resp.content if resp.request.method == "GET" else ssrf_proxy.get(url).content
 
         try:
             upload_file = FileService.upload_file(
                 filename=file_info.filename,
                 content=content,
                 mimetype=file_info.mimetype,
-                user=end_user,  # Use end_user instead of current_user
+                user=end_user,
                 source_url=url,
             )
-        except Exception as e:
-            return {"error": str(e)}, 400
+        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 {
             "id": upload_file.id,