Add maxQueryNodes limit (#98)

Author: ivome

README.md

@@ -31,7 +31,7 @@ import {
 
 const rule = createComplexityRule({
   // The maximum allowed query complexity, queries above this threshold will be rejected
-  maximumComplexity: 1000,
+  maximumComplexity: 1_000,
 
   // The query variables. This is needed because the variables are not available
   // in the visitor of the graphql-js library
@@ -40,9 +40,16 @@ const rule = createComplexityRule({
   // The context object for the request (optional)
   context: {}
 
-  // specify operation name only when pass multi-operation documents
+  // Specify operation name when evaluating multi-operation documents
   operationName?: string,
 
+  // The maximum number of query nodes to evaluate (fields, fragments, composite types).
+  // If a query contains more than the specified number of nodes, the complexity rule will
+  // throw an error, regardless of the complexity of the query.
+  //
+  // Default: 10_000
+  maxQueryNodes?: 10_000,
+
   // Optional callback function to retrieve the determined query complexity
   // Will be invoked whether the query is rejected or not
   // This can be used for logging or to implement rate limiting

src/QueryComplexity.ts

@@ -85,6 +85,11 @@ export interface QueryComplexityOptions {
 
   // Pass request context to the estimators via estimationContext
   context?: Record<string, any>;
+
+  // The maximum number of nodes to evaluate. If this is set, the query will be
+  // rejected if it exceeds this number. (Includes fields, fragments, inline fragments, etc.)
+  // Defaults to 10_000.
+  maxQueryNodes?: number;
 }
 
 function queryComplexityMessage(max: number, actual: number): string {
@@ -101,6 +106,7 @@ export function getComplexity(options: {
   variables?: Record<string, any>;
   operationName?: string;
   context?: Record<string, any>;
+  maxQueryNodes?: number;
 }): number {
   const typeInfo = new TypeInfo(options.schema);
 
@@ -118,6 +124,7 @@ export function getComplexity(options: {
     variables: options.variables,
     operationName: options.operationName,
     context: options.context,
+    maxQueryNodes: options.maxQueryNodes,
   });
 
   visit(options.query, visitWithTypeInfo(typeInfo, visitor));
@@ -140,6 +147,8 @@ export default class QueryComplexity {
   skipDirectiveDef: GraphQLDirective;
   variableValues: Record<string, any>;
   requestContext?: Record<string, any>;
+  evaluatedNodes: number;
+  maxQueryNodes: number;
 
   constructor(context: ValidationContext, options: QueryComplexityOptions) {
     if (
@@ -154,7 +163,8 @@ export default class QueryComplexity {
     this.context = context;
     this.complexity = 0;
     this.options = options;
-
+    this.evaluatedNodes = 0;
+    this.maxQueryNodes = options.maxQueryNodes ?? 10_000;
     this.includeDirectiveDef = this.context.getSchema().getDirective('include');
     this.skipDirectiveDef = this.context.getSchema().getDirective('skip');
     this.estimators = options.estimators;
@@ -274,7 +284,12 @@ export default class QueryComplexity {
             complexities: ComplexityMap,
             childNode: FieldNode | FragmentSpreadNode | InlineFragmentNode
           ): ComplexityMap => {
-            // let nodeComplexity = 0;
+            this.evaluatedNodes++;
+            if (this.evaluatedNodes >= this.maxQueryNodes) {
+              throw new GraphQLError(
+                'Query exceeds the maximum allowed number of nodes.'
+              );
+            }
             let innerComplexities = complexities;
 
             let includeNode = true;

src/__tests__/QueryComplexity-test.ts

@@ -939,4 +939,87 @@ describe('QueryComplexity analysis', () => {
 
     expect(errors).to.have.length(0);
   });
+
+  it('should reject queries that exceed the maximum number of fragment nodes', () => {
+    const query = parse(`
+      query {
+        ...F
+        ...F
+      }
+      fragment F on Query {
+        scalar
+      }
+    `);
+
+    expect(() =>
+      getComplexity({
+        estimators: [simpleEstimator({ defaultComplexity: 1 })],
+        schema,
+        query,
+        maxQueryNodes: 1,
+        variables: {},
+      })
+    ).to.throw('Query exceeds the maximum allowed number of nodes.');
+  });
+
+  it('should reject queries that exceed the maximum number of field nodes', () => {
+    const query = parse(`
+      query {
+        scalar
+        scalar1: scalar
+        scalar2: scalar
+        scalar3: scalar
+        scalar4: scalar
+        scalar5: scalar
+        scalar6: scalar
+        scalar7: scalar
+      }
+    `);
+
+    expect(() =>
+      getComplexity({
+        estimators: [simpleEstimator({ defaultComplexity: 1 })],
+        schema,
+        query,
+        maxQueryNodes: 1,
+        variables: {},
+      })
+    ).to.throw('Query exceeds the maximum allowed number of nodes.');
+  });
+
+  it('should limit the number of query nodes to 10_000 by default', () => {
+    const failingQuery = parse(`
+      query {
+        ${Array.from({ length: 10_000 }, (_, i) => `scalar${i}: scalar`).join(
+          '\n'
+        )}
+      }
+    `);
+
+    expect(() =>
+      getComplexity({
+        estimators: [simpleEstimator({ defaultComplexity: 1 })],
+        schema,
+        query: failingQuery,
+        variables: {},
+      })
+    ).to.throw('Query exceeds the maximum allowed number of nodes.');
+
+    const passingQuery = parse(`
+      query {
+        ${Array.from({ length: 9999 }, (_, i) => `scalar${i}: scalar`).join(
+          '\n'
+        )}
+      }
+    `);
+
+    expect(() =>
+      getComplexity({
+        estimators: [simpleEstimator({ defaultComplexity: 1 })],
+        schema,
+        query: passingQuery,
+        variables: {},
+      })
+    ).not.to.throw();
+  });
 });