feat: add misc graphql validation rules (#686)

rohan-chaturvedi · nimish-ks · web-flow · commit fe603d2f810f · 2025-11-13T13:13:56.000+05:30
* feat: implement custom validation rules for duplicate fields and alias usage for graphene Signed-off-by: rohan <rohan.chaturvedi@protonmail.com> * chore: misc cleanup and documentation, reduce limits Signed-off-by: rohan <rohan.chaturvedi@protonmail.com> * feat: add access control checks for environment in bulk secret mutations Signed-off-by: rohan <rohan.chaturvedi@protonmail.com> * fix: add validation for role changes, invite roles Signed-off-by: rohan <rohan.chaturvedi@protonmail.com> * fix: client IP discovery utils (#687) * refactor: use standardized get_client_ip util Signed-off-by: rohan <rohan.chaturvedi@protonmail.com> * refactor: simplify IP retrieval in IsIPAllowed permission class by using get_client_ip utility * refactor: streamline client IP retrieval in IPWhitelistMiddleware by utilizing get_client_ip utility * refactor: enhance get_client_ip utility to validate IP addresses and improve retrieval logic * debug: add print statement to log raw IP address in get_client_ip utility * chore: add a blank line at the end of ip.py for consistency * refactor: update client IP retrieval logic in get_client_ip utility to prioritize HTTP_X_REAL_IP * refactor: remove debug print statement from get_client_ip utility * fix: remove unnecessary fallbacks Signed-off-by: rohan <rohan.chaturvedi@protonmail.com> --------- Signed-off-by: rohan <rohan.chaturvedi@protonmail.com> Co-authored-by: rohan <rohan.chaturvedi@protonmail.com> * fix: add env access check for bulk secret deletes Signed-off-by: rohan <rohan.chaturvedi@protonmail.com> * fix: validate service account roles during create and update Signed-off-by: rohan <rohan.chaturvedi@protonmail.com> * fix: add safe fallback to raw ip in case header is missing Signed-off-by: rohan <rohan.chaturvedi@protonmail.com> --------- Signed-off-by: rohan <rohan.chaturvedi@protonmail.com> Co-authored-by: Nimish <85357445+nimish-ks@users.noreply.github.com>
diff --git a/backend/api/emails.py b/backend/api/emails.py
@@ -4,11 +4,13 @@
 from datetime import datetime
 import os
 import logging
-from api.utils.rest import encode_string_to_base64, get_client_ip
+from api.utils.rest import encode_string_to_base64
 from api.models import OrganisationMember
 from django.utils import timezone
 from smtplib import SMTPException
 
+from api.utils.access.ip import get_client_ip
+
 logger = logging.getLogger(__name__)
 
 
diff --git a/backend/api/utils/access/ip.py b/backend/api/utils/access/ip.py
@@ -0,0 +1,21 @@
+from ipaddress import ip_address
+
+
+def get_client_ip(request):
+    """
+    Get the client IP address as a single string.
+
+    Args:
+        request: Django request object
+
+    Returns:
+        str | None: The client IP address (IPv4 or IPv6)
+    """
+    raw_ip = (request.META.get("HTTP_X_REAL_IP") or "").strip()
+    if not raw_ip:
+        return None
+    try:
+        ip_address(raw_ip)
+    except ValueError:
+        return None
+    return raw_ip
diff --git a/backend/api/utils/access/middleware.py b/backend/api/utils/access/middleware.py
@@ -1,6 +1,7 @@
 # permissions.py
 
 from api.models import NetworkAccessPolicy, Organisation
+from api.utils.access.ip import get_client_ip
 from rest_framework.permissions import BasePermission
 from itertools import chain
 
@@ -15,10 +16,7 @@ class IsIPAllowed(BasePermission):
     )
 
     def get_client_ip(self, request):
-        x_forwarded_for = request.META.get("HTTP_X_FORWARDED_FOR")
-        if x_forwarded_for:
-            return x_forwarded_for.split(",")[0].strip()
-        return request.META.get("REMOTE_ADDR")
+        return get_client_ip(request)
 
     def has_permission(self, request, view):
         ip = self.get_client_ip(request)
diff --git a/backend/api/utils/rest.py b/backend/api/utils/rest.py
@@ -1,6 +1,7 @@
 from api.models import EnvironmentToken, ServiceAccountToken, ServiceToken, UserToken
 from django.utils import timezone
 import base64
+from api.utils.access.ip import get_client_ip
 
 # Map HTTP methods to permission actions
 METHOD_TO_ACTION = {
@@ -11,15 +12,6 @@
 }
 
 
-def get_client_ip(request):
-    x_forwarded_for = request.META.get("HTTP_X_FORWARDED_FOR")
-    if x_forwarded_for:
-        ip = x_forwarded_for.split(",")[0]
-    else:
-        ip = request.META.get("REMOTE_ADDR")
-    return ip
-
-
 def get_resolver_request_meta(request):
     user_agent = request.META.get("HTTP_USER_AGENT", "Unknown")
     ip_address = get_client_ip(request)
diff --git a/backend/api/views/graphql.py b/backend/api/views/graphql.py
@@ -1,7 +1,15 @@
 from graphene_django.views import GraphQLView
 from django.contrib.auth.mixins import LoginRequiredMixin
+from graphql import specified_rules
+from backend.graphene.validation import DuplicateFieldLimitRule, AliasUsageLimitRule
+
+
+CUSTOM_RULES = tuple(specified_rules) + (
+    DuplicateFieldLimitRule,
+    AliasUsageLimitRule,
+)
 
 
 class PrivateGraphQLView(LoginRequiredMixin, GraphQLView):
     raise_exception = True
-    pass
+    validation_rules = CUSTOM_RULES
diff --git a/backend/api/views/kms.py b/backend/api/views/kms.py
@@ -1,11 +1,9 @@
 from datetime import datetime
-
+from api.utils.access.ip import get_client_ip
 from rest_framework.decorators import api_view, permission_classes
 from rest_framework.permissions import AllowAny
 from django.http import JsonResponse, HttpResponse
-from api.utils.rest import (
-    get_client_ip,
-)
+
 from logs.models import KMSDBLog
 from api.models import (
     App,
diff --git a/backend/backend/graphene/middleware.py b/backend/backend/graphene/middleware.py
@@ -3,6 +3,7 @@
 from api.models import NetworkAccessPolicy, Organisation, OrganisationMember
 
 from itertools import chain
+from api.utils.access.ip import get_client_ip
 
 
 class IPRestrictedError(GraphQLError):
@@ -51,7 +52,7 @@ def resolve(self, next, root, info: GraphQLResolveInfo, **kwargs):
             except OrganisationMember.DoesNotExist:
                 raise GraphQLError("You are not a member of this organisation")
 
-            ip = self.get_client_ip(request)
+            ip = get_client_ip(request)
 
             account_policies = org_member.network_policies.all()
             global_policies = (
@@ -70,7 +71,4 @@ def resolve(self, next, root, info: GraphQLResolveInfo, **kwargs):
             raise IPRestrictedError(org_member.organisation.name)
 
     def get_client_ip(self, request):
-        x_forwarded_for = request.META.get("HTTP_X_FORWARDED_FOR")
-        if x_forwarded_for:
-            return x_forwarded_for.split(",")[0].strip()
-        return request.META.get("REMOTE_ADDR")
+        return get_client_ip(request)
diff --git a/backend/backend/graphene/mutations/environment.py b/backend/backend/graphene/mutations/environment.py
@@ -784,6 +784,9 @@ def mutate(cls, root, info, secrets_data):
                     "You don't have permission to create secrets in this organisation"
                 )
 
+            if not user_can_access_environment(info.context.user.userId, env.id):
+                raise GraphQLError("You don't have access to this environment")
+
             tags = SecretTag.objects.filter(id__in=secret_data.tags)
 
             path = (
@@ -846,6 +849,9 @@ def mutate(cls, root, info, id, secret_data):
                 "You don't have permission to update secrets in this organisation"
             )
 
+        if not user_can_access_environment(info.context.user.userId, env.id):
+            raise GraphQLError("You don't have access to this environment")
+
         tags = SecretTag.objects.filter(id__in=secret_data.tags)
 
         path = (
@@ -905,6 +911,9 @@ def mutate(cls, root, info, secrets_data):
                     "You don't have permission to update secrets in this organisation"
                 )
 
+            if not user_can_access_environment(info.context.user.userId, env.id):
+                raise GraphQLError("You don't have access to this environment")
+
             tags = SecretTag.objects.filter(id__in=secret_data.tags)
 
             path = (
@@ -1003,6 +1012,9 @@ def mutate(cls, root, info, ids):
                     "You don't have permission to delete secrets in this organisation"
                 )
 
+            if not user_can_access_environment(info.context.user.userId, env.id):
+                raise GraphQLError("You don't have access to this environment")
+
             secret.updated_at = timezone.now()
             secret.deleted_at = timezone.now()
             secret.save()
diff --git a/backend/backend/graphene/mutations/organisation.py b/backend/backend/graphene/mutations/organisation.py
@@ -161,6 +161,14 @@ def mutate(cls, root, info, org_id, invites):
 
             app_scope = App.objects.filter(id__in=apps)
 
+            # Restrict roles that can be assigned via invites
+            allowed_invite_roles = ["developer", "service"]
+            role = Role.objects.get(organisation=org, id=role_id)
+            if role.name.lower() not in allowed_invite_roles:
+                raise GraphQLError(
+                    f"You can only invite members with the following roles: {', '.join(allowed_invite_roles)}"
+                )
+
             new_invite = OrganisationMemberInvite.objects.create(
                 organisation=org,
                 role_id=role_id,
@@ -317,6 +325,9 @@ def mutate(cls, root, info, member_id, role_id):
         ):
             raise GraphQLError("You dont have permission to change member roles")
 
+        if org_member.user == info.context.user:
+            raise GraphQLError("You can't change your own role in an organisation")
+
         active_user_role = OrganisationMember.objects.get(
             user=info.context.user,
             organisation=org_member.organisation,
diff --git a/backend/backend/graphene/mutations/service_accounts.py b/backend/backend/graphene/mutations/service_accounts.py
@@ -9,7 +9,11 @@
     ServiceAccountToken,
     Identity,
 )
-from api.utils.access.permissions import user_has_permission, user_is_org_member
+from api.utils.access.permissions import (
+    role_has_global_access,
+    user_has_permission,
+    user_is_org_member,
+)
 from backend.graphene.types import ServiceAccountTokenType, ServiceAccountType
 from datetime import datetime
 from django.conf import settings
@@ -58,10 +62,17 @@ def mutate(
         if handlers is None or len(handlers) == 0:
             raise GraphQLError("At least one service account handler must be provided")
 
+        role = Role.objects.get(id=role_id, organisation=org)
+
+        if role_has_global_access(role):
+            raise GraphQLError(
+                f"Service Accounts cannot be assigned the '{role.name}' role."
+            )
+
         service_account = ServiceAccount.objects.create(
             name=name,
             organisation=org,
-            role=Role.objects.get(id=role_id),
+            role=role,
             identity_key=identity_key,
             server_wrapped_keyring=server_wrapped_keyring,
             server_wrapped_recovery=server_wrapped_recovery,
@@ -168,7 +179,12 @@ def mutate(cls, root, info, service_account_id, name, role_id, identity_ids=None
                 "You don't have the permissions required to update Service Accounts in this organisation"
             )
 
-        role = Role.objects.get(id=role_id)
+        role = Role.objects.get(id=role_id, organisation=service_account.organisation)
+
+        if role_has_global_access(role):
+            raise GraphQLError(
+                f"Service Accounts cannot be assigned the '{role.name}' role."
+            )
         service_account.name = name
         service_account.role = role
         if identity_ids is not None:
diff --git a/backend/backend/graphene/queries/access.py b/backend/backend/graphene/queries/access.py
@@ -6,6 +6,7 @@
     Role,
     Identity,
 )
+from api.utils.access.ip import get_client_ip
 from graphql import GraphQLError
 from django.db import transaction
 from api.utils.access.roles import default_roles
@@ -96,12 +97,7 @@ def resolve_network_access_policies(root, info, organisation_id):
 
 def resolve_client_ip(root, info):
     request = info.context
-    # Use common headers to support reverse proxies
-    x_forwarded_for = request.META.get("HTTP_X_FORWARDED_FOR")
-    if x_forwarded_for:
-        ip = x_forwarded_for.split(",")[0].strip()
-    else:
-        ip = request.META.get("REMOTE_ADDR")
+    ip = get_client_ip(request)
     return ip
 
 
diff --git a/backend/backend/graphene/validation.py b/backend/backend/graphene/validation.py
@@ -0,0 +1,89 @@
+from collections import Counter
+from graphql import GraphQLError
+from graphql.language.ast import FieldNode
+from graphql.validation import ValidationRule
+
+
+class DuplicateFieldLimitRule(ValidationRule):
+    """Limits how many times the same response field (name or alias) can appear
+    in a single selection set. Uses a Counter stack to track counts per nested
+    selection set and reports an error if any exceeds MAX_DUPLICATE_FIELDS."""
+
+    MAX_DUPLICATE_FIELDS = 20
+
+    def __init__(self, context):
+        super().__init__(context)
+        self.max_duplicates = self.MAX_DUPLICATE_FIELDS
+        self._stack = []  # Stack of Counters for nested selection sets
+
+    def enter_selection_set(self, *_):
+        """Push a new Counter for a nested selection set."""
+        self._stack.append(Counter())
+
+    def leave_selection_set(self, node, *_):
+        """Pop Counter, emit an error for any response name exceeding limit."""
+        counts = self._stack.pop()
+        for response_name, hits in counts.items():
+            if hits > self.max_duplicates:
+                offending = [
+                    selection
+                    for selection in node.selections
+                    if isinstance(selection, FieldNode)
+                    and (
+                        selection.alias.value
+                        if selection.alias
+                        else selection.name.value
+                    )
+                    == response_name
+                ]
+                self.context.report_error(
+                    GraphQLError(
+                        f"Field '{response_name}' requested {hits} times; limit is {self.max_duplicates}.",
+                        nodes=offending or [node],
+                    )
+                )
+
+    def enter_field(self, node, *_):
+        """Increment count for this field’s response name (alias or original)."""
+        response_name = node.alias.value if node.alias else node.name.value
+        if not self._stack:
+            self._stack.append(Counter())
+        self._stack[-1][response_name] += 1
+
+
+class AliasUsageLimitRule(ValidationRule):
+    """Caps the number of aliases used within a single operation definition.
+    Tracks alias count per operation; reports an error when exceeding
+    MAX_ALIAS_FIELDS."""
+
+    MAX_ALIAS_FIELDS = 20
+
+    def __init__(self, context):
+        super().__init__(context)
+        self.max_aliases = self.MAX_ALIAS_FIELDS
+        self._operation_alias_counts = (
+            []
+        )  # Stack for nested operations (fragments not counted)
+
+    def enter_operation_definition(self, *_):
+        """Start alias count for a new operation."""
+        self._operation_alias_counts.append(0)
+
+    def leave_operation_definition(self, *_):
+        """End alias count scope for the operation."""
+        self._operation_alias_counts.pop()
+
+    def enter_field(self, node, *_):
+        """Increment alias counter when a field has an alias; error if limit exceeded."""
+        if node.alias:
+            if not self._operation_alias_counts:
+                self._operation_alias_counts.append(0)
+            alias_count = self._operation_alias_counts[-1] + 1
+            self._operation_alias_counts[-1] = alias_count
+            if alias_count > self.max_aliases:
+                self.context.report_error(
+                    GraphQLError(
+                        f"Alias limit of {self.max_aliases} exceeded in a single operation.",
+                        nodes=[node],
+                    )
+                )
diff --git a/backend/requirements.txt b/backend/requirements.txt
@@ -25,7 +25,7 @@ djangorestframework==3.15.2
 djangorestframework-camel-case==1.4.2
 freezegun==1.5.1
 graphene==3.2.1
-graphene-django==3.0.0
+graphene-django==3.2.0
 graphql-core==3.2.3
 graphql-relay==3.2.0
 gunicorn==23.0.0
