Codex Release Bot
45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
"""Rate limits for unauthenticated child-host script job callbacks.
|
|
|
|
Uses a Lua script for atomic INCR + conditional EXPIRE:
|
|
- On first request (count==1): set TTL to start the fixed window
|
|
- On subsequent requests: do NOT reset TTL (prevents window sliding)
|
|
This ensures the rate limit window is fixed, not sliding, while also
|
|
guaranteeing the key always gets a TTL (no orphan keys if process crashes).
|
|
"""
|
|
|
|
from fastapi import HTTPException
|
|
|
|
from server.infrastructure.redis.client import get_redis
|
|
|
|
_JOB_LIMIT = 10
|
|
_IP_LIMIT = 60
|
|
_WINDOW_SECONDS = 60
|
|
|
|
# Lua script: atomic INCR + conditional EXPIRE
|
|
# Returns the new count after increment.
|
|
# Sets TTL only when the key is new (count after INCR == 1),
|
|
# so the window doesn't slide on every request.
|
|
_RATE_INCR_SCRIPT = """
|
|
local count = redis.call('INCR', KEYS[1])
|
|
if count == 1 then
|
|
redis.call('EXPIRE', KEYS[1], ARGV[1])
|
|
end
|
|
return count
|
|
"""
|
|
|
|
|
|
async def check_script_callback_rate(job_id: str, client_ip: str) -> None:
|
|
redis = get_redis()
|
|
ip = (client_ip or "unknown").strip() or "unknown"
|
|
checks = (
|
|
(f"script_cb:job:{job_id}", _JOB_LIMIT),
|
|
(f"script_cb:ip:{ip}", _IP_LIMIT),
|
|
)
|
|
for key, limit in checks:
|
|
count = await redis.eval(_RATE_INCR_SCRIPT, 1, key, str(_WINDOW_SECONDS))
|
|
if count > limit:
|
|
raise HTTPException(
|
|
status_code=429,
|
|
detail="Too many script callback requests",
|
|
)
|