fix: add RemoteFileUploadError for better error handling in remote fi… (#11933)

Signed-off-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
-LAN- 2024-12-21 21:22:57 +08:00 committed by GitHub
parent b8d42cdea7
commit 0b06235527
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 23 additions and 8 deletions

View File

@ -4,3 +4,8 @@ from werkzeug.exceptions import HTTPException
class FilenameNotExistsError(HTTPException): class FilenameNotExistsError(HTTPException):
code = 400 code = 400
description = "The specified filename does not exist." description = "The specified filename does not exist."
class RemoteFileUploadError(HTTPException):
code = 400
description = "Error uploading remote file."

View File

@ -7,6 +7,7 @@ from flask_restful import Resource, marshal_with, reqparse
import services import services
from controllers.common import helpers from controllers.common import helpers
from controllers.common.errors import RemoteFileUploadError
from core.file import helpers as file_helpers from core.file import helpers as file_helpers
from core.helper import ssrf_proxy from core.helper import ssrf_proxy
from fields.file_fields import file_fields_with_signed_url, remote_file_info_fields from fields.file_fields import file_fields_with_signed_url, remote_file_info_fields
@ -43,10 +44,14 @@ class RemoteFileUploadApi(Resource):
url = args["url"] url = args["url"]
resp = ssrf_proxy.head(url=url) try:
if resp.status_code != httpx.codes.OK: resp = ssrf_proxy.head(url=url)
resp = ssrf_proxy.get(url=url, timeout=3, follow_redirects=True) if resp.status_code != httpx.codes.OK:
resp.raise_for_status() resp = ssrf_proxy.get(url=url, timeout=3, follow_redirects=True)
if resp.status_code != httpx.codes.OK:
raise RemoteFileUploadError(f"Failed to fetch file from {url}: {resp.text}")
except httpx.RequestError as e:
raise RemoteFileUploadError(f"Failed to fetch file from {url}: {str(e)}")
file_info = helpers.guess_file_info_from_response(resp) file_info = helpers.guess_file_info_from_response(resp)

View File

@ -5,6 +5,7 @@ from flask_restful import marshal_with, reqparse
import services import services
from controllers.common import helpers from controllers.common import helpers
from controllers.common.errors import RemoteFileUploadError
from controllers.web.wraps import WebApiResource from controllers.web.wraps import WebApiResource
from core.file import helpers as file_helpers from core.file import helpers as file_helpers
from core.helper import ssrf_proxy from core.helper import ssrf_proxy
@ -38,10 +39,14 @@ class RemoteFileUploadApi(WebApiResource):
url = args["url"] url = args["url"]
resp = ssrf_proxy.head(url=url) try:
if resp.status_code != httpx.codes.OK: resp = ssrf_proxy.head(url=url)
resp = ssrf_proxy.get(url=url, timeout=3) if resp.status_code != httpx.codes.OK:
resp.raise_for_status() resp = ssrf_proxy.get(url=url, timeout=3, follow_redirects=True)
if resp.status_code != httpx.codes.OK:
raise RemoteFileUploadError(f"Failed to fetch file from {url}: {resp.text}")
except httpx.RequestError as e:
raise RemoteFileUploadError(f"Failed to fetch file from {url}: {str(e)}")
file_info = helpers.guess_file_info_from_response(resp) file_info = helpers.guess_file_info_from_response(resp)