Feat/flatten fragments cost plugin (#724)

* Update cost-limit.md Signed-off-by: ideeeveee <[email protected]> * add flattenFragments flag to cost-limit plugin Introduce config to control if depthCostFactor should be applied to fragments when calculating the cost Signed-off-by: ideeeveee <[email protected]> * Add tests Signed-off-by: ideeeveee <[email protected]> * feat: changeset --------- Signed-off-by: ideeeveee <[email protected]> Co-authored-by: nullswan <[email protected]>
Escape-Technologies · Oct 25, 2024 · bc9e2d0 · bc9e2d0
1 parent 9c39455
commit bc9e2d0
Show file tree

Hide file tree

Showing 4 changed files with 96 additions and 8 deletions.
diff --git a/.changeset/bright-kings-think.md b/.changeset/bright-kings-think.md
@@ -0,0 +1,6 @@
+---
+'@escape.tech/graphql-armor-cost-limit': minor
+'@escape.tech/graphql-armor': patch
+---
+
+cost limit: flatten fragment option support
diff --git a/packages/plugins/cost-limit/src/index.ts b/packages/plugins/cost-limit/src/index.ts
@@ -16,6 +16,7 @@ export type CostLimitOptions = {
   objectCost?: number;
   scalarCost?: number;
   depthCostFactor?: number;
+  flattenFragments?: boolean;
   ignoreIntrospection?: boolean;
   fragmentRecursionCost?: number;
   exposeLimits?: boolean;
@@ -27,6 +28,7 @@ const costLimitDefaultOptions: Required<CostLimitOptions> = {
   objectCost: 2,
   scalarCost: 1,
   depthCostFactor: 1.5,
+  flattenFragments: false,
   fragmentRecursionCost: 1000,
   ignoreIntrospection: true,
   exposeLimits: true,
@@ -95,11 +97,16 @@ class CostLimitVisitor {
     if ('selectionSet' in node && node.selectionSet) {
       cost = this.config.objectCost;
       for (const child of node.selectionSet.selections) {
-        cost += this.config.depthCostFactor * this.computeComplexity(child, depth + 1);
+        if (
+          this.config.flattenFragments &&
+          (child.kind === Kind.INLINE_FRAGMENT || child.kind === Kind.FRAGMENT_SPREAD)
+        ) {
+          cost += this.computeComplexity(child, depth);
+        } else {
+          cost += this.config.depthCostFactor * this.computeComplexity(child, depth + 1);
+        }
       }
-    }
-
-    if (node.kind === Kind.FRAGMENT_SPREAD) {
+    } else if (node.kind === Kind.FRAGMENT_SPREAD) {
       if (this.visitedFragments.has(node.name.value)) {
         const visitCost = this.visitedFragments.get(node.name.value) ?? 0;
         return cost + this.config.depthCostFactor * visitCost;
@@ -108,14 +115,19 @@ class CostLimitVisitor {
       }
       const fragment = this.context.getFragment(node.name.value);
       if (fragment) {
-        const additionalCost = this.computeComplexity(fragment, depth + 1);
+        let fragmentCost;
+        if (this.config.flattenFragments) {
+          fragmentCost = this.computeComplexity(fragment, depth);
+          cost += fragmentCost;
+        } else {
+          fragmentCost = this.computeComplexity(fragment, depth + 1);
+          cost += this.config.depthCostFactor * fragmentCost;
+        }
         if (this.visitedFragments.get(node.name.value) === -1) {
-          this.visitedFragments.set(node.name.value, additionalCost);
+          this.visitedFragments.set(node.name.value, fragmentCost);
         }
-        cost += this.config.depthCostFactor * additionalCost;
       }
     }
-
     return cost;
   }
 }

diff --git a/packages/plugins/cost-limit/test/index.spec.ts b/packages/plugins/cost-limit/test/index.spec.ts
@@ -141,6 +141,73 @@ describe('costLimitPlugin', () => {
     ]);
   });
 
+  it('should flatten fragments', async () => {
+    const maxCost = 57;
+    const testkit = createTestkit(
+      [
+        costLimitPlugin({
+          maxCost: maxCost,
+          objectCost: 4,
+          scalarCost: 2,
+          depthCostFactor: 2,
+          flattenFragments: true,
+          ignoreIntrospection: true,
+        }),
+      ],
+      schema,
+    );
+    const result = await testkit.execute(`
+    query {
+      ...BookFragment
+    }
+
+    fragment BookFragment on Query {
+      books {
+        title
+        author
+      }
+    }
+    `);
+
+    assertSingleExecutionValue(result);
+    expect(result.errors).toBeUndefined();
+  });
+
+  it('should reject flattened fragments if the maxCost limit is exceeded', async () => {
+    const maxCost = 5;
+    const testkit = createTestkit(
+      [
+        costLimitPlugin({
+          maxCost: maxCost,
+          objectCost: 4,
+          scalarCost: 2,
+          depthCostFactor: 2,
+          flattenFragments: true,
+          ignoreIntrospection: true,
+        }),
+      ],
+      schema,
+    );
+    const result = await testkit.execute(`
+    query {
+      ...BookFragment
+    }
+
+    fragment BookFragment on Query {
+      books {
+        title
+        author
+      }
+    }
+    `);
+
+    assertSingleExecutionValue(result);
+    expect(result.errors).toBeDefined();
+    expect(result.errors?.map((error) => error.message)).toEqual([
+      `Syntax Error: Query Cost limit of ${maxCost} exceeded, found 30.`,
+    ]);
+  });
+
   it('should not crash on recursive fragment', async () => {
     const testkit = createTestkit(
       [

diff --git a/services/docs/docs/plugins/cost-limit.md b/services/docs/docs/plugins/cost-limit.md
@@ -29,6 +29,9 @@ GraphQLArmorConfig({
     // Factorial applied to nested operator | default: 1.5
     depthCostFactor?: int,
 
+    // Flatten frament spreads and inline framents for the cost calculation | default: false
+    flattenFragments?: boolean,
+
     // Ignore the cost of introspection queries | default: true
     ignoreIntrospection?: boolean,