feat: implement plan-based API rate limits

backend/api/auth.py

@@ -19,7 +19,7 @@
 
 class ServiceAccountUser:
     """Mock ServiceAccount user"""
-    
+
     def __init__(self, service_account):
         self.userId = service_account.id
         self.id = service_account.id
@@ -63,14 +63,20 @@ def authenticate(self, request):
         if secret_id:
             found = False
             try:
-                secret = Secret.objects.get(id=secret_id)
+                # Pre-fetch environment, app, and organisation
+                secret = Secret.objects.select_related(
+                    "environment__app__organisation"
+                ).get(id=secret_id)
                 env = secret.environment
                 found = True
             except Secret.DoesNotExist:
                 pass
             if not found:
                 try:
-                    dyn_secret = DynamicSecret.objects.get(id=secret_id)
+                    # Pre-fetch environment, app, and organisation
+                    dyn_secret = DynamicSecret.objects.select_related(
+                        "environment__app__organisation"
+                    ).get(id=secret_id)
                     env = dyn_secret.environment
                     found = True
                 except DynamicSecret.DoesNotExist:
@@ -84,7 +90,10 @@ def authenticate(self, request):
             # Try resolving env from header
             if env_id:
                 try:
-                    env = Environment.objects.get(id=env_id)
+                    # Pre-fetch app and organisation
+                    env = Environment.objects.select_related("app__organisation").get(
+                        id=env_id
+                    )
                 except Environment.DoesNotExist:
                     raise exceptions.AuthenticationFailed("Environment not found")
 
@@ -99,7 +108,10 @@ def authenticate(self, request):
                         )
                     if not env_name:
                         raise exceptions.AuthenticationFailed("Missing env parameter")
-                    env = Environment.objects.get(app_id=app_id, name__iexact=env_name)
+                    # Pre-fetch app and organisation
+                    env = Environment.objects.select_related("app__organisation").get(
+                        app_id=app_id, name__iexact=env_name
+                    )
                 except Environment.DoesNotExist:
                     # Check if the app exists to give a more specific error
                     App = apps.get_model("api", "App")
@@ -134,14 +146,14 @@ def authenticate(self, request):
 
             try:
                 service_token = get_service_token(auth_token)
-                service_account = get_service_account_from_token(auth_token)          
-                
+                service_account = get_service_account_from_token(auth_token)
+
                 creator = getattr(service_token, "created_by", None)
                 if creator:
                     user = creator.user
                 else:
                     user = ServiceAccountUser(service_account)
-                    
+
                 auth["service_account"] = service_account
                 auth["service_account_token"] = service_token
 

backend/api/throttling.py

@@ -0,0 +1,56 @@
+from rest_framework.throttling import SimpleRateThrottle
+from django.conf import settings
+
+CLOUD_HOSTED = settings.APP_HOST == "cloud"
+
+
+class PlanBasedRateThrottle(SimpleRateThrottle):
+    """
+    Limits the rate of API calls based on the Organisation's plan.
+    Uses the pre-fetched organisation data from request.auth to avoid DB lookups.
+    """
+
+    scope = "plan_based"
+
+    def get_cache_key(self, request, view):
+        # Identify the user or service account
+        ident = self.get_ident(request)
+
+        if request.user.is_authenticated and request.auth:
+            if request.auth.get("org_member"):
+                ident = f"user_{request.auth['org_member'].id}"
+            elif request.auth.get("service_account"):
+                ident = f"sa_{request.auth['service_account'].id}"
+        else:
+            ident = f"anon_{ident}"
+
+        return self.cache_format % {"scope": self.scope, "ident": ident}
+
+    def allow_request(self, request, view):
+        """
+        Override allow_request to dynamically set the rate based on the request user's plan.
+        """
+        # Default fallback (reads from REST_FRAMEWORK['DEFAULT_THROTTLE_RATES']['plan_based'])
+        new_rate = self.get_rate()
+
+        if request.user.is_authenticated and request.auth:
+            env = request.auth.get("environment")
+            if env:
+                try:
+                    plan = env.app.organisation.plan
+                    new_rate = self.get_rate_for_plan(plan)
+                except AttributeError:
+                    pass
+
+        # Update the throttle configuration for this specific request
+        self.rate = new_rate
+        self.num_requests, self.duration = self.parse_rate(self.rate)
+
+        return super().allow_request(request, view)
+
+    @staticmethod
+    def get_rate_for_plan(plan):
+        # If self-hosted return the default rate limit. If not set, this will disable throttling
+        if not CLOUD_HOSTED:
+            return settings.PLAN_RATE_LIMITS["DEFAULT"]
+        return settings.PLAN_RATE_LIMITS.get(plan, settings.PLAN_RATE_LIMITS["DEFAULT"])

