feat restrict shield to certain routes fix #26

Author: rrd108

package.json

@@ -33,14 +33,14 @@
     "release:minor": "yarn lint && yarn test && yarn prepack && changelogen --release --minor && yarn publish && git push --follow-tags",
     "lint": "eslint .",
     "test": "vitest run --testTimeout 15000 --reporter=basic --disable-console-intercept",
-    "test:watch": "vitest watch  --testTimeout 15000 --reporter=basic --disable-console-intercept "
+    "test:watch": "vitest watch  --testTimeout 15000 --reporter=basic --disable-console-intercept"
   },
   "dependencies": {
     "@nuxt/kit": "^3.12.2",
     "defu": "^6.1.4"
   },
   "devDependencies": {
-    "@nuxt/devtools": "latest",
+    "@nuxt/devtools": "^1.4.2",
     "@nuxt/eslint-config": "^0.5.0",
     "@nuxt/module-builder": "^0.8.3",
     "@nuxt/schema": "^3.11.1",

src/runtime/server/middleware/shield.ts

@@ -19,16 +19,15 @@ export default defineEventHandler(async (event) => {
 
 
   // console.log(
-  //   `👉 Handling request for URL: ${event.node.req.url} from IP: ${
-  //     getRequestIP(event, { xForwardedFor: true }) || "unKnownIP"
+  // `👉 Handling request for URL: ${url} from IP: ${getRequestIP(event, { xForwardedFor: true }) || "unKnownIP"
   //   }`
   // );
 
   const shieldStorage = useStorage("shield");
   const requestIP = getRequestIP(event, { xForwardedFor: true }) || "unKnownIP";
 
   if (!(await shieldStorage.hasItem(`ip:${requestIP}`))) {
-    //console.log("IP not found in storage, setting initial count.", requestIP);
+    // console.log("  IP not found in storage, setting initial count.", requestIP);
     return await shieldStorage.setItem(`ip:${requestIP}`, {
       count: 1,
       time: Date.now(),
@@ -37,34 +36,29 @@ export default defineEventHandler(async (event) => {
 
   const req = (await shieldStorage.getItem(`ip:${requestIP}`)) as RateLimit;
   req.count++;
-  //console.log(`Set count for IP ${requestIP}: ${req.count}`);
+  // console.log(`  Set count for IP ${requestIP}: ${req.count}`);
 
-  shieldLog(req, requestIP, event.node.req.url);
+  shieldLog(req, requestIP, url);
 
   if (!isRateLimited(req)) {
-    //console.log("Request not rate-limited, updating storage.");
+    // console.log("  Request not rate-limited, updating storage.");
     return await shieldStorage.setItem(`ip:${requestIP}`, {
       count: req.count,
       time: req.time,
     });
   }
 
+  // console.log("  Request is rate-limited.");
+
   if (isBanExpired(req)) {
-    //console.log("Ban expired, resetting count.");
+    // console.log("  Ban expired, resetting count.");
     return await shieldStorage.setItem(`ip:${requestIP}`, {
       count: 1,
       time: Date.now(),
     });
   }
 
-  // console.log(
-  //   "setItem for IP:",
-  //   requestIP,
-  //   "with count:",
-  //   req.count,
-  //   "and time:",
-  //   req.time
-  // );
+  // console.log("  setItem for IP:", requestIP, "with count:", req.count, "and time:", req.time);
   shieldStorage.setItem(`ip:${requestIP}`, {
     count: req.count,
     time: req.time,
@@ -75,7 +69,7 @@ export default defineEventHandler(async (event) => {
   const options = useRuntimeConfig().public.nuxtApiShield;
 
   if (options.retryAfterHeader) {
-    //console.log("Setting Retry-After header", req.count + 1);
+    // console.log("  Setting Retry-After header", req.count + 1);
     event.node.res.setHeader("Retry-After", req.count + 1); // and extra second is added
   }
 
@@ -89,24 +83,24 @@ export default defineEventHandler(async (event) => {
 const isRateLimited = (req: RateLimit) => {
   const options = useRuntimeConfig().public.nuxtApiShield;
 
-  //console.log(`count: ${req.count} <= limit: ${options.limit.max}`);
+  // console.log(`  count: ${req.count} <= limit: ${options.limit.max}`);
   if (req.count <= options.limit.max) {
     return false;
   }
-  //console.log((Date.now() - req.time) / 1000, "<", options.limit.duration);
+  // console.log("  ", (Date.now() - req.time) / 1000, "<", options.limit.duration);
   return (Date.now() - req.time) / 1000 < options.limit.duration;
 };
 
 const banDelay = async (req: RateLimit) => {
   const options = useRuntimeConfig().public.nuxtApiShield;
-  //console.log("delayOnBan: " + options.delayOnBan);
+  // console.log("  delayOnBan is: " + options.delayOnBan);
   if (options.delayOnBan && req.count > options.limit.max) {
     // INFO Nuxt Devtools will send a new request if the response is slow,
     // so we get the count incremented twice or more times, based on the ban delay time
-    //console.log(`Applying ban delay for ${req.count - options.limit.max} sec`);
+    // console.log(`  Applying ban delay for ${(req.count - options.limit.max) * options.limit.ban} sec (${Date.now()})`);
     await new Promise((resolve) =>
-      setTimeout(resolve, (req.count - options.limit.max) * 1000)
+      setTimeout(resolve, (req.count - options.limit.max) * options.limit.ban * 1000)
     );
-    //console.log(`Ban delay completed`);
+    // console.log(`  Ban delay completed (${Date.now()})`);
   }
 };

test/basic.test.ts

@@ -1,35 +1,35 @@
-import { describe, it, expect } from "vitest";
+import { beforeEach, describe, it, expect } from "vitest";
 import { fileURLToPath } from "node:url";
 import { setup, $fetch } from "@nuxt/test-utils/e2e";
-import { beforeEach } from "vitest";
 import { readFile, rm } from "node:fs/promises";
 
 // TODO get these from the config
 const nuxtConfigDuration = 3;
 const nuxtConfigBan = 10;
 
-beforeEach(async () => {
-  // await useStorage("shield").clear(); TODO waiting for https://github.com/nuxt/test-utils/issues/531
-  // this is a workaround to clean the storage
-  const storagePath = fileURLToPath(new URL("../_testBasciShield", import.meta.url));
-  await rm(storagePath, { recursive: true, force: true });
-});
 
 describe("shield", async () => {
   await setup({
     rootDir: fileURLToPath(new URL("./fixtures/basic", import.meta.url)),
   });
 
+  beforeEach(async () => {
+    // await useStorage("shield").clear(); TODO waiting for https://github.com/nuxt/test-utils/issues/531
+    // this is a workaround to clean the storage
+    const storagePath = fileURLToPath(new URL("../_testBasicShield", import.meta.url));
+    await rm(storagePath, { recursive: true, force: true });
+  });
+
   it("respond to api call 2 times (limit.max, limit.duration) and rejects the 3rd call", async () => {
     // req.count = 1
-    let response = await $fetch("/api/example?c=1/1", {
+    let response = await $fetch("/api/basicexample?c=1/1", {
       method: "GET",
       retryStatusCodes: [],
     });
     expect((response as any).name).toBe("Gauranga");
 
     // req.count = 2
-    response = await $fetch("/api/example?c=1/2", {
+    response = await $fetch("/api/basicexample?c=1/2", {
       method: "GET",
       retryStatusCodes: [],
     });
@@ -39,7 +39,7 @@ describe("shield", async () => {
       // req.count = 3
       // as limit.max = 2, this should throw 429 and ban for 3 seconds (limit.ban)
       expect(async () =>
-        $fetch("/api/example?c=1/3", { method: "GET", retryStatusCodes: [] })
+        $fetch("/api/basicexample?c=1/3", { method: "GET", retryStatusCodes: [] })
       ).rejects.toThrowError();
     } catch (err) {
       const typedErr = err as { statusCode: number; statusMessage: string };
@@ -49,18 +49,15 @@ describe("shield", async () => {
   });
 
   it("respond to the 2nd api call when more then limit.duration time passes", async () => {
-    // here we should wait for the 3 sec ban to expire
-    await new Promise((resolve) => setTimeout(resolve, (nuxtConfigBan + 1) * 1000));
-
     // see #13
     // req.count = 1
-    let response = await $fetch("/api/example?c=2/1", {
+    let response = await $fetch("/api/basicexample?c=2/1", {
       method: "GET",
       retryStatusCodes: [],
     });
 
     // req.count = 2
-    response = await $fetch("/api/example?c=2/2", {
+    response = await $fetch("/api/basicexample?c=2/2", {
       method: "GET",
       retryStatusCodes: [],
     });
@@ -69,12 +66,12 @@ describe("shield", async () => {
 
   it("respond to api call after limit.ban expires", async () => {
     // req.count reset here
-    await $fetch("/api/example?c=3/1", { method: "GET", retryStatusCodes: [] }); // req.count = 1
-    await $fetch("/api/example?c=3/2", { method: "GET", retryStatusCodes: [] }); // req.count = 2
+    await $fetch("/api/basicexample?c=3/1", { method: "GET", retryStatusCodes: [] }); // req.count = 1
+    await $fetch("/api/basicexample?c=3/2", { method: "GET", retryStatusCodes: [] }); // req.count = 2
     try {
       // req.count = 3
       expect(async () =>
-        $fetch("/api/example?c=3/3", { method: "GET", retryStatusCodes: [] })
+        $fetch("/api/basicexample?c=3/3", { method: "GET", retryStatusCodes: [] })
       ).rejects.toThrowError();
     } catch (err) {
       const typedErr = err as {
@@ -90,7 +87,7 @@ describe("shield", async () => {
 
     // here we should wait for the 3 sec ban to expire
     await new Promise((resolve) => setTimeout(resolve, nuxtConfigBan * 1000));
-    const response = await $fetch("/api/example?c=3/4", {
+    const response = await $fetch("/api/basicexample?c=3/4", {
       method: "GET",
       retryStatusCodes: [],
     });