diff --git a/blueprints/api.py b/blueprints/api.py index f9785cf..2e4511b 100644 --- a/blueprints/api.py +++ b/blueprints/api.py @@ -1,5 +1,7 @@ from flask import Blueprint, request, jsonify, session, current_app -from extensions import db, redis_client +from urllib.parse import unquote, urlparse +from config import Config +from extensions import db, redis_client, s3_client from models import User, SavedPrompt from middlewares.auth import login_required from services.logger import system_logger @@ -244,12 +246,39 @@ def download_proxy(): return jsonify({"error": "缺少 URL 参数"}), 400 try: + parsed_url = urlparse(url) + public_prefix = Config.MINIO["public_url"].rstrip("/") + "/" + is_same_origin_file = ( + parsed_url.path.startswith(public_prefix) + and (not parsed_url.scheme or parsed_url.netloc == request.host) + ) + if is_same_origin_file: + object_key = unquote(parsed_url.path[len(public_prefix):]) + file_obj = s3_client.get_object(Bucket=Config.MINIO["bucket"], Key=object_key) + headers = { + "Content-Type": file_obj.get("ContentType") or "application/octet-stream", + "Content-Disposition": f'attachment; filename="{filename}"', + "Cache-Control": "private, max-age=0", + } + + def generate_minio_stream(): + try: + for chunk in file_obj["Body"].iter_chunks(chunk_size=4096): + yield chunk + finally: + file_obj["Body"].close() + + return current_app.response_class(generate_minio_stream(), headers=headers) + req, headers = get_remote_file_stream(url) headers["Content-Disposition"] = f'attachment; filename="{filename}"' def generate_stream(): - for chunk in req.iter_content(chunk_size=4096): - yield chunk + try: + for chunk in req.iter_content(chunk_size=4096): + yield chunk + finally: + req.close() return current_app.response_class(generate_stream(), headers=headers) except Exception as e: