80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
import hashlib
|
|
import time
|
|
from django.conf import settings
|
|
from django.core.cache import cache
|
|
from django.http import JsonResponse
|
|
|
|
|
|
class RateLimitMiddleware:
|
|
def __init__(self, get_response):
|
|
self.get_response = get_response
|
|
|
|
def __call__(self, request):
|
|
if not getattr(settings, "RATE_LIMIT_ENABLE", False):
|
|
return self.get_response(request)
|
|
|
|
path = request.path or ""
|
|
rate_limit_paths = getattr(settings, "RATE_LIMIT_PATHS", ["/api/"])
|
|
if not any(path.startswith(prefix) for prefix in rate_limit_paths):
|
|
return self.get_response(request)
|
|
|
|
window = int(getattr(settings, "RATE_LIMIT_WINDOW_SECONDS", 60))
|
|
max_requests = int(getattr(settings, "RATE_LIMIT_REQUESTS", 120))
|
|
now = int(time.time())
|
|
window_key = now // max(window, 1)
|
|
|
|
ident = request.META.get("HTTP_X_FORWARDED_FOR", "").split(",")[0].strip()
|
|
if not ident:
|
|
ident = request.META.get("REMOTE_ADDR", "unknown")
|
|
cache_key = f"rate:{ident}:{window_key}"
|
|
|
|
try:
|
|
current = cache.incr(cache_key)
|
|
except ValueError:
|
|
cache.add(cache_key, 1, timeout=window)
|
|
current = 1
|
|
|
|
if current > max_requests:
|
|
return JsonResponse(
|
|
{"message": "请求过于频繁，请稍后再试", "success": False},
|
|
status=429,
|
|
)
|
|
|
|
return self.get_response(request)
|
|
|
|
|
|
class ETagMiddleware:
|
|
def __init__(self, get_response):
|
|
self.get_response = get_response
|
|
|
|
def __call__(self, request):
|
|
response = self.get_response(request)
|
|
|
|
if request.method not in {"GET", "HEAD"}:
|
|
return response
|
|
if response.status_code != 200:
|
|
return response
|
|
if response.has_header("ETag"):
|
|
return response
|
|
|
|
content_type = response.get("Content-Type", "")
|
|
if "application/json" not in content_type:
|
|
return response
|
|
|
|
content = getattr(response, "content", b"") or b""
|
|
max_bytes = int(getattr(settings, "ETAG_MAX_BYTES", 2 * 1024 * 1024))
|
|
if len(content) > max_bytes:
|
|
return response
|
|
|
|
etag = hashlib.sha256(content).hexdigest()
|
|
etag_value = f'W/"{etag}"'
|
|
response["ETag"] = etag_value
|
|
response["Cache-Control"] = "private, max-age=0"
|
|
|
|
if_none_match = request.META.get("HTTP_IF_NONE_MATCH")
|
|
if if_none_match and if_none_match == etag_value:
|
|
response.status_code = 304
|
|
response.content = b""
|
|
|
|
return response
|