backend/api/views/identities/aws/iam.py

@@ -7,13 +7,14 @@
 from defusedxml.ElementTree import parse
 from django.http import JsonResponse
 from django.utils import timezone
-from rest_framework.decorators import api_view, permission_classes
+from rest_framework.decorators import api_view, permission_classes, throttle_classes
 from rest_framework.permissions import AllowAny
 
 from api.utils.identity.common import (
     resolve_service_account,
     mint_service_account_token,
 )
+from api.throttling import PlanBasedRateThrottle
 
 
 def get_normalized_host(uri):
@@ -27,6 +28,7 @@ def get_normalized_host(uri):
 
 @api_view(["POST"])
 @permission_classes([AllowAny])
+@throttle_classes([PlanBasedRateThrottle])
 def aws_iam_auth(request):
     """Accepts SigV4-signed STS GetCallerIdentity request and issues a ServiceAccount token if trusted."""
     try:

backend/api/views/secrets.py

@@ -32,6 +32,7 @@
 import json
 from api.content_negotiation import CamelCaseContentNegotiation
 from api.utils.access.middleware import IsIPAllowed
+from api.throttling import PlanBasedRateThrottle
 from ee.integrations.secrets.dynamic.exceptions import (
     DynamicSecretError,
     PlanRestrictionError,
@@ -59,6 +60,7 @@
 class E2EESecretsView(APIView):
     authentication_classes = [PhaseTokenAuthentication]
     permission_classes = [IsAuthenticated, IsIPAllowed]
+    throttle_classes = [PlanBasedRateThrottle]
     content_negotiation_class = CamelCaseContentNegotiation
 
     def initial(self, request, *args, **kwargs):
@@ -488,6 +490,7 @@ def delete(self, request, *args, **kwargs):
 class PublicSecretsView(APIView):
     authentication_classes = [PhaseTokenAuthentication]
     permission_classes = [IsAuthenticated, IsIPAllowed]
+    throttle_classes = [PlanBasedRateThrottle]
     renderer_classes = [
         CamelCaseJSONRenderer,
     ]

backend/backend/settings.py

@@ -241,13 +241,25 @@ def get_version():
     "USER_DETAILS_SERIALIZER": "api.serializers.CustomUserSerializer"
 }
 
+# Define rate limits per plan (keys match Organisation model constants: FR, PR, EN)
+PLAN_RATE_LIMITS = {
+    "FR": os.getenv("RATE_LIMIT_FREE", "120/min"),
+    "PR": os.getenv("RATE_LIMIT_PRO", "240/min"),
+    "EN": os.getenv("RATE_LIMIT_ENTERPRISE", "1000/min"),
+    "DEFAULT": os.getenv("RATE_LIMIT_DEFAULT"),
+}
+
 REST_FRAMEWORK = {
     "DEFAULT_PERMISSION_CLASSES": [
         "rest_framework.permissions.IsAuthenticated",
     ],
     "DEFAULT_AUTHENTICATION_CLASSES": (
         "rest_framework.authentication.SessionAuthentication",
     ),
+    "DEFAULT_THROTTLE_CLASSES": [],
+    "DEFAULT_THROTTLE_RATES": {
+        "plan_based": PLAN_RATE_LIMITS["DEFAULT"],
+    },
     "EXCEPTION_HANDLER": "backend.exceptions.custom_exception_handler",
     "DEFAULT_RENDERER_CLASSES": [
         "rest_framework.renderers.JSONRenderer",
@@ -294,6 +306,44 @@ def get_version():
     },
 }
 
