Merge pull request #614 from TaloDev/fix-rate-limiting

Use rate-limiter-flexible for rate-limiting

Author: tudddorrr

package-lock.json

@@ -93,6 +93,7 @@
     "otplib": "^12.0.1",
     "qrcode": "^1.5.0",
     "qs": "^6.11.0",
+    "rate-limiter-flexible": "^7.3.0",
     "stripe": "^18.0.0",
     "uuid": "^9.0.0",
     "ws": "^8.18.0",

package.json

@@ -1,46 +1,38 @@
 import { Redis } from 'ioredis'
+import { RateLimiterRedis } from 'rate-limiter-flexible'
 
-const cache = new Map<string, { count: number, expires: number }>()
+const rateLimiters = new Map<string, RateLimiterRedis>()
 
-setInterval(() => {
-  const now = Date.now()
-  for (const [key, value] of cache.entries()) {
-    if (now > value.expires) {
-      cache.delete(key)
-    }
+function getRateLimiter(redis: Redis, maxRequests: number, duration = 60): RateLimiterRedis {
+  const limiterKey = `${maxRequests}_${duration}`
+  if (!rateLimiters.has(limiterKey)) {
+    rateLimiters.set(limiterKey, new RateLimiterRedis({
+      storeClient: redis,
+      keyPrefix: `rl_${maxRequests}_${duration}`,
+      points: maxRequests,
+      duration: duration,
+      blockDuration: duration
+    }))
   }
-}, 5000)
-
-const script = `
-  local current = redis.call('INCR', KEYS[1])
-  if current == 1 then
-    redis.call('EXPIRE', KEYS[1], ARGV[1])
-  end
-  return current
-`
+  return rateLimiters.get(limiterKey)!
+}
 
 export default async function checkRateLimitExceeded(
   redis: Redis,
   key: string,
-  maxRequests: number
+  maxRequests: number,
+  duration = 60
 ): Promise<boolean> {
-  // Skip cache in test environment for predictable behavior
-  if (process.env.NODE_ENV !== 'test') {
-    const cached = cache.get(key)
-    if (cached && Date.now() < cached.expires) {
-      return cached.count > maxRequests
-    }
-  }
-
-  const current = await redis.eval(script, 1, key, 1) as number
+  const rateLimiter = getRateLimiter(redis, maxRequests, duration)
 
-  // Only cache in production
-  if (process.env.NODE_ENV !== 'test') {
-    cache.set(key, {
-      count: current,
-      expires: Date.now() + 500
-    })
+  try {
+    await rateLimiter.consume(key)
+    return false
+  } catch (err) {
+    if (err && typeof err === 'object' && 'remainingPoints' in err) {
+      return true
+    }
+    // re-throw actual errors
+    throw err
   }
-
-  return current > maxRequests
 }

src/lib/errors/checkRateLimitExceeded.ts

@@ -2,8 +2,15 @@ import { APIKeyScope } from '../../src/entities/api-key'
 import createSocketIdentifyMessage from '../utils/createSocketIdentifyMessage'
 import GameChannelFactory from '../fixtures/GameChannelFactory'
 import createTestSocket from '../utils/createTestSocket'
+import * as checkRateLimitExceeded from '../../src/lib/errors/checkRateLimitExceeded'
 
 describe('Socket rate limiting', () => {
+  const checkRateLimitExceededMock = vi.spyOn(checkRateLimitExceeded, 'default')
+
+  afterEach(() => {
+    checkRateLimitExceededMock.mockClear()
+  })
+
   it('should return a rate limiting error', async () => {
     const { identifyMessage, ticket, player } = await createSocketIdentifyMessage([
       APIKeyScope.READ_PLAYERS,
@@ -14,11 +21,10 @@ describe('Socket rate limiting', () => {
     channel.members.add(player.aliases[0])
     await em.persistAndFlush(channel)
 
-    await createTestSocket(`/?ticket=${ticket}`, async (client, socket) => {
+    await createTestSocket(`/?ticket=${ticket}`, async (client) => {
       await client.identify(identifyMessage)
 
-      const conn = socket.findConnections((conn) => conn.playerAliasId === player.aliases[0].id)[0]
-      await redis.set(conn.rateLimitKey, 999)
+      checkRateLimitExceededMock.mockResolvedValueOnce(true)
 
       client.sendJson({
         req: 'v1.channels.message',
@@ -56,7 +62,7 @@ describe('Socket rate limiting', () => {
       const conn = socket.findConnections((conn) => conn.playerAliasId === player.aliases[0].id)[0]
       conn.rateLimitWarnings = 3
 
-      await redis.set(conn.rateLimitKey, 999)
+      checkRateLimitExceededMock.mockResolvedValueOnce(true)
 
       client.sendJson({
         req: 'v1.channels.message',