fix: inline fragments visitor (#375)

* fix: inline fragments
Escape-Technologies · Mar 28, 2023 · 8d0ab85 · 8d0ab85
1 parent c321159
commit 8d0ab85
Show file tree

Hide file tree

Showing 9 changed files with 54 additions and 29 deletions.
diff --git a/.changeset/six-readers-juggle.md b/.changeset/six-readers-juggle.md
@@ -0,0 +1,9 @@
+---
+'@escape.tech/graphql-armor-max-directives': patch
+'@escape.tech/graphql-armor-max-aliases': patch
+'@escape.tech/graphql-armor-cost-limit': patch
+'@escape.tech/graphql-armor-max-depth': patch
+'@escape.tech/graphql-armor': patch
+---
+
+Fix: Inline fragment visitor (Thanks @simoncrypta @dthyresson)
diff --git a/examples/apollo-v3/test/index.spec.ts b/examples/apollo-v3/test/index.spec.ts
@@ -55,9 +55,7 @@ describe('startup', () => {
       }`,
     });
     expect(query.errors).toBeDefined();
-    expect(query.errors?.map((e) => e.message)).toContain(
-      'Syntax Error: Query Cost limit of 100 exceeded, found 5023.',
-    );
+    expect(query.errors?.map((e) => e.message)).toContain('Syntax Error: Query Cost limit of 100 exceeded, found 138.');
   });
 
   it('should block field suggestion', async () => {

diff --git a/examples/apollo/test/index.spec.ts b/examples/apollo/test/index.spec.ts
@@ -62,9 +62,7 @@ describe('startup', () => {
 
     const query = body.singleResult;
     expect(query.errors).toBeDefined();
-    expect(query.errors?.map((e) => e.message)).toContain(
-      'Syntax Error: Query Cost limit of 100 exceeded, found 5023.',
-    );
+    expect(query.errors?.map((e) => e.message)).toContain('Syntax Error: Query Cost limit of 100 exceeded, found 138.');
   });
 
   it('should block field suggestion', async () => {

diff --git a/examples/yoga/test/index.spec.ts b/examples/yoga/test/index.spec.ts
@@ -68,7 +68,7 @@ describe('startup', () => {
     const body = JSON.parse(response.text);
     expect(body.data?.books).toBeUndefined();
     expect(body.errors).toBeDefined();
-    expect(body.errors?.map((e) => e.message)).toContain('Syntax Error: Query Cost limit of 100 exceeded, found 5023.');
+    expect(body.errors?.map((e) => e.message)).toContain('Syntax Error: Query Cost limit of 100 exceeded, found 138.');
   });
 
   it('should disable field suggestion', async () => {

diff --git a/packages/plugins/cost-limit/src/index.ts b/packages/plugins/cost-limit/src/index.ts
@@ -36,7 +36,7 @@ class CostLimitVisitor {
 
   private readonly context: ValidationContext;
   private readonly config: Required<CostLimitOptions>;
-  private readonly visitedFragments: Set<string> = new Set();
+  private readonly visitedFragments: Map<string, number>;
 
   constructor(context: ValidationContext, options?: CostLimitOptions) {
     this.context = context;
@@ -45,6 +45,7 @@ class CostLimitVisitor {
       costLimitDefaultOptions,
       ...Object.entries(options ?? {}).map(([k, v]) => (v === undefined ? {} : { [k]: v })),
     );
+    this.visitedFragments = new Map();
 
     this.OperationDefinition = {
       enter: this.onOperationDefinitionEnter,
@@ -84,9 +85,6 @@ class CostLimitVisitor {
       return node.selectionSet.selections.reduce((v, child) => v + this.computeComplexity(child, depth + 1), 0);
     }
 
-    // const typeDefs: GraphQLObjectType | GraphQLInterfaceType | GraphQLUnionType = this.context
-    // .getSchema()
-    // .getQueryType();
     let cost = this.config.scalarCost;
     if ('selectionSet' in node && node.selectionSet) {
       cost = this.config.objectCost;
@@ -97,12 +95,18 @@ class CostLimitVisitor {
 
     if (node.kind == Kind.FRAGMENT_SPREAD) {
       if (this.visitedFragments.has(node.name.value)) {
-        return this.config.fragmentRecursionCost;
+        const visitCost = this.visitedFragments.get(node.name.value) ?? 0;
+        return cost + this.config.depthCostFactor * visitCost;
+      } else {
+        this.visitedFragments.set(node.name.value, -1);
       }
-      this.visitedFragments.add(node.name.value);
       const fragment = this.context.getFragment(node.name.value);
       if (fragment) {
-        cost += this.config.depthCostFactor * this.computeComplexity(fragment, depth + 1);
+        const additionalCost = this.computeComplexity(fragment, depth + 1);
+        if (this.visitedFragments.get(node.name.value) === -1) {
+          this.visitedFragments.set(node.name.value, additionalCost);
+        }
+        cost += this.config.depthCostFactor * additionalCost;
       }
     }
 

diff --git a/packages/plugins/cost-limit/test/index.spec.ts b/packages/plugins/cost-limit/test/index.spec.ts
@@ -167,8 +167,6 @@ describe('global', () => {
     `);
     assertSingleExecutionValue(result);
     expect(result.errors).toBeDefined();
-    expect(result.errors?.map((error) => error.message)).toContain(
-      'Syntax Error: Query Cost limit of 50 exceeded, found 16050.',
-    );
+    expect(result.errors?.map((error) => error.message)).toContain('Cannot spread fragment "A" within itself via "B".');
   });
 });
