feat: enable ratelimiting at query level, with dedicated query/unit i…

…n the configuration
authdog · Dec 29, 2023 · 7959cd6 · 7959cd6
1 parent f68ff69
commit 7959cd6
Show file tree

Hide file tree

Showing 5 changed files with 85 additions and 43 deletions.
diff --git a/.changeset/wicked-spiders-vanish.md b/.changeset/wicked-spiders-vanish.md
@@ -0,0 +1,6 @@
+---
+"@authdog/hydra-core": patch
+"@authdog/hydra-cli": patch
+---
+
+enable ratelimiting at query level, with dedicated query/unit in the configuration
diff --git a/packages/cli/src/utils/validateConfig.ts b/packages/cli/src/utils/validateConfig.ts
@@ -38,8 +38,7 @@ const hydraConfigSchema = z.object({
 export const validateConfig = (config: IHydraConfig): IHydraConfig => {
   try {
     // Validate the provided config object against the Zod schema
-    const validatedConfig = hydraConfigSchema.parse(config);
-    return validatedConfig;
+    return hydraConfigSchema.parse(config);
   } catch (error) {
     throw error;
   }

diff --git a/packages/core/src/handlers/hydra.ts b/packages/core/src/handlers/hydra.ts
@@ -21,7 +21,6 @@ export const HydraHandler = async (req, env, ctx): Promise<Response> => {
 
   const { kv, hydraConfig, rateLimiter } = ctx;
   let cacheKey = null;
-  // const { RateLimiter } = req;
   // get ip address
   const ip =
     req.headers.get("cf-connecting-ip") ||
@@ -43,11 +42,10 @@ export const HydraHandler = async (req, env, ctx): Promise<Response> => {
     ],
   };
 
-  const facetId = ip;
+  const facetId = ip || "localhost" // TODO: extend to more facets combinations
   const defaultRateLimitingBudget = hydraConfig.rateLimiting.default.budget;
   let remainingRateBudget = -1;
 
-
   const kvNamespace = kv;
   const requestHeaders = req.clone()?.headers;
   const requestBody = await req.clone()?.json();
@@ -76,33 +74,64 @@ export const HydraHandler = async (req, env, ctx): Promise<Response> => {
     );
 
     if (rateLimiter && !isIntrospection && !isMutation) {
-
-      const timestamp = new Date().toISOString().slice(0, 16).replace(/[-:T]/g, "");
       const facetQueriesIds = extractedQueries.map((query) => {
-        return `${facetId}_${query}`;
-      })
 
-      // TODO: generate map query -> rateCount
-      const rateCounts = await Promise.all(facetQueriesIds.map(async (facetQueryId) => {
-        const rateCount = await fetchRateLimiterWithFacet(req, rateLimiter, facetQueryId, timestamp);
-        return rateCount;
-      }));
+        const facetQueryId = `${facetId}_${query}`;
+        const queryId = hydraConfig.rateLimiting?.queries?.find((queryConfig) => {
+          return queryConfig.id === query;
+        })?.id || "default";
+        const queryBudget = hydraConfig.rateLimiting?.queries?.find((queryConfig) => {
+          return queryConfig.id === query;
+        })?.budget || defaultRateLimitingBudget;
+        const queryBudgetUnit = hydraConfig.rateLimiting?.queries?.find((queryConfig) => {
+          return queryConfig.id === query;
+        })?.unit || "minute";
+
+        return {
+          facetQueryId,
+          queryId,
+          queryBudget,
+          queryBudgetUnit,
+        };
 
+      })
       let hasAtLeastOneExceeded = false;
 
-      rateCounts.map((rateCount) => {
-        // TODO: read from config budget for given query
-        if (Number(rateCount) > defaultRateLimitingBudget) {
+      const rateCountsReports = await Promise.all(facetQueriesIds.map(async (facetObj) => {
+        let timestamp;
+        if (facetObj.queryBudgetUnit === "hour") {
+          // generate timestamp removing minutes and seconds
+          timestamp = new Date().toISOString().slice(0, 16).replace(/[-:T]/g, "").slice(0, 10);
+        } else {
+          // generate timestamp removing seconds
+          timestamp = new Date().toISOString().slice(0, 16).replace(/[-:T]/g, "").slice(0, 12);
+        }
+        const rateCount = await fetchRateLimiterWithFacet(req, rateLimiter, facetObj.facetQueryId, timestamp);
+
+        if (rateCount > facetObj.queryBudget) {
           hasAtLeastOneExceeded = true;
-          return true;
         }
-      })
+
+        return {
+          facetQueryId: facetObj.facetQueryId,
+          queryBudget: facetObj.queryBudget,
+          queryBudgetUnit: facetObj.queryBudgetUnit,
+          rateCount
+        }
+      }));
+
+      console.log("rateCountsReports", JSON.stringify(rateCountsReports, null, 2));
+
+      const excedeedRateCountReports = rateCountsReports.filter((report) => {
+        return report.rateCount > report.queryBudget;
+      });
+
 
-      if (hasAtLeastOneExceeded) {
+      if (excedeedRateCountReports?.length > 0) {
         const errorResponse = {
           errors: [
             {
-              message: "Too many requests",
+              message: `Too many requests for ${excedeedRateCountReports[0]?.facetQueryId}`,
               extensions: {
                 code: "TOO_MANY_REQUESTS",
                 statusCode: 429,

diff --git a/services/itty-hydra/hydra.config.ts b/services/itty-hydra/hydra.config.ts
@@ -13,10 +13,18 @@ export const HydraConfigAcme = {
     default: {
       budget: 20,
     },
-    health: {
-      budget: 5,
-      unit: "minute",
-    }
+    queries: [
+      {
+        id: "health",
+        budget: 5,
+        unit: "minute",
+      },
+      {
+        id: "hydraDevQuery",
+        budget: 15,
+        unit: "minute",
+      },
+    ]
   },
   publicQueries: [
     {

diff --git a/services/itty-hydra/main.ts b/services/itty-hydra/main.ts
@@ -27,24 +27,24 @@ router
   .get("/health", Health)
   .get("/graphql", GraphQLHandler)
   .post("/graphql", HydraHandler)
-  .get("/counter", async (req, { HydraRateLimiter }) => {
-    try {
-      const jsonResponse = await fetchWithRateLimiter(
-        req,
-        HydraRateLimiter,
-        "testFacet",
-      );
-      return new Response(JSON.stringify(jsonResponse), {
-        status: 200,
-        headers: {
-          "content-type": "text/plain",
-        },
-      });
-    } catch (error) {
-      // Handle any errors that might occur during the process
-      return new Response("Internal Server Error", { status: 500 });
-    }
-  })
+  // .get("/counter", async (req, { HydraRateLimiter }) => {
+  //   try {
+  //     const jsonResponse = await fetchWithRateLimiter(
+  //       req,
+  //       HydraRateLimiter,
+  //       "testFacet",
+  //     );
+  //     return new Response(JSON.stringify(jsonResponse), {
+  //       status: 200,
+  //       headers: {
+  //         "content-type": "text/plain",
+  //       },
+  //     });
+  //   } catch (error) {
+  //     // Handle any errors that might occur during the process
+  //     return new Response("Internal Server Error", { status: 500 });
+  //   }
+  // })
   .all("*", NotFound);
 
 const handleRequest = async (req, env, ctx) => {