feat: add helper function to include rate limit headers and enhance tests for key generation errors

Author: jkyberneees

lib/middleware/rate-limit.js

@@ -74,6 +74,29 @@ function createRateLimit(options = {}) {
     skip,
   } = options
 
+  /**
+   * Helper function to add rate limit headers to a response
+   * @param {Response} response - Response object to add headers to
+   * @param {number} totalHits - Current hit count
+   * @param {Date} resetTime - When the rate limit resets
+   * @returns {Response} Response with headers added
+   */
+  const addRateLimitHeaders = (response, totalHits, resetTime) => {
+    if (standardHeaders && response && response.headers) {
+      response.headers.set('X-RateLimit-Limit', max.toString())
+      response.headers.set(
+        'X-RateLimit-Remaining',
+        Math.max(0, max - totalHits).toString(),
+      )
+      response.headers.set(
+        'X-RateLimit-Reset',
+        Math.ceil(resetTime.getTime() / 1000).toString(),
+      )
+      response.headers.set('X-RateLimit-Used', totalHits.toString())
+    }
+    return response
+  }
+
   return async function rateLimitMiddleware(req, next) {
     // Allow test to inject a fresh store for isolation
     const activeStore = req && req.rateLimitStore ? req.rateLimitStore : store
@@ -92,22 +115,6 @@ function createRateLimit(options = {}) {
       const key = await keyGenerator(req)
       const {totalHits, resetTime} = await activeStore.increment(key, windowMs)
 
-      const createResponse = (response) => {
-        if (standardHeaders && response && response.headers) {
-          response.headers.set('X-RateLimit-Limit', max.toString())
-          response.headers.set(
-            'X-RateLimit-Remaining',
-            Math.max(0, max - totalHits).toString(),
-          )
-          response.headers.set(
-            'X-RateLimit-Reset',
-            Math.ceil(resetTime.getTime() / 1000).toString(),
-          )
-          response.headers.set('X-RateLimit-Used', totalHits.toString())
-        }
-        return response
-      }
-
       if (totalHits > max) {
         let response
 
@@ -121,7 +128,7 @@ function createRateLimit(options = {}) {
           }
         }
 
-        return createResponse(response)
+        return addRateLimitHeaders(response, totalHits, resetTime)
       }
 
       // Set rate limit context
@@ -143,7 +150,7 @@ function createRateLimit(options = {}) {
 
       const response = await next()
       if (response instanceof Response) {
-        return createResponse(response)
+        return addRateLimitHeaders(response, totalHits, resetTime)
       }
       return response
     } catch (error) {
@@ -173,7 +180,7 @@ function createRateLimit(options = {}) {
 
         const response = await next()
         if (response instanceof Response) {
-          return response
+          return addRateLimitHeaders(response, totalHits, resetTime)
         }
         return response
       } catch (e) {

test/unit/rate-limit.test.js

@@ -432,6 +432,43 @@ describe('Rate Limit Middleware', () => {
       expect(req.rateLimit).toBeDefined() // Confirms we're in the fallback path
       expect(req.rateLimit.current).toBe(1) // Fallback incremented with 'unknown' key
     })
+
+    it('should include rate-limit headers in fallback path after key generation error', async () => {
+      const faultyKeyGenerator = jest.fn(() => {
+        throw new Error('Key generation error')
+      })
+
+      // Create a Response object to be returned by next()
+      const testResponse = new Response('test content', {
+        status: 200,
+        headers: {'X-Test': 'value'},
+      })
+
+      const nextFunc = jest.fn(() => testResponse)
+
+      const middleware = rateLimit({
+        windowMs: 60000,
+        max: 5,
+        keyGenerator: faultyKeyGenerator,
+        standardHeaders: true, // Ensure headers are enabled
+      })
+
+      const result = await middleware(req, nextFunc)
+
+      // Should return the response with added rate-limit headers
+      expect(result).toBe(testResponse)
+      expect(result.status).toBe(200)
+
+      // Verify rate-limit headers were added to the response
+      expect(result.headers.get('X-RateLimit-Limit')).toBe('5')
+      expect(result.headers.get('X-RateLimit-Remaining')).toBe('4') // 5 - 1 = 4
+      expect(result.headers.get('X-RateLimit-Used')).toBe('1')
+      expect(result.headers.get('X-RateLimit-Reset')).toBeTruthy()
+
+      // Verify we're in the fallback path
+      expect(req.rateLimit).toBeDefined()
+      expect(req.rateLimit.current).toBe(1) // Fallback incremented with 'unknown' key
+    })
   })
 
   describe('MemoryStore', () => {