fix: dos in annotation import (#29470)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
@@ -1,6 +1,6 @@
|
||||
from typing import Any, Literal
|
||||
|
||||
from flask import request
|
||||
from flask import abort, request
|
||||
from flask_restx import Resource, fields, marshal, marshal_with
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
||||
@@ -8,6 +8,8 @@ from controllers.common.errors import NoFileUploadedError, TooManyFilesError
|
||||
from controllers.console import console_ns
|
||||
from controllers.console.wraps import (
|
||||
account_initialization_required,
|
||||
annotation_import_concurrency_limit,
|
||||
annotation_import_rate_limit,
|
||||
cloud_edition_billing_resource_check,
|
||||
edit_permission_required,
|
||||
setup_required,
|
||||
@@ -314,18 +316,25 @@ class AnnotationUpdateDeleteApi(Resource):
|
||||
@console_ns.route("/apps/<uuid:app_id>/annotations/batch-import")
|
||||
class AnnotationBatchImportApi(Resource):
|
||||
@console_ns.doc("batch_import_annotations")
|
||||
@console_ns.doc(description="Batch import annotations from CSV file")
|
||||
@console_ns.doc(description="Batch import annotations from CSV file with rate limiting and security checks")
|
||||
@console_ns.doc(params={"app_id": "Application ID"})
|
||||
@console_ns.response(200, "Batch import started successfully")
|
||||
@console_ns.response(403, "Insufficient permissions")
|
||||
@console_ns.response(400, "No file uploaded or too many files")
|
||||
@console_ns.response(413, "File too large")
|
||||
@console_ns.response(429, "Too many requests or concurrent imports")
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@cloud_edition_billing_resource_check("annotation")
|
||||
@annotation_import_rate_limit
|
||||
@annotation_import_concurrency_limit
|
||||
@edit_permission_required
|
||||
def post(self, app_id):
|
||||
from configs import dify_config
|
||||
|
||||
app_id = str(app_id)
|
||||
|
||||
# check file
|
||||
if "file" not in request.files:
|
||||
raise NoFileUploadedError()
|
||||
@@ -335,9 +344,27 @@ class AnnotationBatchImportApi(Resource):
|
||||
|
||||
# get file from request
|
||||
file = request.files["file"]
|
||||
|
||||
# check file type
|
||||
if not file.filename or not file.filename.lower().endswith(".csv"):
|
||||
raise ValueError("Invalid file type. Only CSV files are allowed")
|
||||
|
||||
# Check file size before processing
|
||||
file.seek(0, 2) # Seek to end of file
|
||||
file_size = file.tell()
|
||||
file.seek(0) # Reset to beginning
|
||||
|
||||
max_size_bytes = dify_config.ANNOTATION_IMPORT_FILE_SIZE_LIMIT * 1024 * 1024
|
||||
if file_size > max_size_bytes:
|
||||
abort(
|
||||
413,
|
||||
f"File size exceeds maximum limit of {dify_config.ANNOTATION_IMPORT_FILE_SIZE_LIMIT}MB. "
|
||||
f"Please reduce the file size and try again.",
|
||||
)
|
||||
|
||||
if file_size == 0:
|
||||
raise ValueError("The uploaded file is empty")
|
||||
|
||||
return AppAnnotationService.batch_import_app_annotations(app_id, file)
|
||||
|
||||
|
||||
|
||||
@@ -331,3 +331,91 @@ def is_admin_or_owner_required(f: Callable[P, R]):
|
||||
return f(*args, **kwargs)
|
||||
|
||||
return decorated_function
|
||||
|
||||
|
||||
def annotation_import_rate_limit(view: Callable[P, R]):
|
||||
"""
|
||||
Rate limiting decorator for annotation import operations.
|
||||
|
||||
Implements sliding window rate limiting with two tiers:
|
||||
- Short-term: Configurable requests per minute (default: 5)
|
||||
- Long-term: Configurable requests per hour (default: 20)
|
||||
|
||||
Uses Redis ZSET for distributed rate limiting across multiple instances.
|
||||
"""
|
||||
|
||||
@wraps(view)
|
||||
def decorated(*args: P.args, **kwargs: P.kwargs):
|
||||
_, current_tenant_id = current_account_with_tenant()
|
||||
current_time = int(time.time() * 1000)
|
||||
|
||||
# Check per-minute rate limit
|
||||
minute_key = f"annotation_import_rate_limit:{current_tenant_id}:1min"
|
||||
redis_client.zadd(minute_key, {current_time: current_time})
|
||||
redis_client.zremrangebyscore(minute_key, 0, current_time - 60000)
|
||||
minute_count = redis_client.zcard(minute_key)
|
||||
redis_client.expire(minute_key, 120) # 2 minutes TTL
|
||||
|
||||
if minute_count > dify_config.ANNOTATION_IMPORT_RATE_LIMIT_PER_MINUTE:
|
||||
abort(
|
||||
429,
|
||||
f"Too many annotation import requests. Maximum {dify_config.ANNOTATION_IMPORT_RATE_LIMIT_PER_MINUTE} "
|
||||
f"requests per minute allowed. Please try again later.",
|
||||
)
|
||||
|
||||
# Check per-hour rate limit
|
||||
hour_key = f"annotation_import_rate_limit:{current_tenant_id}:1hour"
|
||||
redis_client.zadd(hour_key, {current_time: current_time})
|
||||
redis_client.zremrangebyscore(hour_key, 0, current_time - 3600000)
|
||||
hour_count = redis_client.zcard(hour_key)
|
||||
redis_client.expire(hour_key, 7200) # 2 hours TTL
|
||||
|
||||
if hour_count > dify_config.ANNOTATION_IMPORT_RATE_LIMIT_PER_HOUR:
|
||||
abort(
|
||||
429,
|
||||
f"Too many annotation import requests. Maximum {dify_config.ANNOTATION_IMPORT_RATE_LIMIT_PER_HOUR} "
|
||||
f"requests per hour allowed. Please try again later.",
|
||||
)
|
||||
|
||||
return view(*args, **kwargs)
|
||||
|
||||
return decorated
|
||||
|
||||
|
||||
def annotation_import_concurrency_limit(view: Callable[P, R]):
|
||||
"""
|
||||
Concurrency control decorator for annotation import operations.
|
||||
|
||||
Limits the number of concurrent import tasks per tenant to prevent
|
||||
resource exhaustion and ensure fair resource allocation.
|
||||
|
||||
Uses Redis ZSET to track active import jobs with automatic cleanup
|
||||
of stale entries (jobs older than 2 minutes).
|
||||
"""
|
||||
|
||||
@wraps(view)
|
||||
def decorated(*args: P.args, **kwargs: P.kwargs):
|
||||
_, current_tenant_id = current_account_with_tenant()
|
||||
current_time = int(time.time() * 1000)
|
||||
|
||||
active_jobs_key = f"annotation_import_active:{current_tenant_id}"
|
||||
|
||||
# Clean up stale entries (jobs that should have completed or timed out)
|
||||
stale_threshold = current_time - 120000 # 2 minutes ago
|
||||
redis_client.zremrangebyscore(active_jobs_key, 0, stale_threshold)
|
||||
|
||||
# Check current active job count
|
||||
active_count = redis_client.zcard(active_jobs_key)
|
||||
|
||||
if active_count >= dify_config.ANNOTATION_IMPORT_MAX_CONCURRENT:
|
||||
abort(
|
||||
429,
|
||||
f"Too many concurrent import tasks. Maximum {dify_config.ANNOTATION_IMPORT_MAX_CONCURRENT} "
|
||||
f"concurrent imports allowed per workspace. Please wait for existing imports to complete.",
|
||||
)
|
||||
|
||||
# Allow the request to proceed
|
||||
# The actual job registration will happen in the service layer
|
||||
return view(*args, **kwargs)
|
||||
|
||||
return decorated
|
||||
|
||||
Reference in New Issue
Block a user