feat(cu): application level rate limiting for the cu

Author: VinceJuliano

servers/cu/src/config.js

@@ -47,7 +47,13 @@ const serverConfigSchema = domainConfigSchema.extend({
    */
   UNIT_MODE: z.enum(['cu', 'ru']),
   port: positiveIntSchema,
-  ENABLE_METRICS_ENDPOINT: z.preprocess((val) => !!val, z.boolean())
+  ENABLE_METRICS_ENDPOINT: z.preprocess((val) => !!val, z.boolean()),
+  /** 
+   * Rate limiting settings
+   */
+  RATE_LIMIT_WINDOW: positiveIntSchema,
+  RATE_LIMIT_MAX_REQUESTS: positiveIntSchema,
+  RATE_LIMIT_CLEANUP_INTERVAL: positiveIntSchema
 })
 
 /**
@@ -183,7 +189,10 @@ const CONFIG_ENVS = {
     DRY_RUN_DEFAULT_MAX_PROCESS_AGE: process.env.DRY_RUN_DEFAULT_MAX_PROCESS_AGE || 100,
     DRY_RUN_PROCESS_CACHE_TTL: process.env.DRY_RUN_PROCESS_CACHE_TTL || 2000,
     DRY_RUN_RESULT_MAX_AGE: process.env.DRY_RUN_RESULT_MAX_AGE || 60000,
-    LOAD_MESSAGES_PAGE_SIZE: process.env.LOAD_MESSAGES_PAGE_SIZE || 1000
+    LOAD_MESSAGES_PAGE_SIZE: process.env.LOAD_MESSAGES_PAGE_SIZE || 1000,
+    RATE_LIMIT_WINDOW: process.env.RATE_LIMIT_WINDOW,
+    RATE_LIMIT_MAX_REQUESTS: process.env.RATE_LIMIT_MAX_REQUESTS,
+    RATE_LIMIT_CLEANUP_INTERVAL: process.env.RATE_LIMIT_CLEANUP_INTERVAL
   },
   production: {
     MODE,
@@ -246,7 +255,10 @@ const CONFIG_ENVS = {
     DRY_RUN_DEFAULT_MAX_PROCESS_AGE: process.env.DRY_RUN_DEFAULT_MAX_PROCESS_AGE || 100,
     DRY_RUN_PROCESS_CACHE_TTL: process.env.DRY_RUN_PROCESS_CACHE_TTL || 2000,
     DRY_RUN_RESULT_MAX_AGE: process.env.DRY_RUN_RESULT_MAX_AGE || 60000,
-    LOAD_MESSAGES_PAGE_SIZE: process.env.LOAD_MESSAGES_PAGE_SIZE || 1000
+    LOAD_MESSAGES_PAGE_SIZE: process.env.LOAD_MESSAGES_PAGE_SIZE || 1000,
+    RATE_LIMIT_WINDOW: process.env.RATE_LIMIT_WINDOW,
+    RATE_LIMIT_MAX_REQUESTS: process.env.RATE_LIMIT_MAX_REQUESTS,
+    RATE_LIMIT_CLEANUP_INTERVAL: process.env.RATE_LIMIT_CLEANUP_INTERVAL
   }
 }
 

servers/cu/src/domain/api/dryRun.js

@@ -182,7 +182,7 @@ export function dryRunWith (env) {
         'Using recently cached dry-run result for dry-run to process "%s"',
         processId
       )
-      return Resolved(cached.ctx)
+      return Resolved({ result: cached.ctx, wasCached: true })
     }
 
     return of({ processId, messageTxId })
@@ -252,7 +252,7 @@ export function dryRunWith (env) {
         const omitted = omit(['Memory'], res.output)
         const cached = { age: new Date().getTime(), ctx: omitted }
         dryRunResultCache.set(dryRunHash, cached, DRY_RUN_RESULT_MAX_AGE)
-        return omitted
+        return { result: omitted, wasCached: false }
       })
   }
 }

servers/cu/src/routes/dryRun.js

@@ -39,6 +39,11 @@ export const withDryRunRoutes = app => {
           domain: { BUSY_THRESHOLD, apis: { dryRun } }
         } = req
 
+        if(req.limitRequest(req)) {
+          res.status(429).send({'error': 'Rate limit exceeded'})
+          return
+        }
+
         const input = inputSchema.parse({ processId, messageTxId, maxProcessAge, dryRun: body })
 
         await busyIn(
@@ -48,7 +53,12 @@ export const withDryRunRoutes = app => {
             res.status(202)
             return { message: `Evaluation of process "${input.processId}" to "${input.messageTxId || 'latest'}" is in progress.` }
           }
-        ).then((output) => res.send(output))
+        ).then((output) => {
+          res.send(output.result)
+          if(!output.wasCached) {
+            req.recordRequest(req)
+          }
+        })
       })
     )()
   )

servers/cu/src/routes/middleware/index.js

@@ -2,16 +2,22 @@ import { compose } from 'ramda'
 
 import { withErrorHandler } from './withErrorHandler.js'
 import { withDomain } from './withDomain.js'
+import { createRateLimitMiddleware } from './withRateLimits.js'
+import { config } from '../../config.js'
 
 export * from './withProcessRestriction.js'
 export * from './withMetrics.js'
 export * from './withCuMode.js'
 
+// Initialize rate limiting with server config
+const withRateLimits = createRateLimitMiddleware(config)
+
 /**
  * A convenience method that composes common middleware needed on most routes,
  * such that other routes can simply compose this one middleware.
  */
 export const withMiddleware = compose(
   withDomain,
-  withErrorHandler
+  withErrorHandler,
+  withRateLimits
 )