+REDIS_HOST = os.getenv("REDIS_HOST")
+REDIS_PORT = os.getenv("REDIS_PORT")
+REDIS_PASSWORD = get_secret("REDIS_PASSWORD")
+REDIS_SSL = os.getenv("REDIS_SSL", "False").lower() == "true"
+REDIS_PROTOCOL = "rediss" if REDIS_SSL else "redis"
+REDIS_AUTH = f":{REDIS_PASSWORD}@" if REDIS_PASSWORD else ""
+
+CACHES = {
+    "default": {
+        "BACKEND": "django.core.cache.backends.redis.RedisCache",
+        "LOCATION": f"{REDIS_PROTOCOL}://{REDIS_AUTH}{REDIS_HOST}:{REDIS_PORT}/1",
+        "OPTIONS": (
+            {
+                "ssl_cert_reqs": None,
+            }
+            if REDIS_SSL
+            else {}
+        ),
+    }
+}
+
+RQ_QUEUES = {
+    "default": {
+        "HOST": REDIS_HOST,
+        "PORT": REDIS_PORT,
+        "PASSWORD": REDIS_PASSWORD,
+        "SSL": REDIS_SSL,
+        "DB": 0,
+    },
+    "scheduled-jobs": {
+        "HOST": REDIS_HOST,
+        "PORT": REDIS_PORT,
+        "PASSWORD": REDIS_PASSWORD,
+        "SSL": REDIS_SSL,
+        "DB": 0,
+    },
+}
+
 DYNAMODB = {
     "TABLE": os.getenv("DYNAMODB_LOGS_TABLE"),
     "INDEX": os.getenv("DYNAMODB_LOGS_TIMESTAMP_INDEX"),
@@ -358,22 +408,6 @@ def get_version():
 except:
     APP_HOST = "self"
 
-RQ_QUEUES = {
-    "default": {
-        "HOST": os.getenv("REDIS_HOST"),
-        "PORT": os.getenv("REDIS_PORT"),
-        "PASSWORD": get_secret("REDIS_PASSWORD"),
-        "SSL": os.getenv("REDIS_SSL", None),
-        "DB": 0,
-    },
-    "scheduled-jobs": {
-        "HOST": os.getenv("REDIS_HOST"),
-        "PORT": os.getenv("REDIS_PORT"),
-        "PASSWORD": get_secret("REDIS_PASSWORD"),
-        "SSL": os.getenv("REDIS_SSL", None),
-        "DB": 0,
-    },
-}
 
 PHASE_LICENSE = check_license(get_secret("PHASE_LICENSE_OFFLINE"))
 

backend/conftest.py

@@ -0,0 +1,23 @@
+import os
+import django
+
+# Set environment variables required for settings.py to import successfully
+os.environ.setdefault("ALLOWED_HOSTS", "localhost")
+os.environ.setdefault("ALLOWED_ORIGINS", "http://localhost")
+
+# Set dummy Redis values so settings.py generates a valid URL (e.g. redis://localhost:6379/1)
+os.environ.setdefault("REDIS_HOST", "localhost")
+os.environ.setdefault("REDIS_PORT", "6379")
+
+# Set dummy database config
+os.environ.setdefault("DATABASE_HOST", "localhost")
+os.environ.setdefault("DATABASE_PORT", "5432")
+os.environ.setdefault("DATABASE_NAME", "dummy_db")
+os.environ.setdefault("DATABASE_USER", "dummy_user")
+os.environ.setdefault("DATABASE_PASSWORD", "dummy_password")
+
+os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings")
+
+
+def pytest_configure():
+    django.setup()

backend/dev-requirements.txt

@@ -1,4 +1,5 @@
 pytest==8.3.4
+pytest-django==4.11.1
 pytest-cov==7.0.0
 Faker==37.4.0
 colorama==0.4.6

backend/ee/integrations/secrets/dynamic/rest/views.py

@@ -18,6 +18,7 @@
 )
 
 from api.utils.access.middleware import IsIPAllowed
+from api.throttling import PlanBasedRateThrottle
 from ee.integrations.secrets.dynamic.aws.utils import (
     revoke_aws_dynamic_secret_lease,
 )
@@ -50,6 +51,7 @@
 class DynamicSecretsView(APIView):
     authentication_classes = [PhaseTokenAuthentication]
     permission_classes = [IsAuthenticated, IsIPAllowed]
+    throttle_classes = [PlanBasedRateThrottle]
     renderer_classes = [
         CamelCaseJSONRenderer,
     ]
@@ -200,6 +202,7 @@ def get(self, request, *args, **kwargs):
 class DynamicSecretLeaseView(APIView):
     authentication_classes = [PhaseTokenAuthentication]
     permission_classes = [IsAuthenticated, IsIPAllowed]
+    throttle_classes = [PlanBasedRateThrottle]
     renderer_classes = [CamelCaseJSONRenderer]
 
     def _get_account_and_org(self, request):

backend/pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+python_files = tests.py test_*.py *_tests.py