Spaces:

alpgul
/

proxy

Sleeping

proxy / app.py

alpgul

fix: Correctly extract client IP for rate limiting

03b52dc 11 months ago

5.71 kB


	from fastapi import FastAPI, HTTPException, BackgroundTasks, Request
	from fastapi.responses import HTMLResponse, StreamingResponse
	from curl_cffi import requests
	import os
	import time
	import hashlib
	import threading
	from dotenv import load_dotenv
	from collections import defaultdict

	load_dotenv()

	app = FastAPI()

	# --- Rate Limiting Configuration ---
	RATE_LIMIT_REQUESTS = 20
	RATE_LIMIT_WINDOW = 60 # 60 seconds
	rate_limit_data = defaultdict(list)

	# --- Cache Configuration ---
	CACHE_DIR = "cache"
	CACHE_MAX_SIZE = 100 * 1024 * 1024 # 100 MB
	CACHE_MAX_AGE = 3600 # 1 hour in seconds
	MAX_FILE_SIZE = 10 * 1024 * 1024 # 10 MB
	CACHE_CLEANUP_INTERVAL = 600 # 10 minutes

	# --- Cache Initialization ---
	if not os.path.exists(CACHE_DIR):
	os.makedirs(CACHE_DIR)

	# --- Cache Cleanup Functions ---
	def clean_cache():
	"""Cleans the cache directory based on size and age.
	This function is designed to be run in a background thread.
	"""
	while True:
	try:
	total_size = 0
	files = []
	for filename in os.listdir(CACHE_DIR):
	filepath = os.path.join(CACHE_DIR, filename)
	if not os.path.isfile(filepath):
	continue

	stat = os.stat(filepath)
	total_size += stat.st_size
	files.append((filepath, stat.st_mtime, stat.st_size))

	# 1. Remove expired files
	now = time.time()
	for filepath, mtime, _ in files:
	if now - mtime > CACHE_MAX_AGE:
	os.remove(filepath)
	print(f"Removed expired cache file: {filepath}")

	# 2. Enforce total cache size limit (LRU)
	if total_size > CACHE_MAX_SIZE:
	# Sort files by last modification time (oldest first)
	files.sort(key=lambda x: x[1])
	while total_size > CACHE_MAX_SIZE and files:
	filepath, _, size = files.pop(0)
	os.remove(filepath)
	total_size -= size
	print(f"Removed old cache file to free space: {filepath}")

	except Exception as e:
	print(f"Error during cache cleanup: {e}")

	time.sleep(CACHE_CLEANUP_INTERVAL)

	# --- Background Thread for Cache Cleanup ---
	cache_cleanup_thread = threading.Thread(target=clean_cache, daemon=True)
	cache_cleanup_thread.start()

	# --- API Endpoints ---
	@app.get("/", response_class=HTMLResponse)
	async def root():
	return """
	<html>
	<head>
	<title>Proxy Usage</title>
	</head>
	<body>
	<h1>Proxy Usage</h1>
	<p>To use the proxy, make a GET request to the following URL:</p>
	<p><code>/{api_key}?url={your_url}</code></p>
	<p>Replace <code>{api_key}</code> with your API key and <code>{your_url}</code> with the URL you want to proxy.</p>
	<h2>Example with cURL:</h2>
	<p><code>curl -X GET "http://your-hugging-face-space-url/your_api_key?url=https://example.com"</code></p>
	</body>
	</html>
	"""

	@app.get("/{api_key}")
	async def proxy(api_key: str, url: str, request: Request, background_tasks: BackgroundTasks):
	# --- Rate Limiting Logic ---
	client_ip = request.headers.get("X-Forwarded-For")
	if client_ip:
	client_ip = client_ip.split(',')[0].strip()
	else:
	client_ip = request.headers.get("X-Real-IP")
	if not client_ip:
	client_ip = request.client.host
	now = time.time()
	# Remove old timestamps
	rate_limit_data[client_ip] = [t for t in rate_limit_data[client_ip] if now - t < RATE_LIMIT_WINDOW]
	# Check request limit
	if len(rate_limit_data[client_ip]) >= RATE_LIMIT_REQUESTS:
	raise HTTPException(status_code=429, detail="Too Many Requests")
	# Add current timestamp
	rate_limit_data[client_ip].append(now)

	# --- API Key Check ---
	expected_api_key = os.environ.get("API_KEY")
	if not expected_api_key or api_key != expected_api_key:
	raise HTTPException(status_code=401, detail="Invalid API key")

	url_hash = hashlib.md5(url.encode()).hexdigest()
	cache_path = os.path.join(CACHE_DIR, url_hash)

	# Check for valid cached file
	if os.path.exists(cache_path):
	if now - os.path.getmtime(cache_path) < CACHE_MAX_AGE:
	def file_iterator(file_path, chunk_size=8192):
	with open(file_path, "rb") as f:
	while True:
	chunk = f.read(chunk_size)
	if not chunk:
	break
	yield chunk
	return StreamingResponse(file_iterator(cache_path))

	# If not cached or expired, fetch from URL
	try:
	response = requests.get(url, impersonate="chrome110", stream=True, timeout=15)
	response.raise_for_status() # Raise an exception for bad status codes

	async def stream_and_cache():
	total_size = 0
	do_cache = True
	temp_cache_path = cache_path + ".tmp"

	with open(temp_cache_path, "wb") as f:
	for chunk in response.iter_content(chunk_size=8192):
	total_size += len(chunk)
	if total_size > MAX_FILE_SIZE:
	do_cache = False
	if do_cache:
	f.write(chunk)
	yield chunk

	if do_cache:
	os.rename(temp_cache_path, cache_path)
	else:
	os.remove(temp_cache_path)

	return StreamingResponse(stream_and_cache())

	except Exception as e:
	raise HTTPException(status_code=500, detail=str(e))