From 63fd9974ac20a11ddec6533856da030104d66df6 Mon Sep 17 00:00:00 2001 From: Daniel McKnight <34697904+NeonDaniel@users.noreply.github.com> Date: Thu, 25 Jan 2024 16:35:54 -0800 Subject: [PATCH] Configurable Authorization Request Limits (#9) * 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 --- README.md | 1 + neon_hana/auth/client_manager.py | 18 ++++++++++++------ 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 38001d6..02e3cd1 100644 --- a/README.md +++ b/README.md @@ -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" diff --git a/neon_hana/auth/client_manager.py b/neon_hana/auth/client_manager.py index 1005a11..253b825 100644 --- a/neon_hana/auth/client_manager.py +++ b/neon_hana/auth/client_manager.py @@ -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" @@ -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