feat: Mcp rate Limiter (#16)

Co-authored-by: John Leider <[email protected]>

Author: Haviles04

src/index.ts

@@ -11,11 +11,11 @@ import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'
 
 import { intro } from './cli/intro.js'
 import packageJson from '../package.json' with { type: 'json' }
+import { startHttpServer } from './transports/http.js'
 
 import { registerPrompts } from '#prompts/index'
 import { registerResources } from '#resources/index'
 import { registerTools } from '#tools/index'
-import { startHttpServer } from '#transports/http'
 
 function parseArgs () {
   const args = process.argv.slice(2)

src/transports/http.ts

@@ -13,11 +13,14 @@ import { registerPrompts } from '#prompts/index'
 import { registerResources } from '#resources/index'
 import { registerTools } from '#tools/index'
 import packageJson from '../../package.json' with { type: 'json' }
+import { RateLimiter } from '../utils/rate-limiter.js'
+import type { RateLimiterOptions } from '../utils/rate-limiter.js'
 
 export interface HttpServerOptions {
   port?: number
   host?: string
   path?: string
+  rateLimit?: RateLimiterOptions
 }
 
 async function createMcpServer () {
@@ -49,10 +52,12 @@ export async function startHttpServer (options: HttpServerOptions = {}): Promise
   const host = options.host ?? 'localhost'
   const path = options.path ?? '/mcp'
 
+  const rateLimiter = options.rateLimit ? new RateLimiter(options.rateLimit) : null
+
   return new Promise((resolve, reject) => {
     const httpServer = createServer(async (req, res) => {
       try {
-        await handleRequest(req, res, path)
+        await handleRequest(req, res, path, rateLimiter, options.rateLimit)
       } catch (error) {
         console.error('Error handling request:', error)
         if (!res.headersSent) {
@@ -71,10 +76,27 @@ export async function startHttpServer (options: HttpServerOptions = {}): Promise
   })
 }
 
+function getClientIdentifier (req: IncomingMessage): string {
+  // Use session ID if available (for stateful mode)
+  const sessionId = req.headers['mcp-session-id']
+  if (sessionId && typeof sessionId === 'string') {
+    return `session:${sessionId}`
+  }
+
+  // Fall back to IP address
+  const forwarded = req.headers['x-forwarded-for']
+  const ip = forwarded
+    ? (typeof forwarded === 'string' ? forwarded.split(',')[0] : forwarded[0])
+    : req.socket.remoteAddress
+  return `ip:${ip ?? 'unknown'}`
+}
+
 async function handleRequest (
   req: IncomingMessage,
   res: ServerResponse,
   mcpPath: string,
+  rateLimiter: RateLimiter | null,
+  rateLimitOptions?: RateLimiterOptions,
 ): Promise<void> {
   console.error(`[${new Date().toISOString()}] ${req.method} ${req.url}`)
 
@@ -89,6 +111,31 @@ async function handleRequest (
     return
   }
 
+  if (rateLimiter && req.url !== '/health' && req.url !== '/') {
+    const clientId = getClientIdentifier(req)
+    const rateLimitResult = rateLimiter.check(clientId)
+
+    if (rateLimitOptions) {
+      res.setHeader('X-RateLimit-Limit', rateLimitOptions.maxRequests.toString())
+      res.setHeader('X-RateLimit-Remaining', rateLimitResult.remaining.toString())
+      res.setHeader('X-RateLimit-Reset', new Date(rateLimitResult.resetTime).toISOString())
+    }
+
+    if (!rateLimitResult.allowed) {
+      res.writeHead(429, {
+        'Content-Type': 'application/json',
+        'Retry-After': rateLimitResult.retryAfter?.toString() ?? '60',
+      })
+      res.end(JSON.stringify({
+        error: 'Too Many Requests',
+        message: 'Rate limit exceeded. Please try again later.',
+        retryAfter: rateLimitResult.retryAfter,
+        resetTime: new Date(rateLimitResult.resetTime).toISOString(),
+      }))
+      return
+    }
+  }
+
   // Health check endpoint
   if (req.url === '/health' && req.method === 'GET') {
     res.writeHead(200, { 'Content-Type': 'application/json' })

src/utils/rate-limiter.ts

@@ -0,0 +1,129 @@
+/**
+ * Rate Limiter Implementation
+ *
+ * Implements a sliding window rate limiter to control request rates
+ */
+
+export interface RateLimiterOptions {
+  /** Maximum number of requests allowed in the time window */
+  maxRequests: number
+  /** Time window in milliseconds */
+  windowMs: number
+  /** Optional key generator function for identifying clients */
+  keyGenerator?: (identifier: string) => string
+}
+
+interface RateLimitRecord {
+  timestamps: number[]
+}
+
+export class RateLimiter {
+  private records: Map<string, RateLimitRecord> = new Map()
+  private options: Required<RateLimiterOptions>
+  private cleanupInterval: NodeJS.Timeout | null = null
+
+  constructor (options: RateLimiterOptions) {
+    this.options = {
+      maxRequests: options.maxRequests,
+      windowMs: options.windowMs,
+      keyGenerator: options.keyGenerator ?? ((id: string) => id),
+    }
+
+    // Clean up old records every minute to prevent memory leaks
+    this.cleanupInterval = setInterval(() => {
+      this.cleanup()
+    }, 60_000)
+  }
+
+  /**
+   * Check if a request should be allowed
+   * @param identifier - Client identifier (e.g., IP address, session ID)
+   * @returns Object with allowed status and retry information
+   */
+  check (identifier: string): {
+    allowed: boolean
+    remaining: number
+    resetTime: number
+    retryAfter?: number
+  } {
+    const key = this.options.keyGenerator(identifier)
+    const now = Date.now()
+    const windowStart = now - this.options.windowMs
+
+    let record = this.records.get(key)
+    if (!record) {
+      record = { timestamps: [] }
+      this.records.set(key, record)
+    }
+
+    // Remove timestamps outside the current window
+    record.timestamps = record.timestamps.filter(ts => ts > windowStart)
+
+    // Check if limit is exceeded
+    const allowed = record.timestamps.length < this.options.maxRequests
+    const remaining = Math.max(0, this.options.maxRequests - record.timestamps.length)
+
+    // Calculate reset time (when the oldest request will fall outside the window)
+    const resetTime = record.timestamps.length > 0
+      ? record.timestamps[0] + this.options.windowMs
+      : now + this.options.windowMs
+
+    let retryAfter: number | undefined
+    if (!allowed && record.timestamps.length > 0) {
+      // Calculate how long to wait until the oldest request expires
+      retryAfter = Math.ceil((record.timestamps[0] + this.options.windowMs - now) / 1000)
+    }
+
+    // If allowed, record this request
+    if (allowed) {
+      record.timestamps.push(now)
+    }
+
+    return {
+      allowed,
+      remaining: allowed ? remaining - 1 : remaining,
+      resetTime,
+      retryAfter,
+    }
+  }
+
+  /**
+   * Reset rate limit for a specific identifier
+   */
+  reset (identifier: string): void {
+    const key = this.options.keyGenerator(identifier)
+    this.records.delete(key)
+  }
+
+  /**
+   * Clear all rate limit records
+   */
+  clear (): void {
+    this.records.clear()
+  }
+
+  destroy (): void {
+    if (this.cleanupInterval) {
+      clearInterval(this.cleanupInterval)
+      this.cleanupInterval = null
+    }
+  }
+
+  /**
+   * Clean up expired records
+   */
+  private cleanup (): void {
+    const now = Date.now()
+    const windowStart = now - this.options.windowMs
+
+    for (const [key, record] of this.records.entries()) {
+      // Remove timestamps outside the current window
+      record.timestamps = record.timestamps.filter(ts => ts > windowStart)
+
+      // Remove the record entirely if it has no timestamps
+      if (record.timestamps.length === 0) {
+        this.records.delete(key)
+      }
+    }
+  }
+}