@creator.co/wapi 1.8.4 → 1.8.5

package/tests/Server/lib/container/RateLimit.test.ts

@@ -0,0 +1,772 @@
+import { jest } from '@jest/globals'
+import { APIGatewayProxyEvent, Context } from 'aws-lambda'
+import { expect as c_expect } from 'chai'
+import request from 'supertest'
+
+import Globals from '../../../../src/Globals.js'
+import Proxy from '../../../../src/Server/lib/container/Proxy.js'
+import { GlobalRateLimitConfig } from '../../../../src/Server/Router.js'
+import { defaultUrl } from '../../../Test.utils.js'
+
+describe('Rate Limiting', () => {
+  // @ts-ignore
+  let mockExit = jest.spyOn(process, 'exit').mockImplementation(() => {})
+  let proxy: Proxy | null = null
+
+  beforeAll(() => {
+    // @ts-ignore
+    mockExit = jest.spyOn(process, 'exit').mockImplementation(() => {})
+  })
+
+  afterAll(() => {
+    mockExit.mockRestore()
+  })
+
+  beforeEach(() => {
+    mockExit.mockReset()
+  })
+
+  afterEach(async () => {
+    if (proxy) {
+      await proxy.unload()
+      proxy = null
+    }
+    // Give ports time to release
+    await new Promise(resolve => setTimeout(resolve, 100))
+  })
+
+  test('Rate limit is applied with default settings', async () => {
+    const port = 57001
+    const url = defaultUrl.replace(String(Globals.Listener_HTTP_DefaultPort), String(port))
+
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: true,
+      windowMs: 1000, // 1 second window for testing
+      limit: 3, // Allow only 3 requests
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    // First 3 requests should succeed
+    for (let i = 0; i < 3; i++) {
+      const res = await request(url).get('/test').expect(200)
+      c_expect(res.body).to.deep.equal({ success: true })
+    }
+
+    // 4th request should be rate limited
+    const rateLimitedRes = await request(url).get('/test').expect(429)
+    c_expect(rateLimitedRes.body).to.have.property('error', 'rate_limit_exceeded')
+    c_expect(rateLimitedRes.body).to.have.property('message')
+
+    // Wait for window to reset
+    await new Promise(resolve => setTimeout(resolve, 1100))
+
+    // Should work again after window reset
+    const afterResetRes = await request(url).get('/test').expect(200)
+    c_expect(afterResetRes.body).to.deep.equal({ success: true })
+  }, 10000)
+
+  test('Rate limit headers are present', async () => {
+    const port = 57002
+    const url = defaultUrl.replace(String(Globals.Listener_HTTP_DefaultPort), String(port))
+
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: true,
+      windowMs: 60000,
+      limit: 10,
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    const res = await request(url).get('/test').expect(200)
+
+    // Check for rate limit headers
+    c_expect(res.headers).to.have.property('ratelimit-limit')
+    c_expect(res.headers).to.have.property('ratelimit-remaining')
+    c_expect(res.headers).to.have.property('ratelimit-reset')
+
+    // Verify header values
+    c_expect(res.headers['ratelimit-limit']).to.equal('10')
+    c_expect(parseInt(res.headers['ratelimit-remaining'])).to.be.at.most(9)
+  })
+
+  test('Rate limit can be disabled', async () => {
+    const port = 57003
+    const url = defaultUrl.replace(String(Globals.Listener_HTTP_DefaultPort), String(port))
+
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: false,
+      windowMs: 100,
+      limit: 2,
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    // Should be able to make many requests without hitting limit
+    for (let i = 0; i < 5; i++) {
+      const res = await request(url).get('/test').expect(200)
+      c_expect(res.body).to.deep.equal({ success: true })
+    }
+  })
+
+  test('Rate limit with custom handler', async () => {
+    const port = 57004
+    const url = defaultUrl.replace(String(Globals.Listener_HTTP_DefaultPort), String(port))
+    let customHandlerCalled = false
+
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: true,
+      windowMs: 1000,
+      limit: 2,
+      handler: (req, res) => {
+        customHandlerCalled = true
+        res.status(429).json({
+          custom: true,
+          message: 'Custom rate limit message',
+        })
+      },
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    // First 2 requests succeed
+    await request(url).get('/test').expect(200)
+    await request(url).get('/test').expect(200)
+
+    // 3rd request hits custom handler
+    const rateLimitedRes = await request(url).get('/test').expect(429)
+    c_expect(customHandlerCalled).to.be.true
+    c_expect(rateLimitedRes.body).to.have.property('custom', true)
+    c_expect(rateLimitedRes.body).to.have.property('message', 'Custom rate limit message')
+  }, 10000)
+
+  test('Rate limit with skip function', async () => {
+    const port = 57005
+    const url = defaultUrl.replace(String(Globals.Listener_HTTP_DefaultPort), String(port))
+
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: true,
+      windowMs: 1000,
+      limit: 2,
+      skip: req => {
+        // Skip rate limiting for requests with special header
+        return req.headers['x-skip-rate-limit'] === 'true'
+      },
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    // Use up the rate limit
+    await request(url).get('/test').expect(200)
+    await request(url).get('/test').expect(200)
+
+    // Normal request should be rate limited
+    await request(url).get('/test').expect(429)
+
+    // Request with skip header should succeed
+    const skippedRes = await request(url).get('/test').set('x-skip-rate-limit', 'true').expect(200)
+    c_expect(skippedRes.body).to.deep.equal({ success: true })
+  }, 10000)
+
+  test('Rate limit with custom key generator', async () => {
+    const port = 57006
+    const url = defaultUrl.replace(String(Globals.Listener_HTTP_DefaultPort), String(port))
+
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: true,
+      windowMs: 1000,
+      limit: 2,
+      keyGenerator: req => {
+        // Use custom header for rate limiting key instead of IP
+        return (req.headers['x-user-id'] as string) || 'anonymous'
+      },
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    // User 1 makes 2 requests (uses up their limit)
+    await request(url).get('/test').set('x-user-id', 'user1').expect(200)
+    await request(url).get('/test').set('x-user-id', 'user1').expect(200)
+
+    // User 1's 3rd request should be rate limited
+    await request(url).get('/test').set('x-user-id', 'user1').expect(429)
+
+    // User 2 should still have their own limit
+    const user2Res = await request(url).get('/test').set('x-user-id', 'user2').expect(200)
+    c_expect(user2Res.body).to.deep.equal({ success: true })
+  }, 10000)
+
+  test('Rate limit without config does not apply limiting', async () => {
+    const port = 57007
+    const url = defaultUrl.replace(String(Globals.Listener_HTTP_DefaultPort), String(port))
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        // No rateLimit config
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    // Should be able to make many requests
+    for (let i = 0; i < 10; i++) {
+      const res = await request(url).get('/test').expect(200)
+      c_expect(res.body).to.deep.equal({ success: true })
+    }
+  })
+
+  test('Multiple requests from different IPs have separate limits', async () => {
+    const port = 57008
+    const url = defaultUrl.replace(String(Globals.Listener_HTTP_DefaultPort), String(port))
+
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: true,
+      windowMs: 1000,
+      limit: 1, // Very strict: only 1 request per window
+      keyGenerator: req => {
+        // Use x-forwarded-for header to simulate different IPs
+        return (req.headers['x-forwarded-for'] as string) || req.ip || 'unknown'
+      },
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    // IP1 makes a request (uses their limit)
+    await request(url).get('/test').set('x-forwarded-for', '192.168.1.1').expect(200)
+
+    // IP1's second request should be rate limited
+    await request(url).get('/test').set('x-forwarded-for', '192.168.1.1').expect(429)
+
+    // IP2 should still have their own limit available
+    const ip2Res = await request(url).get('/test').set('x-forwarded-for', '192.168.1.2').expect(200)
+    c_expect(ip2Res.body).to.deep.equal({ success: true })
+
+    // IP2's second request should also be rate limited
+    await request(url).get('/test').set('x-forwarded-for', '192.168.1.2').expect(429)
+  }, 10000)
+
+  test('Rate limit with containerSetupHook', async () => {
+    const port = 57009
+    const url = defaultUrl.replace(String(Globals.Listener_HTTP_DefaultPort), String(port))
+    let hookCalled = false
+
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: true,
+      windowMs: 1000,
+      limit: 5,
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+        containerSetupHook: async (listener, app) => {
+          hookCalled = true
+          // Hook can be used to add custom middleware or configure the server
+          c_expect(listener).to.exist
+          c_expect(app).to.exist
+        },
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    // Verify hook was called
+    c_expect(hookCalled).to.be.true
+
+    // Verify rate limiting still works after hook
+    const res = await request(url).get('/test').expect(200)
+    c_expect(res.body).to.deep.equal({ success: true })
+
+    // Check rate limit headers are present
+    c_expect(res.headers).to.have.property('ratelimit-limit')
+  })
+
+  test('Rate limit with Redis store configuration', async () => {
+    const port = 57010
+
+    // Capture console logs to verify Redis store is configured
+    const originalLog = console.log
+    const logs: string[] = []
+    console.log = (...args: any[]) => {
+      logs.push(args.join(' '))
+      originalLog.apply(console, args)
+    }
+
+    // Mock Redis client with proper typing
+    const mockRedisClient: any = {
+      sendCommand: jest.fn(async () => 'OK'),
+      connect: jest.fn(async () => undefined),
+      disconnect: jest.fn(async () => undefined),
+      isOpen: true,
+      isReady: true,
+    }
+
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: true,
+      windowMs: 1000,
+      limit: 3,
+      store: 'redis',
+      redis: {
+        client: mockRedisClient as any,
+        prefix: 'test:rl:',
+      },
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    // Restore console.log
+    console.log = originalLog
+
+    // Verify Redis store logging message
+    const redisStoreLog = logs.find(log => log.includes('[RATE-LIMIT] - Using Redis store'))
+    c_expect(redisStoreLog).to.exist
+
+    // Verify rate limiting is enabled
+    const rateLimitEnabledLog = logs.find(log =>
+      log.includes('[RATE-LIMIT] - Global rate limiting enabled')
+    )
+    c_expect(rateLimitEnabledLog).to.exist
+
+    // Note: We don't make actual requests in this test because the mock Redis client
+    // doesn't fully implement the RedisStore interface. This test verifies that:
+    // 1. Redis store configuration is recognized
+    // 2. Correct logging messages are emitted
+    // 3. The proxy loads successfully with Redis config
+  })
+
+  test('Rate limit logs indicate store type', async () => {
+    const port = 57011
+    const url = defaultUrl.replace(String(Globals.Listener_HTTP_DefaultPort), String(port))
+
+    // Capture console logs
+    const originalLog = console.log
+    const logs: string[] = []
+    console.log = (...args: any[]) => {
+      logs.push(args.join(' '))
+      originalLog.apply(console, args)
+    }
+
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: true,
+      windowMs: 1000,
+      limit: 5,
+      store: 'memory', // Explicitly set to memory
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    // Restore console.log
+    console.log = originalLog
+
+    // Verify logging messages
+    const rateLimitEnabledLog = logs.find(log =>
+      log.includes('[RATE-LIMIT] - Global rate limiting enabled')
+    )
+    const memoryStoreLog = logs.find(log => log.includes('[RATE-LIMIT] - Using in-memory store'))
+
+    c_expect(rateLimitEnabledLog).to.exist
+    c_expect(memoryStoreLog).to.exist
+
+    // Make a request to verify it works
+    const res = await request(url).get('/test').expect(200)
+    c_expect(res.body).to.deep.equal({ success: true })
+  })
+
+  test('Rate limit uses default keyGenerator with IP fallback chain', async () => {
+    const port = 57012
+    const url = defaultUrl.replace(String(Globals.Listener_HTTP_DefaultPort), String(port))
+
+    // Don't provide a custom keyGenerator - should use default
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: true,
+      windowMs: 1000,
+      limit: 2,
+      // NO keyGenerator provided - will use default IP-based logic
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    // Make requests without x-forwarded-for header
+    // Should use req.ip or req.socket.remoteAddress
+    const res1 = await request(url).get('/test').expect(200)
+    c_expect(res1.body).to.deep.equal({ success: true })
+
+    const res2 = await request(url).get('/test').expect(200)
+    c_expect(res2.body).to.deep.equal({ success: true })
+
+    // 3rd request should be rate limited (covers default handler lines 208-209)
+    const res3 = await request(url).get('/test').expect(429)
+    c_expect(res3.body).to.have.property('error', 'rate_limit_exceeded')
+    c_expect(res3.body).to.have.property('message', 'Too many requests. Please try again later.')
+  }, 10000)
+
+  test('Rate limit default keyGenerator uses x-forwarded-for', async () => {
+    const port = 57013
+    const url = defaultUrl.replace(String(Globals.Listener_HTTP_DefaultPort), String(port))
+
+    // Don't provide a custom keyGenerator
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: true,
+      windowMs: 1000,
+      limit: 1, // Very strict
+      // NO keyGenerator - uses default with x-forwarded-for fallback
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    // Test with x-forwarded-for header (covers line 197-198)
+    const res1 = await request(url)
+      .get('/test')
+      .set('x-forwarded-for', '10.0.0.1, 10.0.0.2')
+      .expect(200)
+    c_expect(res1.body).to.deep.equal({ success: true })
+
+    // Second request from same IP should be rate limited
+    const res2 = await request(url)
+      .get('/test')
+      .set('x-forwarded-for', '10.0.0.1, 10.0.0.3')
+      .expect(429)
+    c_expect(res2.body).to.have.property('error', 'rate_limit_exceeded')
+
+    // Wait for window to reset
+    await new Promise(resolve => setTimeout(resolve, 1100))
+
+    // Different IP should work (covers the split and trim logic)
+    const res3 = await request(url).get('/test').set('x-forwarded-for', '10.0.0.99').expect(200)
+    c_expect(res3.body).to.deep.equal({ success: true })
+  }, 10000)
+
+  test('Rate limit default handler with warning log', async () => {
+    const port = 57014
+    const url = defaultUrl.replace(String(Globals.Listener_HTTP_DefaultPort), String(port))
+
+    // Capture console.warn to verify default handler logging
+    const originalWarn = console.warn
+    const warnings: any[] = []
+    console.warn = (...args: any[]) => {
+      warnings.push(args)
+      originalWarn.apply(console, args)
+    }
+
+    // Don't provide custom handler - should use default
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: true,
+      windowMs: 1000,
+      limit: 1,
+      // NO custom handler - will use default (lines 208-209)
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    // Use up the limit
+    await request(url).get('/test-endpoint').expect(200)
+
+    // Trigger rate limit (should call default handler and log warning)
+    const rateLimitedRes = await request(url).get('/test-endpoint').expect(429)
+
+    // Restore console.warn
+    console.warn = originalWarn
+
+    // Verify default handler was called (lines 208-209, 220-225)
+    c_expect(rateLimitedRes.body).to.deep.equal({
+      error: 'rate_limit_exceeded',
+      message: 'Too many requests. Please try again later.',
+    })
+
+    // Verify warning was logged
+    const rateLimitWarning = warnings.find(w => w[0] === '[Proxy] - [RATE-LIMIT] - Limit exceeded')
+    c_expect(rateLimitWarning).to.exist
+    c_expect(rateLimitWarning[1]).to.have.property('path', '/test-endpoint')
+    c_expect(rateLimitWarning[1]).to.have.property('method', 'GET')
+    c_expect(rateLimitWarning[1]).to.have.property('timestamp')
+  }, 10000)
+
+  test('Rate limit with Redis store sendCommand and prefix', async () => {
+    const port = 57015
+
+    // Capture console logs
+    const originalLog = console.log
+    const logs: string[] = []
+    console.log = (...args: any[]) => {
+      logs.push(args.join(' '))
+      originalLog.apply(console, args)
+    }
+
+    // Mock Redis client
+    const sendCommandCalls: any[] = []
+    const mockRedisClient: any = {
+      sendCommand: jest.fn(async (...args: any[]) => {
+        sendCommandCalls.push(args)
+        return 'OK'
+      }),
+    }
+
+    // Test with custom prefix (covers lines 248-252)
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: true,
+      windowMs: 1000,
+      limit: 3,
+      store: 'redis',
+      redis: {
+        client: mockRedisClient,
+        prefix: 'custom:prefix:', // Custom prefix (line 252)
+      },
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    await proxy.load()
+
+    // Restore console.log
+    console.log = originalLog
+
+    // Verify Redis store was created with correct logging (line 249)
+    const redisLog = logs.find(log => log.includes('[RATE-LIMIT] - Using Redis store'))
+    c_expect(redisLog).to.exist
+
+    // Verify RedisStore was instantiated (lines 250-252)
+    // The store object should have been created with sendCommand and prefix
+    // We can't directly test the RedisStore internals, but we verified:
+    // 1. Redis store logging message appears
+    // 2. The configuration is accepted without errors
+    // 3. Proxy loads successfully with Redis config
+  })
+
+  test('Rate limit with Redis store default prefix', async () => {
+    const port = 57016
+
+    const mockRedisClient: any = {
+      sendCommand: jest.fn(async () => 'OK'),
+    }
+
+    // Don't provide prefix - should use default 'wapi:rl:' (line 252)
+    const rateLimitConfig: GlobalRateLimitConfig = {
+      enabled: true,
+      windowMs: 1000,
+      limit: 5,
+      store: 'redis',
+      redis: {
+        client: mockRedisClient,
+        // NO prefix - should default to 'wapi:rl:' (line 252)
+      },
+    }
+
+    proxy = new Proxy(
+      {
+        routes: [],
+        port,
+        rateLimit: rateLimitConfig,
+      },
+      async (event: APIGatewayProxyEvent, context: Context) => {
+        context.succeed({
+          body: JSON.stringify({ success: true }),
+          statusCode: 200,
+        })
+      }
+    )
+
+    // Should load successfully with default prefix
+    await proxy.load()
+
+    // Verify proxy is running
+    c_expect(proxy).to.exist
+  })
+})
+
+export {}