Factorize rate limit check

This commit is contained in:
Alexandre Aubin 2019-07-25 11:07:39 +02:00
parent 7d868ec5dd
commit dea0590318

View file

@ -32,6 +32,21 @@ def clear_rate_limit_db(now):
del RATE_LIMIT_DB[key] del RATE_LIMIT_DB[key]
def check_rate_limit(key, now):
if key in RATE_LIMIT_DB:
since_last_attempt = now - RATE_LIMIT_DB[key]
if since_last_attempt < RATE_LIMIT_SECONDS:
logger.info(f"Rate limit reached for {key}, can retry in {int(RATE_LIMIT_SECONDS - since_last_attempt)} seconds")
return json_response({
"status": "error",
"code": "error_rate_limit",
"content": f"Rate limit reached for this domain or ip, retry in {int(RATE_LIMIT_SECONDS - since_last_attempt)} seconds",
}, status=400)
RATE_LIMIT_DB[key] = time.time()
async def query_dns(host, dns_entry_type): async def query_dns(host, dns_entry_type):
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
dns_resolver = aiodns.DNSResolver(loop=loop) dns_resolver = aiodns.DNSResolver(loop=loop)
@ -65,32 +80,24 @@ async def check_http(request):
- answer saying if the domain can be reached - answer saying if the domain can be reached
""" """
# this is supposed to be a fast operation if run enough # this is supposed to be a fast operation if run often enough
now = time.time() now = time.time()
clear_rate_limit_db(now) clear_rate_limit_db(now)
ip = request.ip ip = request.ip
if ip in RATE_LIMIT_DB: check_rate_limit_ip = check_rate_limit(ip, now)
since_last_attempt = now - RATE_LIMIT_DB[ip] if check_rate_limit_ip:
if since_last_attempt < RATE_LIMIT_SECONDS: return check_rate_limit_ip
logger.info(f"Rate limite {ip}, can retry in {int(RATE_LIMIT_SECONDS - since_last_attempt)} seconds")
return json_response({
"status": "error",
"code": "error_rate_limit",
"content": f"Rate limit on ip, retry in {int(RATE_LIMIT_SECONDS - since_last_attempt)} seconds",
}, status=400)
RATE_LIMIT_DB[ip] = time.time()
try: try:
data = request.json data = request.json
except InvalidUsage: except InvalidUsage:
logger.info(f"Unvalid json in request, body is : {request.body}") logger.info(f"Invalid json in request, body is : {request.body}")
return json_response({ return json_response({
"status": "error", "status": "error",
"code": "error_bad_json", "code": "error_bad_json",
"content": "InvalidUsage, body isn't proper json", "content": "Invalid usage, body isn't proper json",
}, status=400) }, status=400)
if not data or "domain" not in data: if not data or "domain" not in data:
@ -98,22 +105,14 @@ async def check_http(request):
return json_response({ return json_response({
"status": "error", "status": "error",
"code": "error_no_domain", "code": "error_no_domain",
"content": "request must specify a domain", "content": "Request must specify a domain",
}, status=400) }, status=400)
domain = data["domain"] domain = data["domain"]
if domain in RATE_LIMIT_DB: check_rate_limit_domain = check_rate_limit(domain, now)
since_last_attempt = now - RATE_LIMIT_DB[domain] if check_rate_limit_domain:
if since_last_attempt < RATE_LIMIT_SECONDS: return check_rate_limit_domain
logger.info(f"Rate limite {domain}, can retry in {int(RATE_LIMIT_SECONDS - since_last_attempt)} seconds")
return json_response({
"status": "error",
"code": "error_rate_limit",
"content": f"Rate limit on domain, retry in {int(RATE_LIMIT_SECONDS - since_last_attempt)} seconds",
}, status=400)
RATE_LIMIT_DB[domain] = time.time()
if not validators.domain(domain): if not validators.domain(domain):
logger.info(f"Invalid request, is not in the right format (domain is : {domain})") logger.info(f"Invalid request, is not in the right format (domain is : {domain})")