Add DepthAnalyzer: disallow graphql queries over depth 7

graphy/utils/graphql.py

@@ -1,5 +1,5 @@
 from graphene_django.views import GraphQLView
-from graphql import GraphQLError
+from graphql import GraphQLError, GraphQLCoreBackend
 
 from rest_framework import (
     permissions,
@@ -16,6 +16,42 @@ class GraphqlAuthenticationError(GraphQLError):
     pass
 
 
+def measure_depth(definition, level=1):
+    depth = level
+
+    selection_set = getattr(definition, 'selection_set', None)
+    if selection_set:
+        selections = getattr(definition.selection_set, 'selections', [])
+    else:
+        selections = getattr(definition, 'selections', [])
+
+    for field in selections:
+        if hasattr(field, 'selection_set'):
+            new_depth = measure_depth(field.selection_set, level=level + 1)
+            if new_depth > depth:
+                depth = new_depth
+    return depth
+
+
+class DepthAnalysisBackend(GraphQLCoreBackend):
+    MAX_DEPTH = 7
+
+    def document_from_string(self, schema, document_string):
+        document = super().document_from_string(schema, document_string)
+        ast_definitions = document.document_ast.definitions
+
+        for definition in ast_definitions:
+            # We are only interested in queries
+            if getattr(definition, 'operation', None) != 'query':
+                continue
+
+            depth = measure_depth(definition)
+            if depth > self.MAX_DEPTH:
+                raise Exception('Query is too complex')
+
+        return document
+
+
 def auth_required(fn):
     def wrapper(*args, **kwargs):
         *_, info = args
@@ -45,7 +81,8 @@ def format_error(error):
 
     @classmethod
     def as_view(cls, *args, **kwargs):
-        view = super(GraphQLView, cls).as_view(*args, **kwargs)
+        backend = DepthAnalysisBackend()
+        view = super(GraphQLView, cls).as_view(backend=backend, *args, **kwargs)
         view = permission_classes((permissions.AllowAny,))(view)
         view = api_view(['GET', 'POST'])(view)
         return view
@@ -62,7 +99,8 @@ def parse_body(self, request):
 
     @classmethod
     def as_view(cls, *args, **kwargs):
-        view = super(GraphQLView, cls).as_view(*args, **kwargs)
+        backend = DepthAnalysisBackend()
+        view = super(GraphQLView, cls).as_view(backend=backend, *args, **kwargs)
         view = permission_classes((permissions.AllowAny,))(view)
         view = api_view(['GET', 'POST'])(view)
         return view

graphy/utils/tests/test_depth_analyzer.py

@@ -0,0 +1,85 @@
+import json
+
+from graphy.customers.tests.conftest import *  # noqa
+from graphy.leads.tests.conftest import *  # noqa
+
+
+def test_can_query_with_depth_7(drf_client_staff, address, customer, lead):
+    """
+    This test looks a bit funky. Our gql_client fixture does not include the
+    depth analyzer. The endpoint on our server does however
+    """
+    lead.address = address
+    lead.save()
+    customer.leads.add(lead)
+    expected_county_id = customer.home_address.zip_code.municipality.county_id
+    query = """
+        query {
+          customers {
+            homeAddress {
+              zipCode {
+                municipality {
+                  county {
+                    id
+                  }
+                }
+              }
+            }
+          }
+        }
+    """
+
+    result = drf_client_staff.post(
+        '/graphql/?',
+        data={'operationName': 'null', 'query': query, 'variables': 'null'},
+    )
+    assert json.loads(result._container[0])['data']['customers'] == [
+        {
+            'homeAddress': {
+                'zipCode': {
+                    'municipality': {
+                        'county': {
+                            'id': str(expected_county_id)
+                        }
+                    }
+                }
+            }
+        }
+    ]
+
+
+def test_can_not_query_with_depth_8(drf_client_staff, address, customer, lead):
+    """
+    This test looks a bit funky. Our gql_client fixture does not include the
+    depth analyzer. The endpoint on our server does however
+    """
+    lead.address = address
+    lead.save()
+    customer.leads.add(lead)
+
+    query = """
+        query {
+          customers {
+            leads {
+              address {
+                zipCode {
+                  municipality {
+                    county {
+                      id
+                    }
+                  }
+                }
+              }
+            }
+          }
+        }
+    """
+
+    result = drf_client_staff.post(
+        '/graphql/?',
+        data={'operationName': 'null', 'query': query, 'variables': 'null'},
+    )
+
+    assert json.loads(result._container[0]) == {
+        'errors': [{'message': 'Query is too complex'}]
+    }