Count rate limits for GraphQL fields

Author: timokoessler

library/agent/Agent.ts

@@ -573,6 +573,15 @@ export class Agent {
     });
   }
 
+  onGraphQLFieldRateLimited(
+    method: string,
+    path: string,
+    type: "query" | "mutation",
+    field: string
+  ) {
+    this.routes.countGraphQLFieldRateLimited(method, path, type, field);
+  }
+
   getRoutes() {
     return this.routes;
   }

library/agent/Routes.test.ts

@@ -980,3 +980,56 @@ t.test("it counts rate limited requests", async (t) => {
     },
   ]);
 });
+
+t.test("it counts rate limited GraphQL fields", async (t) => {
+  const routes = new Routes(200);
+  routes.addRoute(getContext("POST", "/graphql"));
+  routes.addGraphQLField("POST", "/graphql", "query", "user");
+
+  routes.countGraphQLFieldRateLimited("POST", "/graphql", "query", "user");
+
+  t.same(routes.asArray(), [
+    {
+      method: "POST",
+      path: "/graphql",
+      hits: 1,
+      rateLimitedCount: 0,
+      graphql: undefined,
+      apispec: {},
+      graphQLSchema: undefined,
+    },
+    {
+      method: "POST",
+      path: "/graphql",
+      hits: 1,
+      rateLimitedCount: 1,
+      graphql: { type: "query", name: "user" },
+      apispec: {},
+      graphQLSchema: undefined,
+    },
+  ]);
+
+  routes.addGraphQLField("POST", "/graphql", "query", "user");
+  routes.countGraphQLFieldRateLimited("POST", "/graphql", "query", "user");
+
+  t.same(routes.asArray(), [
+    {
+      method: "POST",
+      path: "/graphql",
+      hits: 1,
+      rateLimitedCount: 0,
+      graphql: undefined,
+      apispec: {},
+      graphQLSchema: undefined,
+    },
+    {
+      method: "POST",
+      path: "/graphql",
+      hits: 2,
+      rateLimitedCount: 2,
+      graphql: { type: "query", name: "user" },
+      apispec: {},
+      graphQLSchema: undefined,
+    },
+  ]);
+});

library/agent/Routes.ts

@@ -140,6 +140,21 @@ export class Routes {
     }
   }
 
+  countGraphQLFieldRateLimited(
+    method: string,
+    path: string,
+    type: "query" | "mutation",
+    name: string
+  ) {
+    const key = this.getGraphQLKey(method, path, type, name);
+    const existing = this.routes.get(key);
+
+    if (existing) {
+      existing.rateLimitedCount++;
+      return;
+    }
+  }
+
   clear() {
     this.routes.clear();
   }

library/sources/GraphQL.test.ts

@@ -11,7 +11,9 @@ function getTestContext() {
     method: "POST",
     url: "http://localhost:4000/graphql",
     query: {},
-    headers: {},
+    headers: {
+      "content-type": "application/json",
+    },
     body: { query: '{ getFile(path: "/etc/bashrc") }' },
     cookies: {},
     routeParams: {},
@@ -96,11 +98,13 @@ t.test("it works", async () => {
   await runWithContext(getTestContext(), async () => {
     const success = await query("/etc/bashrc");
     t.same(success.data!.getFile, "file content");
-    await query("/etc/bashrc");
-    await query("/etc/bashrc");
+
     const result = await query("/etc/bashrc");
     t.same(result.errors![0].message, "You are rate limited by Zen.");
 
+    await query("/etc/bashrc");
+    await query("/etc/bashrc");
+
     // With operation name
     t.same(
       await graphql({
@@ -124,6 +128,33 @@ t.test("it works", async () => {
     );
   });
 
+  t.same(agent.getRoutes().asArray(), [
+    {
+      method: "POST",
+      path: "/graphql",
+      hits: 6,
+      rateLimitedCount: 3,
+      graphql: {
+        type: "query",
+        name: "getFile",
+      },
+      apispec: {},
+      graphQLSchema: undefined,
+    },
+    {
+      method: "POST",
+      path: "/graphql",
+      hits: 1,
+      rateLimitedCount: 0,
+      graphql: {
+        type: "query",
+        name: "anotherQuery",
+      },
+      apispec: {},
+      graphQLSchema: undefined,
+    },
+  ]);
+
   // Empty context
   await runWithContext({} as Context, async () => {
     const response = await query("/etc/bashrc");

library/sources/GraphQL.tools.test.ts

@@ -66,7 +66,7 @@ t.test("it works", async () => {
     method: "POST",
     url: "http://localhost:4000/graphql",
     query: {},
-    headers: {},
+    headers: { "content-type": "application/json" },
     body: { query: '{ getFile(path: "/etc/bashrc") }' },
     cookies: {},
     routeParams: {},
@@ -85,9 +85,27 @@ t.test("it works", async () => {
   await runWithContext(context, async () => {
     const success = await query("/etc/bashrc");
     t.same(success.data.getFile, "file content");
+
     await query("/etc/bashrc");
-    await query("/etc/bashrc");
+
     const result = await query("/etc/bashrc");
     t.same(result.errors[0].message, "You are rate limited by Zen.");
+
+    await query("/etc/bashrc");
+
+    t.same(agent.getRoutes().asArray(), [
+      {
+        method: "POST",
+        path: "/graphql",
+        hits: 5,
+        rateLimitedCount: 2,
+        graphql: {
+          type: "query",
+          name: "getFile",
+        },
+        apispec: {},
+        graphQLSchema: undefined,
+      },
+    ]);
   });
 });

library/sources/GraphQL.ts

@@ -142,6 +142,15 @@ export class GraphQL implements Wrapper {
       agent.getInspectionStatistics().onRateLimitedRequest();
       updateContext(context, "rateLimited", true);
 
+      if (context && context.method && context.route) {
+        agent.onGraphQLFieldRateLimited(
+          context.method,
+          context.route,
+          result.operationType,
+          result.field.name.value
+        );
+      }
+
       return {
         errors: [
           new this.graphqlModule.GraphQLError("You are rate limited by Zen.", {

library/sources/graphql/shouldRateLimitOperation.test.ts

@@ -105,6 +105,9 @@ t.test("it rate limits query", async () => {
 
   t.match(shouldRateLimitOperation(agent, context, args), {
     block: true,
+    operationType: "query",
+    source: "ip",
+    remoteAddress: "1.2.3.4",
   });
 });
 
@@ -170,5 +173,8 @@ t.test("it rate limits mutation", async () => {
 
   t.match(shouldRateLimitOperation(agent, context, args), {
     block: true,
+    operationType: "mutation",
+    source: "ip",
+    remoteAddress: "1.2.3.4",
   });
 });

library/sources/graphql/shouldRateLimitOperation.ts

@@ -14,12 +14,14 @@ type Result =
       field: FieldNode;
       source: "ip";
       remoteAddress: string;
+      operationType: "query" | "mutation";
     }
   | {
       block: true;
       field: FieldNode;
       source: "user";
       userId: string;
+      operationType: "query" | "mutation";
     };
 
 export function shouldRateLimitOperation(
@@ -109,6 +111,7 @@ function shouldRateLimitField(
         field: field,
         source: "ip",
         remoteAddress: context.remoteAddress,
+        operationType: operationType,
       };
     }
   }
@@ -128,6 +131,7 @@ function shouldRateLimitField(
         field: field,
         source: "user",
         userId: context.user.id,
+        operationType: operationType,
       };
     }
   }