mirror of
https://github.com/langgenius/dify.git
synced 2026-01-08 07:14:14 +00:00
feat: implement file extension blacklist for upload security (#27540)
This commit is contained in:
@@ -371,6 +371,12 @@ UPLOAD_IMAGE_FILE_SIZE_LIMIT=10
|
||||
UPLOAD_VIDEO_FILE_SIZE_LIMIT=100
|
||||
UPLOAD_AUDIO_FILE_SIZE_LIMIT=50
|
||||
|
||||
# Comma-separated list of file extensions blocked from upload for security reasons.
|
||||
# Extensions should be lowercase without dots (e.g., exe,bat,sh,dll).
|
||||
# Empty by default to allow all file types.
|
||||
# Recommended: exe,bat,cmd,com,scr,vbs,ps1,msi,dll
|
||||
UPLOAD_FILE_EXTENSION_BLACKLIST=
|
||||
|
||||
# Model configuration
|
||||
MULTIMODAL_SEND_FORMAT=base64
|
||||
PROMPT_GENERATION_MAX_TOKENS=512
|
||||
|
||||
@@ -331,6 +331,31 @@ class FileUploadConfig(BaseSettings):
|
||||
default=10,
|
||||
)
|
||||
|
||||
inner_UPLOAD_FILE_EXTENSION_BLACKLIST: str = Field(
|
||||
description=(
|
||||
"Comma-separated list of file extensions that are blocked from upload. "
|
||||
"Extensions should be lowercase without dots (e.g., 'exe,bat,sh,dll'). "
|
||||
"Empty by default to allow all file types."
|
||||
),
|
||||
validation_alias=AliasChoices("UPLOAD_FILE_EXTENSION_BLACKLIST"),
|
||||
default="",
|
||||
)
|
||||
|
||||
@computed_field # type: ignore[misc]
|
||||
@property
|
||||
def UPLOAD_FILE_EXTENSION_BLACKLIST(self) -> set[str]:
|
||||
"""
|
||||
Parse and return the blacklist as a set of lowercase extensions.
|
||||
Returns an empty set if no blacklist is configured.
|
||||
"""
|
||||
if not self.inner_UPLOAD_FILE_EXTENSION_BLACKLIST:
|
||||
return set()
|
||||
return {
|
||||
ext.strip().lower().strip(".")
|
||||
for ext in self.inner_UPLOAD_FILE_EXTENSION_BLACKLIST.split(",")
|
||||
if ext.strip()
|
||||
}
|
||||
|
||||
|
||||
class HttpConfig(BaseSettings):
|
||||
"""
|
||||
|
||||
@@ -25,6 +25,12 @@ class UnsupportedFileTypeError(BaseHTTPException):
|
||||
code = 415
|
||||
|
||||
|
||||
class BlockedFileExtensionError(BaseHTTPException):
|
||||
error_code = "file_extension_blocked"
|
||||
description = "The file extension is blocked for security reasons."
|
||||
code = 400
|
||||
|
||||
|
||||
class TooManyFilesError(BaseHTTPException):
|
||||
error_code = "too_many_files"
|
||||
description = "Only one file is allowed."
|
||||
|
||||
@@ -8,6 +8,7 @@ import services
|
||||
from configs import dify_config
|
||||
from constants import DOCUMENT_EXTENSIONS
|
||||
from controllers.common.errors import (
|
||||
BlockedFileExtensionError,
|
||||
FilenameNotExistsError,
|
||||
FileTooLargeError,
|
||||
NoFileUploadedError,
|
||||
@@ -83,6 +84,8 @@ class FileApi(Resource):
|
||||
raise FileTooLargeError(file_too_large_error.description)
|
||||
except services.errors.file.UnsupportedFileTypeError:
|
||||
raise UnsupportedFileTypeError()
|
||||
except services.errors.file.BlockedFileExtensionError as blocked_extension_error:
|
||||
raise BlockedFileExtensionError(blocked_extension_error.description)
|
||||
|
||||
return upload_file, 201
|
||||
|
||||
|
||||
@@ -11,3 +11,7 @@ class FileTooLargeError(BaseServiceError):
|
||||
|
||||
class UnsupportedFileTypeError(BaseServiceError):
|
||||
pass
|
||||
|
||||
|
||||
class BlockedFileExtensionError(BaseServiceError):
|
||||
description = "File extension '{extension}' is not allowed for security reasons"
|
||||
|
||||
@@ -23,7 +23,7 @@ from models import Account
|
||||
from models.enums import CreatorUserRole
|
||||
from models.model import EndUser, UploadFile
|
||||
|
||||
from .errors.file import FileTooLargeError, UnsupportedFileTypeError
|
||||
from .errors.file import BlockedFileExtensionError, FileTooLargeError, UnsupportedFileTypeError
|
||||
|
||||
PREVIEW_WORDS_LIMIT = 3000
|
||||
|
||||
@@ -59,6 +59,10 @@ class FileService:
|
||||
if len(filename) > 200:
|
||||
filename = filename.split(".")[0][:200] + "." + extension
|
||||
|
||||
# check if extension is in blacklist
|
||||
if extension and extension in dify_config.UPLOAD_FILE_EXTENSION_BLACKLIST:
|
||||
raise BlockedFileExtensionError(f"File extension '.{extension}' is not allowed for security reasons")
|
||||
|
||||
if source == "datasets" and extension not in DOCUMENT_EXTENSIONS:
|
||||
raise UnsupportedFileTypeError()
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ from configs import dify_config
|
||||
from models import Account, Tenant
|
||||
from models.enums import CreatorUserRole
|
||||
from models.model import EndUser, UploadFile
|
||||
from services.errors.file import FileTooLargeError, UnsupportedFileTypeError
|
||||
from services.errors.file import BlockedFileExtensionError, FileTooLargeError, UnsupportedFileTypeError
|
||||
from services.file_service import FileService
|
||||
|
||||
|
||||
@@ -943,3 +943,150 @@ class TestFileService:
|
||||
|
||||
# Should have the signed URL when source_url is empty
|
||||
assert upload_file2.source_url == "https://example.com/signed-url"
|
||||
|
||||
# Test file extension blacklist
|
||||
def test_upload_file_blocked_extension(
|
||||
self, db_session_with_containers, engine, mock_external_service_dependencies
|
||||
):
|
||||
"""
|
||||
Test file upload with blocked extension.
|
||||
"""
|
||||
fake = Faker()
|
||||
account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
|
||||
|
||||
# Mock blacklist configuration by patching the inner field
|
||||
with patch.object(dify_config, "inner_UPLOAD_FILE_EXTENSION_BLACKLIST", "exe,bat,sh"):
|
||||
filename = "malware.exe"
|
||||
content = b"test content"
|
||||
mimetype = "application/x-msdownload"
|
||||
|
||||
with pytest.raises(BlockedFileExtensionError):
|
||||
FileService(engine).upload_file(
|
||||
filename=filename,
|
||||
content=content,
|
||||
mimetype=mimetype,
|
||||
user=account,
|
||||
)
|
||||
|
||||
def test_upload_file_blocked_extension_case_insensitive(
|
||||
self, db_session_with_containers, engine, mock_external_service_dependencies
|
||||
):
|
||||
"""
|
||||
Test file upload with blocked extension (case insensitive).
|
||||
"""
|
||||
fake = Faker()
|
||||
account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
|
||||
|
||||
# Mock blacklist configuration by patching the inner field
|
||||
with patch.object(dify_config, "inner_UPLOAD_FILE_EXTENSION_BLACKLIST", "exe,bat"):
|
||||
# Test with uppercase extension
|
||||
filename = "malware.EXE"
|
||||
content = b"test content"
|
||||
mimetype = "application/x-msdownload"
|
||||
|
||||
with pytest.raises(BlockedFileExtensionError):
|
||||
FileService(engine).upload_file(
|
||||
filename=filename,
|
||||
content=content,
|
||||
mimetype=mimetype,
|
||||
user=account,
|
||||
)
|
||||
|
||||
def test_upload_file_not_in_blacklist(self, db_session_with_containers, engine, mock_external_service_dependencies):
|
||||
"""
|
||||
Test file upload with extension not in blacklist.
|
||||
"""
|
||||
fake = Faker()
|
||||
account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
|
||||
|
||||
# Mock blacklist configuration by patching the inner field
|
||||
with patch.object(dify_config, "inner_UPLOAD_FILE_EXTENSION_BLACKLIST", "exe,bat,sh"):
|
||||
filename = "document.pdf"
|
||||
content = b"test content"
|
||||
mimetype = "application/pdf"
|
||||
|
||||
upload_file = FileService(engine).upload_file(
|
||||
filename=filename,
|
||||
content=content,
|
||||
mimetype=mimetype,
|
||||
user=account,
|
||||
)
|
||||
|
||||
assert upload_file is not None
|
||||
assert upload_file.name == filename
|
||||
assert upload_file.extension == "pdf"
|
||||
|
||||
def test_upload_file_empty_blacklist(self, db_session_with_containers, engine, mock_external_service_dependencies):
|
||||
"""
|
||||
Test file upload with empty blacklist (default behavior).
|
||||
"""
|
||||
fake = Faker()
|
||||
account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
|
||||
|
||||
# Mock empty blacklist configuration by patching the inner field
|
||||
with patch.object(dify_config, "inner_UPLOAD_FILE_EXTENSION_BLACKLIST", ""):
|
||||
# Should allow all file types when blacklist is empty
|
||||
filename = "script.sh"
|
||||
content = b"#!/bin/bash\necho test"
|
||||
mimetype = "application/x-sh"
|
||||
|
||||
upload_file = FileService(engine).upload_file(
|
||||
filename=filename,
|
||||
content=content,
|
||||
mimetype=mimetype,
|
||||
user=account,
|
||||
)
|
||||
|
||||
assert upload_file is not None
|
||||
assert upload_file.extension == "sh"
|
||||
|
||||
def test_upload_file_multiple_blocked_extensions(
|
||||
self, db_session_with_containers, engine, mock_external_service_dependencies
|
||||
):
|
||||
"""
|
||||
Test file upload with multiple blocked extensions.
|
||||
"""
|
||||
fake = Faker()
|
||||
account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
|
||||
|
||||
# Mock blacklist with multiple extensions by patching the inner field
|
||||
blacklist_str = "exe,bat,cmd,com,scr,vbs,ps1,msi,dll"
|
||||
with patch.object(dify_config, "inner_UPLOAD_FILE_EXTENSION_BLACKLIST", blacklist_str):
|
||||
for ext in blacklist_str.split(","):
|
||||
filename = f"malware.{ext}"
|
||||
content = b"test content"
|
||||
mimetype = "application/octet-stream"
|
||||
|
||||
with pytest.raises(BlockedFileExtensionError):
|
||||
FileService(engine).upload_file(
|
||||
filename=filename,
|
||||
content=content,
|
||||
mimetype=mimetype,
|
||||
user=account,
|
||||
)
|
||||
|
||||
def test_upload_file_no_extension_with_blacklist(
|
||||
self, db_session_with_containers, engine, mock_external_service_dependencies
|
||||
):
|
||||
"""
|
||||
Test file upload with no extension when blacklist is configured.
|
||||
"""
|
||||
fake = Faker()
|
||||
account = self._create_test_account(db_session_with_containers, mock_external_service_dependencies)
|
||||
|
||||
# Mock blacklist configuration by patching the inner field
|
||||
with patch.object(dify_config, "inner_UPLOAD_FILE_EXTENSION_BLACKLIST", "exe,bat"):
|
||||
# Files with no extension should not be blocked
|
||||
filename = "README"
|
||||
content = b"test content"
|
||||
mimetype = "text/plain"
|
||||
|
||||
upload_file = FileService(engine).upload_file(
|
||||
filename=filename,
|
||||
content=content,
|
||||
mimetype=mimetype,
|
||||
user=account,
|
||||
)
|
||||
|
||||
assert upload_file is not None
|
||||
assert upload_file.extension == ""
|
||||
|
||||
Reference in New Issue
Block a user