Spaces:
Sleeping
Sleeping
File size: 4,190 Bytes
c109c49 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | import httpx
import asyncio
import urllib.parse
import logging
from collections import OrderedDict
from fastapi import HTTPException, Request, Response
from security_utils import is_safe_url
logger = logging.getLogger(__name__)
# Bounded LRU Cache for proxied images to prevent memory exhaustion DoS
class AsyncLRUCache:
def __init__(self, capacity: int = 100):
self.capacity = capacity
self.cache = OrderedDict()
self.lock = asyncio.Lock()
async def get(self, key: str):
async with self.lock:
if key not in self.cache:
return None
self.cache.move_to_end(key)
return self.cache[key]
async def set(self, key: str, value: bytes):
async with self.lock:
self.cache[key] = value
self.cache.move_to_end(key)
if len(self.cache) > self.capacity:
self.cache.popitem(last=False)
IMAGE_CACHE = AsyncLRUCache(capacity=50)
IN_FLIGHT = {}
IN_FLIGHT_LOCK = asyncio.Lock()
async def proxy_image_logic(request: Request):
raw_query = request.url.query
if not raw_query.startswith("url="):
raise HTTPException(status_code=400, detail="Missing url query parameter.")
# Unquote twice to handle potential double-escaped query params
url = urllib.parse.unquote(urllib.parse.unquote(raw_query[4:]))
# SSRF Protection: Ensure target host is public and safe
if not is_safe_url(url):
logger.warning(f"SSRF attempt blocked. Insecure destination URL: {url}")
raise HTTPException(status_code=400, detail="Forbidden image destination URL host.")
cached_content = await IMAGE_CACHE.get(url)
if cached_content:
return Response(content=cached_content, media_type="image/jpeg")
# Guard in-flight duplicate requests
async with IN_FLIGHT_LOCK:
if url in IN_FLIGHT:
event = IN_FLIGHT[url]
else:
event = asyncio.Event()
IN_FLIGHT[url] = event
if event.is_set():
# Already fetched while we waited for lock
cached_content = await IMAGE_CACHE.get(url)
if cached_content:
return Response(content=cached_content, media_type="image/jpeg")
raise HTTPException(status_code=500, detail="Failed to fetch image.")
try:
# Check if URL belongs to pollination API prompts and trim to prevent HTTP header bloat
if "/prompt/" in url:
parts = url.split("?", 1)
base, params = parts[0], ("?" + parts[1] if len(parts) > 1 else "")
p_start = base.find("/prompt/") + 8
# Limit prompts to 600 chars to avoid command/argument length issues
base = base[:p_start] + base[p_start:][:600]
url = base + params
# Perform fetch with timeout
async with httpx.AsyncClient() as client:
for attempt in range(3):
try:
r = await client.get(url, timeout=15.0, follow_redirects=True)
if r.status_code == 200:
# Validate that returned content looks like an image
content_type = r.headers.get("Content-Type", "")
if not content_type.startswith("image/"):
logger.error(f"Target URL returned non-image content-type: {content_type}")
raise HTTPException(status_code=400, detail="Target URL is not an image.")
await IMAGE_CACHE.set(url, r.content)
return Response(content=r.content, media_type="image/jpeg")
if r.status_code == 429:
await asyncio.sleep(1)
except Exception as e:
logger.debug(f"Attempt {attempt} failed: {e}")
if attempt < 2:
await asyncio.sleep(1)
raise HTTPException(status_code=500, detail="Failed to fetch image.")
finally:
event.set()
async with IN_FLIGHT_LOCK:
if url in IN_FLIGHT:
del IN_FLIGHT[url]
|