Enhance rate limiting with robust IP detection and cleanup mechanism

app.py

@@ -14,6 +14,7 @@ import requests
 from dotenv import load_dotenv
 from PIL import Image
 from collections import defaultdict, deque
+import threading
 
 import gradio as gr
 from supabase import create_client, Client
@@ -84,6 +85,38 @@ RATE_LIMIT_REQUESTS = 30
 RATE_LIMIT_WINDOW = 3600  # 1 hour in seconds
 request_tracker = defaultdict(deque)
 
+def get_client_ip(request: gr.Request) -> str:
+    """Extract client IP with multiple fallback methods."""
+    if not request:
+        return "no-request"
+    
+    # Try multiple sources for IP address
+    ip_sources = [
+        # Direct client host
+        getattr(request.client, 'host', None) if hasattr(request, 'client') and request.client else None,
+        # Common proxy headers
+        request.headers.get('X-Forwarded-For', '').split(',')[0].strip() if hasattr(request, 'headers') else None,
+        request.headers.get('X-Real-IP', '') if hasattr(request, 'headers') else None,
+        request.headers.get('CF-Connecting-IP', '') if hasattr(request, 'headers') else None,  # Cloudflare
+        request.headers.get('True-Client-IP', '') if hasattr(request, 'headers') else None,
+        request.headers.get('X-Client-IP', '') if hasattr(request, 'headers') else None,
+    ]
+    
+    # Return first valid IP found
+    for ip in ip_sources:
+        if ip and ip.strip() and ip != '::1' and not ip.startswith('127.'):
+            return ip.strip()
+    
+    # Final fallbacks
+    if hasattr(request, 'client') and request.client:
+        client_host = getattr(request.client, 'host', None)
+        if client_host:
+            return client_host
+    
+    # If all else fails, use a session-based identifier
+    session_id = getattr(request, 'session_hash', 'unknown-session')
+    return f"session-{session_id}"
+
 def check_rate_limit(client_ip: str) -> bool:
     """Check if IP has exceeded rate limit. Returns True if allowed, False if blocked."""
     current_time = time.time()
@@ -97,8 +130,30 @@ def check_rate_limit(client_ip: str) -> bool:
     if len(user_requests) < RATE_LIMIT_REQUESTS:
         user_requests.append(current_time)
         return True
+    else:
+        # Log rate limit hit for monitoring
+        print(f"[RATE_LIMIT] IP {client_ip} exceeded limit ({len(user_requests)}/{RATE_LIMIT_REQUESTS})")
+        return False
+
+def cleanup_rate_limiter():
+    """Periodic cleanup to prevent memory issues."""
+    current_time = time.time()
+    ips_to_remove = []
+    
+    for ip, requests in request_tracker.items():
+        # Remove old requests
+        while requests and current_time - requests[0] > RATE_LIMIT_WINDOW:
+            requests.popleft()
+        
+        # If no recent requests, mark IP for removal
+        if not requests:
+            ips_to_remove.append(ip)
     
-    return False
+    # Clean up empty entries
+    for ip in ips_to_remove:
+        del request_tracker[ip]
+    
+    print(f"[RATE_LIMITER] Cleaned up {len(ips_to_remove)} inactive IPs. Active IPs: {len(request_tracker)}")
 
 # -----------------------------------------------------------------------------
 # Helper functions
@@ -227,7 +282,7 @@ def generate(base_img: Image.Image, garment_img: Image.Image, workflow_choice: s
         raise gr.Error("Please provide both images.")
     
     # Rate limiting check
-    client_ip = request.client.host if request and request.client else "unknown"
+    client_ip = get_client_ip(request)
     if not check_rate_limit(client_ip):
         raise gr.Error("Rate limit exceeded. Please try again later (30 requests per hour limit).")
 
@@ -462,6 +517,15 @@ with gr.Blocks(title="YOURMIRROR.IO - SM4LL-VTON Demo") as demo:
     product_example_2.select(lambda: Image.open("assets/product_image-2.jpg"), outputs=garment_in)
     product_example_3.select(lambda: Image.open("assets/product_image-3.jpg"), outputs=garment_in)
 
+# Periodic cleanup for rate limiter (runs every 10 minutes)
+def periodic_cleanup():
+    cleanup_rate_limiter()
+    # Schedule next cleanup
+    threading.Timer(600.0, periodic_cleanup).start()  # 10 minutes
+
+# Start cleanup timer
+threading.Timer(600.0, periodic_cleanup).start()
+
 # Run app if executed directly (e.g. `python app.py`). HF Spaces launches via
 # `python app.py` automatically if it finds `app.py` at repo root, but our file
 # lives in a sub-folder, so we keep the guard.