diff --git a/apps/api/plane/api/rate_limit.py b/apps/api/plane/api/rate_limit.py index 0d266e98b5..1cc99bfd48 100644 --- a/apps/api/plane/api/rate_limit.py +++ b/apps/api/plane/api/rate_limit.py @@ -3,6 +3,7 @@ import os # Third party imports from rest_framework.throttling import SimpleRateThrottle +from plane.db.models import APIToken class ApiKeyRateThrottle(SimpleRateThrottle): @@ -85,3 +86,27 @@ class ServiceTokenRateThrottle(SimpleRateThrottle): request.META["X-RateLimit-Reset"] = reset_time return allowed + + +class WorkspaceTokenRateThrottle(SimpleRateThrottle): + scope = "workspace_token" + rate = "60/minute" + + def get_cache_key(self, request, view): + api_key = request.headers.get("X-Api-Key") + if not api_key: + return None + + return f"{self.scope}:{api_key}" + + def allow_request(self, request, view): + api_key = request.headers.get("X-Api-Key") + + if api_key: + token = APIToken.objects.filter(token=api_key).only("allowed_rate_limit").first() + if token and token.allowed_rate_limit: + self.rate = token.allowed_rate_limit + # Must re-parse to update num_requests and duration + self.num_requests, self.duration = self.parse_rate(self.rate) + + return super().allow_request(request, view) diff --git a/apps/api/plane/api/views/base.py b/apps/api/plane/api/views/base.py index f17ae2e328..3cb736df61 100644 --- a/apps/api/plane/api/views/base.py +++ b/apps/api/plane/api/views/base.py @@ -13,8 +13,6 @@ from django.utils import timezone from rest_framework import status from rest_framework.permissions import IsAuthenticated from rest_framework.response import Response -from django_filters.rest_framework import DjangoFilterBackend -from rest_framework.filters import SearchFilter from rest_framework.viewsets import ModelViewSet from rest_framework.exceptions import APIException from rest_framework.generics import GenericAPIView @@ -22,7 +20,7 @@ from rest_framework.generics import GenericAPIView # Module imports from plane.db.models.api import APIToken from plane.api.middleware.api_authentication import APIKeyAuthentication -from plane.api.rate_limit import ApiKeyRateThrottle, ServiceTokenRateThrottle +from plane.api.rate_limit import ApiKeyRateThrottle, ServiceTokenRateThrottle, WorkspaceTokenRateThrottle from plane.utils.exception_logger import log_exception from plane.utils.paginator import BasePaginator from plane.utils.core.mixins import ReadReplicaControlMixin @@ -62,12 +60,20 @@ class BaseAPIView(TimezoneMixin, GenericAPIView, ReadReplicaControlMixin, BasePa api_key = self.request.headers.get("X-Api-Key") if api_key: - service_token = APIToken.objects.filter(token=api_key, is_service=True).first() + api_token = APIToken.objects.filter(token=api_key) + + service_token = api_token.filter(is_service=True).first() + + workspace_token = api_token.filter(workspace_id__isnull=False).first() if service_token: throttle_classes.append(ServiceTokenRateThrottle()) return throttle_classes + if workspace_token: + throttle_classes.append(WorkspaceTokenRateThrottle()) + return throttle_classes + throttle_classes.append(ApiKeyRateThrottle()) return throttle_classes