Spaces:
Running
Running
File size: 5,984 Bytes
f500072 | 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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 | import time
import secrets
import logging
import redis as redis_lib
from fastapi import Request, HTTPException
logger = logging.getLogger(__name__)
# Rate limits
MAX_PASTES_PER_WINDOW = 20
MAX_PASTE_BYTES_PER_WINDOW = 500 * 1024
MAX_PASTE_SIZE_BYTES = 50 * 1024
MAX_GETS_PER_WINDOW = 500
WINDOW_SECONDS = 3600
def _now_ms() -> int:
"""Current time in milliseconds."""
return int(time.time() * 1000)
def _window_start_ms() -> int:
return _now_ms() - (WINDOW_SECONDS * 1000)
def sliding_window_check(
r: redis_lib.Redis, key: str, limit: int, cost: int = 1
) -> tuple[bool, int]:
"""
Check and conditionally increment a sliding window counter.
"""
now = _now_ms()
window = _window_start_ms()
pipe = r.pipeline()
pipe.zremrangebyscore(key, "-inf", window)
pipe.zcard(key)
pipe.expire(key, WINDOW_SECONDS + 60)
results = pipe.execute()
current_count = results[1]
allowed = (current_count + cost) <= limit
if allowed:
member = f"{now}-{secrets.token_hex(8)}"
r.zadd(key, {member: now})
return allowed, current_count
def sliding_window_bytes_check(
r: redis_lib.Redis, key: str, limit_bytes: int, content_bytes: int
) -> tuple[bool, int]:
"""
Check and conditionally record bandwidth usage in a sliding window.
"""
now = _now_ms()
window = _window_start_ms()
pipe = r.pipeline()
pipe.zremrangebyscore(key, "-inf", window)
pipe.zrangebyscore(key, window, "+inf")
pipe.expire(key, WINDOW_SECONDS + 60)
results = pipe.execute()
members_in_window = results[1]
used_bytes = sum(
int(m.split(b":")[1].split(b"-")[0]) if isinstance(m, bytes) else int(m.split(":")[1].split("-")[0])
for m in members_in_window
)
allowed = (used_bytes + content_bytes) <= limit_bytes
if allowed:
member = f"{now}:{content_bytes}-{secrets.token_hex(8)}"
r.zadd(key, {member: now})
return allowed, used_bytes
def get_client_ip(req: Request) -> str:
forwarded_for = req.headers.get("X-Forwarded-For")
if forwarded_for:
return forwarded_for.split(",")[0].strip()
return req.client.host
def check_post_rate_limit(r: redis_lib.Redis, ip: str, content_bytes: int) -> dict:
"""
Enforce all POST rate limits. Returns rate limit metadata for response headers.
Raises HTTPException on violation.
"""
if content_bytes > MAX_PASTE_SIZE_BYTES:
raise HTTPException(
status_code=413,
detail={
"error": "payload_too_large",
"message": f"Individual paste cannot exceed {MAX_PASTE_SIZE_BYTES // 1024}KB.",
"limit_kb": MAX_PASTE_SIZE_BYTES // 1024,
},
)
count_key = f"rate:post:count:{ip}"
allowed_count, current = sliding_window_check(r, count_key, MAX_PASTES_PER_WINDOW)
if not allowed_count:
logger.warning(
"rate_limit_exceeded type=paste_count ip=%s current=%d", ip, current
)
raise HTTPException(
status_code=429,
detail={
"error": "rate_limit_exceeded",
"limit_type": "paste_count",
"message": f"Maximum {MAX_PASTES_PER_WINDOW} pastes per hour.",
"current": current,
"limit": MAX_PASTES_PER_WINDOW,
"retry_after_seconds": WINDOW_SECONDS,
},
headers={
"Retry-After": str(WINDOW_SECONDS),
"X-RateLimit-Limit": str(MAX_PASTES_PER_WINDOW),
"X-RateLimit-Remaining": "0",
},
)
bytes_key = f"rate:post:bytes:{ip}"
allowed_bytes, used = sliding_window_bytes_check(
r, bytes_key, MAX_PASTE_BYTES_PER_WINDOW, content_bytes
)
if not allowed_bytes:
logger.warning(
"rate_limit_exceeded type=bandwidth ip=%s used_bytes=%d", ip, used
)
raise HTTPException(
status_code=429,
detail={
"error": "rate_limit_exceeded",
"limit_type": "bandwidth",
"message": f"Maximum {MAX_PASTE_BYTES_PER_WINDOW // 1024}KB upload per hour.",
"used_bytes": used,
"limit_bytes": MAX_PASTE_BYTES_PER_WINDOW,
"retry_after_seconds": WINDOW_SECONDS,
},
headers={
"Retry-After": str(WINDOW_SECONDS),
"X-RateLimit-Limit": str(MAX_PASTE_BYTES_PER_WINDOW),
"X-RateLimit-Remaining": str(max(0, MAX_PASTE_BYTES_PER_WINDOW - used)),
},
)
return {
"X-RateLimit-Limit": str(MAX_PASTES_PER_WINDOW),
"X-RateLimit-Remaining": str(max(0, MAX_PASTES_PER_WINDOW - current - 1)),
}
def check_get_rate_limit(r: redis_lib.Redis, ip: str) -> dict:
"""
Enforce GET rate limit. Returns rate limit metadata for response headers.
Raises HTTPException on violation.
"""
key = f"rate:get:{ip}"
allowed, current = sliding_window_check(r, key, MAX_GETS_PER_WINDOW)
if not allowed:
logger.warning(
"rate_limit_exceeded type=get_count ip=%s current=%d", ip, current
)
raise HTTPException(
status_code=429,
detail={
"error": "rate_limit_exceeded",
"limit_type": "get_count",
"message": f"Maximum {MAX_GETS_PER_WINDOW} reads per hour.",
"current": current,
"limit": MAX_GETS_PER_WINDOW,
"retry_after_seconds": WINDOW_SECONDS,
},
headers={
"Retry-After": str(WINDOW_SECONDS),
"X-RateLimit-Limit": str(MAX_GETS_PER_WINDOW),
"X-RateLimit-Remaining": "0",
},
)
return {
"X-RateLimit-Limit": str(MAX_GETS_PER_WINDOW),
"X-RateLimit-Remaining": str(max(0, MAX_GETS_PER_WINDOW - current - 1)),
}
|