From 823ae03a0884e1fc0c9fbcc17540f97bfce63d20 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Thu, 7 Nov 2024 14:35:58 +0800 Subject: [PATCH] fix(remote-files): fallback to get when remote server not support head method (#10370) --- api/controllers/console/error.py | 24 ++++++++++ .../console/{files/__init__.py => files.py} | 2 +- api/controllers/console/files/errors.py | 25 ----------- api/controllers/console/remote_files.py | 44 ++++++++++++------- api/controllers/web/remote_files.py | 43 ++++++++++-------- 5 files changed, 77 insertions(+), 61 deletions(-) rename api/controllers/console/{files/__init__.py => files.py} (99%) delete mode 100644 api/controllers/console/files/errors.py diff --git a/api/controllers/console/error.py b/api/controllers/console/error.py index ed6a99a017..e0630ca66c 100644 --- a/api/controllers/console/error.py +++ b/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 diff --git a/api/controllers/console/files/__init__.py b/api/controllers/console/files.py similarity index 99% rename from api/controllers/console/files/__init__.py rename to api/controllers/console/files.py index 6c7bd8acfd..946d3db37f 100644 --- a/api/controllers/console/files/__init__.py +++ b/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, diff --git a/api/controllers/console/files/errors.py b/api/controllers/console/files/errors.py deleted file mode 100644 index 1654ef2cf4..0000000000 --- a/api/controllers/console/files/errors.py +++ /dev/null @@ -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 diff --git a/api/controllers/console/remote_files.py b/api/controllers/console/remote_files.py index 42d6e25416..9b899bef64 100644 --- a/api/controllers/console/remote_files.py +++ b/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, diff --git a/api/controllers/web/remote_files.py b/api/controllers/web/remote_files.py index cf36ae302d..d6b8eb2855 100644 --- a/api/controllers/web/remote_files.py +++ b/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,