Skip to content

Commit

Permalink
Configurable Authorization Request Limits (#9)
Browse files Browse the repository at this point in the history
* Make authorization requests per minute param configurable with comparable defaults (6/60 from 3/30)
Document configuration option in README.md

* Add time to wait to HTTP response

---------

Co-authored-by: Daniel McKnight <daniel@neon.ai>
  • Loading branch information
NeonDaniel and NeonDaniel authored Jan 26, 2024
1 parent e0f534a commit 63fd997
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 6 deletions.
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ hana:
access_token_ttl: 86400 # 1 day
refresh_token_ttl: 604800 # 1 week
requests_per_minute: 60
auth_requests_per_minute: 6 # This counts valid and invalid requests from an IP address
access_token_secret: a800445648142061fc238d1f84e96200da87f4f9fa7835cac90db8b4391b117b
refresh_token_secret: 833d369ac73d883123743a44b4a7fe21203cffc956f4c8fec712e71aafa8e1aa
fastapi_title: "My HANA API Host"
Expand Down
18 changes: 12 additions & 6 deletions neon_hana/auth/client_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ def __init__(self, config: dict):
self._access_secret = config.get("access_token_secret")
self._refresh_secret = config.get("refresh_token_secret")
self._rpm = config.get("requests_per_minute", 60)
self._auth_rpm = config.get("auth_requests_per_minute", 6)
self._disable_auth = config.get("disable_auth")
self._jwt_algo = "HS256"

Expand All @@ -69,14 +70,19 @@ def check_auth_request(self, client_id: str, username: str,
print(f"Using cached client: {self.authorized_clients[client_id]}")
return self.authorized_clients[client_id]

if not self.rate_limiter.get_all_buckets(f"auth{origin_ip}"):
self.rate_limiter.add_bucket(f"auth{origin_ip}",
TokenBucket(replenish_time=30,
max_tokens=3))
if not self.rate_limiter.consume(f"auth{origin_ip}"):
ratelimit_id = f"auth{origin_ip}"
if not self.rate_limiter.get_all_buckets(ratelimit_id):
self.rate_limiter.add_bucket(ratelimit_id,
TokenBucket(replenish_time=60,
max_tokens=self._auth_rpm))
if not self.rate_limiter.consume(ratelimit_id):
bucket = list(self.rate_limiter.get_all_buckets(ratelimit_id).
values())[0]
replenish_time = bucket.last_replenished + bucket.replenish_time
wait_time = round(replenish_time - time())
raise HTTPException(status_code=429,
detail=f"Too many auth requests from: "
f"{origin_ip}. Wait 30 seconds.")
f"{origin_ip}. Wait {wait_time}s.")

if username != "guest":
# TODO: Validate password here
Expand Down

0 comments on commit 63fd997

Please sign in to comment.