diff --git a/apiserver/plane/api/rate_limit.py b/apiserver/plane/api/rate_limit.py index b62936d8e..65cfd6b0c 100644 --- a/apiserver/plane/api/rate_limit.py +++ b/apiserver/plane/api/rate_limit.py @@ -40,3 +40,44 @@ class ApiKeyRateThrottle(SimpleRateThrottle): request.META["X-RateLimit-Reset"] = reset_time return allowed + + +class ServiceTokenRateThrottle(SimpleRateThrottle): + scope = "service_token" + rate = "300/minute" + + def get_cache_key(self, request, view): + # Retrieve the API key from the request header + api_key = request.headers.get("X-Api-Key") + if not api_key: + return None # Allow the request if there's no API key + + # Use the API key as part of the cache key + return f"{self.scope}:{api_key}" + + def allow_request(self, request, view): + allowed = super().allow_request(request, view) + + if allowed: + now = self.timer() + # Calculate the remaining limit and reset time + history = self.cache.get(self.key, []) + + # Remove old histories + while history and history[-1] <= now - self.duration: + history.pop() + + # Calculate the requests + num_requests = len(history) + + # Check available requests + available = self.num_requests - num_requests + + # Unix timestamp for when the rate limit will reset + reset_time = int(now + self.duration) + + # Add headers + request.META["X-RateLimit-Remaining"] = max(0, available) + request.META["X-RateLimit-Reset"] = reset_time + + return allowed \ No newline at end of file