servers/cu/src/routes/middleware/withRateLimits.js

@@ -0,0 +1,94 @@
+/**
+ * Extract the IP address from the request
+ */
+const getClientIp = (req) => {
+  const forwarded = req.headers['x-forwarded-for']
+  if (forwarded) {
+    return forwarded.split(',')[0].trim()
+  }
+
+  const realIp = req.headers['x-real-ip']
+  if (realIp) {
+    return realIp
+  }
+
+  return req.socket.remoteAddress || req.connection.remoteAddress
+}
+
+/**
+ * Factory function that creates rate limiting middleware with config
+ */
+export const createRateLimitMiddleware = (config) => {
+  const RATE_LIMIT_WINDOW = config.RATE_LIMIT_WINDOW
+  const RATE_LIMIT_MAX_REQUESTS = config.RATE_LIMIT_MAX_REQUESTS
+  const RATE_LIMIT_CLEANUP_INTERVAL = config.RATE_LIMIT_CLEANUP_INTERVAL
+
+  if(!RATE_LIMIT_WINDOW || !RATE_LIMIT_MAX_REQUESTS || !RATE_LIMIT_CLEANUP_INTERVAL) {
+    return (handler) => (req, res) => {
+        req.recordRequest = () => {}
+        req.limitRequest = () => {}
+        return handler(req, res)
+    }
+  }
+
+  const requestStore = new Map()
+
+  /*
+   * Clean up IPs that havent made a request
+   */
+  setInterval(() => {
+    const now = Date.now()
+    const cutoff = now - RATE_LIMIT_WINDOW
+
+    for (const [ip, timestamps] of requestStore.entries()) {
+      const validTimestamps = timestamps.filter(ts => ts > cutoff)
+      if (validTimestamps.length === 0) {
+        requestStore.delete(ip)
+      } else {
+        requestStore.set(ip, validTimestamps)
+      }
+    }
+  }, RATE_LIMIT_CLEANUP_INTERVAL)
+
+  /*
+   * Log a request in the limit table
+   */
+  const recordRequest = (req) => {
+    const ip = getClientIp(req)
+    const now = Date.now()
+
+    const timestamps = requestStore.get(ip) || []
+    timestamps.push(now)
+    requestStore.set(ip, timestamps)
+  }
+
+  /*
+   * Determine if a request should be limited
+   */
+  const limitRequest = (req) => {
+    const ip = getClientIp(req)
+    const now = Date.now()
+    const cutoff = now - RATE_LIMIT_WINDOW
+
+    const timestamps = requestStore.get(ip) || []
+
+    const validTimestamps = timestamps.filter(ts => ts > cutoff)
+
+    if (validTimestamps.length > 0) {
+      requestStore.set(ip, validTimestamps)
+    } else {
+      requestStore.delete(ip)
+    }
+
+    return validTimestamps.length >= RATE_LIMIT_MAX_REQUESTS
+  }
+
+  /**
+   * A middleware that exposes rate limiting functionality
+   */
+  return (handler) => (req, res) => {
+    req.recordRequest = recordRequest
+    req.limitRequest = limitRequest
+    return handler(req, res)
+  }
+}