Root Level Limit: different take with onValidate

Author: ardatan

packages/plugins/root-level-limitation/__tests__/root-level-limitation.spec.ts

@@ -64,4 +64,51 @@ describe('root-level-limitation', () => {
     });
     expect(res.status).toBe(200);
   });
+
+  it('should not allow requests with max root level query and nested fragments', async () => {
+    const res = await yoga.fetch('http://yoga/graphql', {
+      method: 'POST',
+      body: JSON.stringify({
+        query: /* GraphQL */ `
+          fragment QueryFragment on Query {
+            topBooks {
+              id
+            }
+            topProducts {
+              id
+            }
+          }
+          {
+            ...QueryFragment
+          }
+        `,
+      }),
+      headers: {
+        'Content-Type': 'application/json',
+      },
+    });
+  })
+
+  it('should allow requests with max root level query in comments', async () => {
+    const res = await yoga.fetch('http://yoga/graphql', {
+      method: 'POST',
+      body: JSON.stringify({
+        query: /* GraphQL */ `
+          {
+            # topBooks {
+            #   id
+            # }
+            topProducts {
+              id
+            }
+          }
+        `,
+      }),
+      headers: {
+        'Content-Type': 'application/json',
+      },
+    });
+
+    expect(res.status).toBe(200);
+  })
 });

packages/plugins/root-level-limitation/package.json

@@ -38,15 +38,14 @@
   },
   "peerDependencies": {
     "@graphql-tools/utils": "^10.1.0",
-    "@whatwg-node/server": "^0.9.32",
-    "graphql": "^15.2.0 || ^16.0.0",
-    "graphql-yoga": "^5.3.0"
+    "@envelop/core": "^5.0.0",
+    "graphql": "^15.2.0 || ^16.0.0"
   },
   "dependencies": {
     "tslib": "^2.5.2"
   },
   "devDependencies": {
-    "@whatwg-node/fetch": "^0.9.17",
+    "@envelop/core": "^5.0.0",
     "graphql": "^16.6.0",
     "graphql-yoga": "5.3.0"
   },

packages/plugins/root-level-limitation/src/index.ts

@@ -1,79 +1,47 @@
-import { createGraphQLError } from '@graphql-tools/utils';
+import { createGraphQLError, getRootTypes } from '@graphql-tools/utils';
+import { Plugin } from '@envelop/core';
+import type { ValidationRule } from 'graphql/validation/ValidationContext';
+import { isObjectType } from 'graphql';
 
-interface GraphQLParams {
-  operationName?: string;
-  query?: string;
+export interface RootLevelQueryLimitOptions {
+  maxRootLevelFields: number;
 }
 
-export function rootLevelQueryLimit({ maxRootLevelFields }: { maxRootLevelFields: number }) {
-  return {
-    onParams({ params }: unknown) {
-      const { query, operationName } = params as GraphQLParams;
-
-      if (operationName?.includes('IntrospectionQuery')) return true;
-
-      const newQuery = formatQuery(query || '');
-      const linesArray = newQuery.split('\n');
-
-      let countLeadingSpacesTwo = 0;
-
-      for (const line of linesArray) {
-        const leadingSpaces = line?.match(/^\s*/)?.[0]?.length || 0;
-
-        if (leadingSpaces === 4 && line[leadingSpaces] !== ')') {
-          countLeadingSpacesTwo++;
-
-          if (countLeadingSpacesTwo > maxRootLevelFields * 2) {
+export function createRootLevelQueryLimitRule(opts: RootLevelQueryLimitOptions): ValidationRule {
+  const { maxRootLevelFields } = opts;
+
+  return function rootLevelQueryLimitRule (context) {
+    const rootTypes = getRootTypes(context.getSchema());
+    let rootFieldCount = 0;
+    return {
+      Field() {
+        const parentType = context.getParentType();
+        if (isObjectType(parentType) && rootTypes.has(parentType)) {
+          rootFieldCount++;
+          if (rootFieldCount > maxRootLevelFields) {
             throw createGraphQLError('Query is too complex.', {
               extensions: {
                 http: {
                   spec: false,
                   status: 400,
-                  headers: {
-                    Allow: 'POST',
-                  },
                 },
               },
             });
           }
         }
-      }
-
-      return true;
-    },
+      },
+    };
   };
-}
 
-function formatQuery(queryString: string) {
-  queryString = queryString.replace(/^\s+/gm, '');
-
-  let indentLevel = 0;
-  let formattedString = '';
-
-  for (let i = 0; i < queryString.length; i++) {
-    const char = queryString[i];
-
-    if (char === '{' || char === '(') {
-      formattedString += char;
-      indentLevel++;
-      // formattedString += ' '.repeat(indentLevel * 4);
-    } else if (char === '}' || char === ')') {
-      indentLevel--;
-
-      if (formattedString[formattedString.length - 1] !== '\n')
-        formattedString = formattedString.trim().replace(/\n$/, '');
-
-      if (char === ')') formattedString += char;
+}
 
-      if (char === '}') formattedString += '\n' + ' '.repeat(indentLevel * 4) + char;
-    } else if (char === '\n') {
-      if (queryString[i + 1] !== '\n' && queryString[i + 1] !== undefined) {
-        formattedString += char + ' '.repeat(indentLevel * 4);
-      }
-    } else {
-      formattedString += char;
+export function rootLevelQueryLimit(opts: RootLevelQueryLimitOptions): Plugin {
+  const rootLevelQueryLimitRule = createRootLevelQueryLimitRule(opts);
+  return {
+    onValidate({ addValidationRule }) {
+      addValidationRule(
+        rootLevelQueryLimitRule
+      )
     }
   }
-
-  return formattedString;
 }