From 11a1b13b722af292cbcaf1636458106e9cb530c6 Mon Sep 17 00:00:00 2001
From: Pascal de Vink <pascal@ticketswap.com>
Date: Mon, 24 Apr 2017 14:00:29 +0200
Subject: [PATCH] Use directives to calculate query complexity

---
 src/Validator/Rules/QueryComplexity.php | 31 ++++++++++++++++
 tests/Validator/QueryComplexityTest.php | 48 +++++++++++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/src/Validator/Rules/QueryComplexity.php b/src/Validator/Rules/QueryComplexity.php
index b290da7..83e9cad 100644
--- a/src/Validator/Rules/QueryComplexity.php
+++ b/src/Validator/Rules/QueryComplexity.php
@@ -12,6 +12,7 @@ use GraphQL\Language\AST\NodeKind;
 use GraphQL\Language\AST\OperationDefinitionNode;
 use GraphQL\Language\AST\SelectionSetNode;
 use GraphQL\Language\Visitor;
+use GraphQL\Type\Definition\Directive;
 use GraphQL\Type\Definition\FieldDefinition;
 use GraphQL\Validator\ValidationContext;
 
@@ -138,6 +139,10 @@ class QueryComplexity extends AbstractQuerySecurity
                 $fieldDef = $astFieldInfo[1];
 
                 if ($fieldDef instanceof FieldDefinition) {
+                    if ($this->directiveExcludesField($node)) {
+                        break;
+                    }
+
                     $args = $this->buildFieldArguments($node);
                     //get complexity fn using fieldDef complexity
                     if (method_exists($fieldDef, 'getComplexityFn')) {
@@ -205,6 +210,32 @@ class QueryComplexity extends AbstractQuerySecurity
         return $args;
     }
 
+    private function directiveExcludesField(FieldNode $node) {
+        foreach ($node->directives as $directiveNode) {
+            if ($directiveNode->name->value === 'deprecated') {
+                return false;
+            }
+
+            $variableValues = Values::getVariableValues(
+                $this->context->getSchema(),
+                $this->variableDefs,
+                $this->getRawVariableValues()
+            );
+
+            if ($directiveNode->name->value === 'include') {
+                $directive = Directive::includeDirective();
+                $directiveArgs = Values::getArgumentValues($directive, $directiveNode, $variableValues);
+
+                return !$directiveArgs['if'];
+            } else {
+                $directive = Directive::skipDirective();
+                $directiveArgs = Values::getArgumentValues($directive, $directiveNode, $variableValues);
+
+                return $directiveArgs['if'];
+            }
+        }
+    }
+
     protected function isEnabled()
     {
         return $this->getMaxQueryComplexity() !== static::DISABLED;
diff --git a/tests/Validator/QueryComplexityTest.php b/tests/Validator/QueryComplexityTest.php
index a915f97..0a75117 100644
--- a/tests/Validator/QueryComplexityTest.php
+++ b/tests/Validator/QueryComplexityTest.php
@@ -86,6 +86,54 @@ class QueryComplexityTest extends AbstractQuerySecurityTest
         $this->assertDocumentValidators($query, 3, 4);
     }
 
+    public function testQueryWithEnabledIncludeDirectives()
+    {
+        $query = 'query MyQuery($withDogs: Boolean!) { human { dogs(name: "Root") @include(if:$withDogs) { name } } }';
+
+        $this->getRule()->setRawVariableValues(['withDogs' => true]);
+
+        $this->assertDocumentValidators($query, 3, 4);
+    }
+
+    public function testQueryWithDisabledIncludeDirectives()
+    {
+        $query = 'query MyQuery($withDogs: Boolean!) { human { dogs(name: "Root") @include(if:$withDogs) { name } } }';
+
+        $this->getRule()->setRawVariableValues(['withDogs' => false]);
+
+        $this->assertDocumentValidators($query, 1, 2);
+    }
+
+    public function testQueryWithEnabledSkipDirectives()
+    {
+        $query = 'query MyQuery($withoutDogs: Boolean!) { human { dogs(name: "Root") @skip(if:$withoutDogs) { name } } }';
+
+        $this->getRule()->setRawVariableValues(['withoutDogs' => true]);
+
+        $this->assertDocumentValidators($query, 1, 2);
+    }
+
+    public function testQueryWithDisabledSkipDirectives()
+    {
+        $query = 'query MyQuery($withoutDogs: Boolean!) { human { dogs(name: "Root") @skip(if:$withoutDogs) { name } } }';
+
+        $this->getRule()->setRawVariableValues(['withoutDogs' => false]);
+
+        $this->assertDocumentValidators($query, 3, 4);
+    }
+
+    public function testQueryWithMultipleDirectives()
+    {
+        $query = 'query MyQuery($withDogs: Boolean!, $withoutDogName: Boolean!) { human { dogs(name: "Root") @include(if:$withDogs) { name @skip(if:$withoutDogName) } } }';
+
+        $this->getRule()->setRawVariableValues([
+            'withDogs' => true,
+            'withoutDogName' => true
+        ]);
+
+        $this->assertDocumentValidators($query, 2, 3);
+    }
+
     public function testComplexityIntrospectionQuery()
     {
         $this->assertIntrospectionQuery(181);