fix(ai): validate generated image result URLs (#4289)

This commit is contained in:
RaresKeY
2026-06-15 17:40:49 +03:00
committed by GitHub
parent f602819523
commit f66a23d19d
2 changed files with 116 additions and 3 deletions
+12 -3
View File
@@ -1613,7 +1613,9 @@ async def do_generate_image(content: str, session_id: Optional[str] = None, owne
"""
import base64
import httpx
import os
from pathlib import Path
from src.url_safety import check_outbound_url
lines = content.strip().split("\n")
prompt = lines[0].strip() if lines else ""
@@ -1779,8 +1781,15 @@ async def do_generate_image(content: str, session_id: Optional[str] = None, owne
elif img.get("url"):
# Download external URL and save locally (DALL-E returns temp URLs)
result_url = img["url"]
ok, reason = check_outbound_url(
result_url,
block_private=os.getenv("IMAGE_BLOCK_PRIVATE_IPS", "false").lower() == "true",
)
if not ok:
return {"error": f"Image API returned unsafe image URL: {reason}"}
try:
dl_resp = httpx.get(img["url"], timeout=60)
dl_resp = httpx.get(result_url, timeout=60)
if dl_resp.status_code == 200:
img_dir = Path(GENERATED_IMAGES_DIR)
img_dir.mkdir(parents=True, exist_ok=True)
@@ -1790,10 +1799,10 @@ async def do_generate_image(content: str, session_id: Optional[str] = None, owne
image_url = f"/api/generated-image/{filename}"
image_id = _save_to_gallery(filename)
else:
image_url = img["url"] # fallback to external URL
image_url = result_url # fallback to external URL
except Exception as _dl_e:
logger.warning(f"Failed to download DALL-E image: {_dl_e}")
image_url = img["url"] # fallback to external URL
image_url = result_url # fallback to external URL
else:
return {"error": "Image API returned unexpected format (no b64_json or url)"}