diff --git a/packages/plugins/max-aliases/src/index.ts b/packages/plugins/max-aliases/src/index.ts
@@ -24,7 +24,7 @@ class MaxAliasesVisitor {
 
   private readonly context: ValidationContext;
   private readonly config: Required<MaxAliasesOptions>;
-  private readonly visitedFragments: Set<string> = new Set();
+  private readonly visitedFragments: Map<string, number>;
 
   constructor(context: ValidationContext, options?: MaxAliasesOptions) {
     this.context = context;
@@ -33,6 +33,7 @@ class MaxAliasesVisitor {
       maxAliasesDefaultOptions,
       ...Object.entries(options ?? {}).map(([k, v]) => (v === undefined ? {} : { [k]: v })),
     );
+    this.visitedFragments = new Map();
 
     this.OperationDefinition = {
       enter: this.onOperationDefinitionEnter,
@@ -71,12 +72,17 @@ class MaxAliasesVisitor {
       }
     } else if (node.kind === Kind.FRAGMENT_SPREAD) {
       if (this.visitedFragments.has(node.name.value)) {
-        return 0;
+        return this.visitedFragments.get(node.name.value) ?? 0;
+      } else {
+        this.visitedFragments.set(node.name.value, -1);
       }
-      this.visitedFragments.add(node.name.value);
       const fragment = this.context.getFragment(node.name.value);
       if (fragment) {
-        aliases += this.countAliases(fragment);
+        const additionalAliases = this.countAliases(fragment);
+        if (this.visitedFragments.get(node.name.value) === -1) {
+          this.visitedFragments.set(node.name.value, additionalAliases);
+        }
+        aliases += additionalAliases;
       }
     }
     return aliases;

diff --git a/packages/plugins/max-depth/src/index.ts b/packages/plugins/max-depth/src/index.ts
@@ -25,7 +25,7 @@ class MaxDepthVisitor {
 
   private readonly context: ValidationContext;
   private readonly config: Required<MaxDepthOptions>;
-  private readonly visitedFragments: Set<string> = new Set();
+  private readonly visitedFragments: Map<string, number>;
 
   constructor(context: ValidationContext, options?: MaxDepthOptions) {
     this.context = context;
@@ -34,6 +34,7 @@ class MaxDepthVisitor {
       maxDepthDefaultOptions,
       ...Object.entries(options ?? {}).map(([k, v]) => (v === undefined ? {} : { [k]: v })),
     );
+    this.visitedFragments = new Map();
 
     this.OperationDefinition = {
       enter: this.onOperationDefinitionEnter,
@@ -74,12 +75,17 @@ class MaxDepthVisitor {
       }
     } else if (node.kind == Kind.FRAGMENT_SPREAD) {
       if (this.visitedFragments.has(node.name.value)) {
-        return 0;
+        return this.visitedFragments.get(node.name.value) ?? 0;
+      } else {
+        this.visitedFragments.set(node.name.value, -1);
       }
-      this.visitedFragments.add(node.name.value);
       const fragment = this.context.getFragment(node.name.value);
       if (fragment) {
-        depth = Math.max(depth, this.countDepth(fragment, parentDepth + 1));
+        const fragmentDepth = this.countDepth(fragment, parentDepth + 1);
+        depth = Math.max(depth, fragmentDepth);
+        if (this.visitedFragments.get(node.name.value) === -1) {
+          this.visitedFragments.set(node.name.value, fragmentDepth);
+        }
       }
     }
     return depth;

diff --git a/packages/plugins/max-directives/src/index.ts b/packages/plugins/max-directives/src/index.ts
@@ -26,7 +26,7 @@ class MaxDirectivesVisitor {
 
   private readonly context: ValidationContext;
   private readonly config: Required<MaxDirectivesOptions>;
-  private readonly visitedFragments: Set<string> = new Set();
+  private readonly visitedFragments: Map<string, number>;
 
   constructor(context: ValidationContext, options?: MaxDirectivesOptions) {
     this.context = context;
@@ -35,6 +35,7 @@ class MaxDirectivesVisitor {
       maxDirectivesDefaultOptions,
       ...Object.entries(options ?? {}).map(([k, v]) => (v === undefined ? {} : { [k]: v })),
     );
+    this.visitedFragments = new Map();
 
     this.OperationDefinition = {
       enter: this.onOperationDefinitionEnter,
@@ -73,12 +74,17 @@ class MaxDirectivesVisitor {
       }
     } else if (node.kind == Kind.FRAGMENT_SPREAD) {
       if (this.visitedFragments.has(node.name.value)) {
-        return 0;
+        return this.visitedFragments.get(node.name.value) ?? 0;
+      } else {
+        this.visitedFragments.set(node.name.value, -1);
       }
-      this.visitedFragments.add(node.name.value);
       const fragment = this.context.getFragment(node.name.value);
       if (fragment) {
-        directives += this.countDirectives(fragment);
+        const additionalDirectives = this.countDirectives(fragment);
+        if (this.visitedFragments.get(node.name.value) === -1) {
+          this.visitedFragments.set(node.name.value, additionalDirectives);
+        }
+        directives += additionalDirectives;
       }
     }
     return directives;