fix: failing test

Author: rrd108

src/runtime/server/middleware/shield.ts

@@ -23,11 +23,7 @@ export default defineEventHandler(async (event) => {
   // );
 
   const shieldStorage = useStorage('shield')
-  // Determine whether to trust X-Forwarded-For header based on config.
-  // The module's default for `config.security.trustXForwardedFor` is true.
-  const trustXForwardedFor = config.security?.trustXForwardedFor;
-
-  const requestIP = getRequestIP(event, { xForwardedFor: trustXForwardedFor }) || 'unKnownIP'
+  const requestIP = getRequestIP(event, { xForwardedFor: true }) || 'unKnownIP'
 
   const banKey = `ban:${requestIP}`
   const bannedUntilRaw = await shieldStorage.getItem(banKey)
@@ -43,79 +39,43 @@ export default defineEventHandler(async (event) => {
     })
   }
   else if (bannedUntilRaw && !Number.isNaN(bannedUntil) && Date.now() >= bannedUntil) {
-    // Ban expired, remove ban and reset IP count
     await shieldStorage.removeItem(banKey)
-    // When a ban expires, we reset the IP's rate limit data as if it's a new IP.
-    // This ensures that the count doesn't carry over from before the ban.
-    await shieldStorage.setItem(`ip:${requestIP}`, { count: 1, time: Date.now() })
-    // Note: The request that triggers this ban removal is counted as the first request in the new window.
-    // Thus, we don't 'return' here but proceed to process this request.
+    await shieldStorage.setItem(`ip:${requestIP}`, {
+      count: 1,
+      time: Date.now(),
+    })
   }
 
-  // IP Rate Limiting Logic
-  const ipKey = `ip:${requestIP}`
-  let currentRateLimitData = (await shieldStorage.getItem(ipKey)) as RateLimit | null
-  const currentTime = Date.now()
-
-  if (!currentRateLimitData) {
-    // First request from this IP (or first after a ban was just cleared above)
-    currentRateLimitData = { count: 1, time: currentTime }
-  }
-  else {
-    // Existing IP, check if window needs reset
-    if ((currentTime - currentRateLimitData.time) / 1000 >= config.limit.duration) {
-      // Window expired, reset
-      currentRateLimitData.count = 1
-      currentRateLimitData.time = currentTime
-    }
-    else {
-      // Window still active, increment count
-      currentRateLimitData.count++
-    }
+  if (!(await shieldStorage.hasItem(`ip:${requestIP}`))) {
+    // console.log("  IP not found in storage, setting initial count.", requestIP);
+    return await shieldStorage.setItem(`ip:${requestIP}`, {
+      count: 1,
+      time: Date.now(),
+    })
   }
 
-  const req = currentRateLimitData // This is the state *after* accounting for the current request
-
-  shieldLog(req, requestIP, url.toString()) // Log the current state
-
-  if (isRateLimited(req)) { // isRateLimited now checks if this current state (req) warrants a ban
-    const banUntilTimestamp = Date.now() + config.limit.ban * 1000
-    await shieldStorage.setItem(banKey, banUntilTimestamp)
-    // When a user is banned, their IP tracking data should also be reset
-    // so that after the ban, they start with a fresh count.
-    await shieldStorage.setItem(ipKey, { count: 1, time: Date.now() })
+  const req = (await shieldStorage.getItem(`ip:${requestIP}`)) as RateLimit
+  req.count++
+  // console.log(`  Set count for IP ${requestIP}: ${req.count}`);
 
-    if (config.retryAfterHeader) {
-      event.node.res.setHeader('Retry-After', config.limit.ban) // Duration of ban in seconds
-    }
+  shieldLog(req, requestIP, url.toString())
 
-    // console.error("Throwing 429 error due to rate limit ban");
-    throw createError({
-      statusCode: 429,
-      message: config.errorMessage,
-    })
-  }
-  else {
-    // Not banned, just update the storage with the new count/time
-    await shieldStorage.setItem(ipKey, {
+  if (!isRateLimited(req)) {
+    // console.log("  Request not rate-limited, updating storage.");
+    return await shieldStorage.setItem(`ip:${requestIP}`, {
       count: req.count,
       time: req.time,
     })
-    // Request is allowed, proceed to the actual API handler
-    return
   }
-})
 
-const isRateLimited = (req: RateLimit) => { // req here is the state *after* current request is counted
-  const options = useRuntimeConfig().public.nuxtApiShield
-  // If count is over limit AND it happened within the duration since req.time
-  if (req.count > options.limit.max) {
-    // Check if these excess requests occurred within the defined duration
-    // from the start of the current window (req.time)
-    return (Date.now() - req.time) / 1000 < options.limit.duration
-  }
-  return false // Count is not over max, so not rate limited
-}
+  const banUntil = Date.now() + config.limit.ban * 1000
+  await shieldStorage.setItem(banKey, banUntil)
+  await shieldStorage.setItem(`ip:${requestIP}`, {
+    count: 1,
+    time: Date.now(),
+  })
+
+  if (config.retryAfterHeader) {
     event.node.res.setHeader('Retry-After', config.limit.ban)
   }