Configurable Authorization Request Limits

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"

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