Add request rate limit to proxy

proxy/README.md

@@ -20,6 +20,7 @@ _See the [full architecture guide](../README.md) first._
 - Proxy removes cookie headers.
 - Proxy set user’s IP in `X-Forwarded-For` header.
 - Proxy has timeout and response size limit.
+- Proxy has rate limit. The rate limiting is implemented using an in-memory map to track the number of requests made from each IP address to each domain and globally.
 
 ## Test Strategy
 

proxy/index.ts

@@ -6,7 +6,7 @@ const PORT = process.env.PORT ?? 5284
 
 let allowsFrom: RegExp[]
 if (process.env.NODE_ENV !== 'production') {
-  allowsFrom = [/^http:\/\/localhost:5173/]
+  allowsFrom = [/^http:\/\/localhost:5284/]
 } else if (process.env.STAGING) {
   allowsFrom = [
     /^https:\/\/dev.slowreader.app/,

proxy/package.json

@@ -8,7 +8,7 @@
   },
   "scripts": {
     "start": "tsx watch index.ts",
-    "test": "FORCE_COLOR=1 pnpm run /^test:/",
+    "test": "pnpm run /^test:/",
     "build": "esbuild index.ts --bundle --platform=node --sourcemap --format=esm --outfile=dist/index.mjs",
     "production": "node --run build && ./scripts/run-image.sh",
     "test:proxy-coverage": "c8 bnt",

proxy/proxy.ts

@@ -1,9 +1,10 @@
-import type { Server } from 'node:http'
+import type { IncomingMessage, Server, ServerResponse } from 'node:http'
 import { createServer } from 'node:http'
 import { isIP } from 'node:net'
+import { setTimeout } from 'node:timers/promises'
 import { styleText } from 'node:util'
 
-class BadRequestError extends Error {
+export class BadRequestError extends Error {
   code: number
 
   constructor(message: string, code = 400) {
@@ -13,16 +14,164 @@ class BadRequestError extends Error {
   }
 }
 
-export function createProxyServer(config: {
+export interface RateLimitInfo {
+  count: number
+  timestamp: number
+}
+
+export interface ProxyConfig {
   allowLocalhost?: boolean
   allowsFrom: RegExp[]
   maxSize: number
   timeout: number
-}): Server {
-  return createServer(async (req, res) => {
-    let sent = false
+}
+
+const RATE_LIMIT = {
+  GLOBAL: {
+    DURATION: 60 * 1000,
+    LIMIT: 5000
+  },
+  PER_DOMAIN: {
+    DURATION: 60 * 1000,
+    LIMIT: 500
+  }
+}
+
+export let rateLimitMap: Map<string, RateLimitInfo> = new Map()
+let requestQueue: Map<string, Promise<void>> = new Map()
+
+export function isRateLimited(
+  key: string,
+  store: Map<string, RateLimitInfo>,
+  limit: { DURATION: number; LIMIT: number }
+): boolean {
+  let now = performance.now()
+  let rateLimitInfo = store.get(key) || { count: 0, timestamp: now }
+
+  if (now - rateLimitInfo.timestamp > limit.DURATION) {
+    rateLimitInfo.count = 0
+    rateLimitInfo.timestamp = now
+  }
+
+  if (rateLimitInfo.count >= limit.LIMIT) {
+    return true
+  }
+
+  rateLimitInfo.count += 1
+  store.set(key, rateLimitInfo)
+
+  return false
+}
+
+export function checkRateLimit(ip: string, domain: string): boolean {
+  return ['domain', 'global'].some(type => {
+    let key = type === 'domain' ? `${ip}:${domain}` : ip
+    let limit = type === 'domain' ? RATE_LIMIT.PER_DOMAIN : RATE_LIMIT.GLOBAL
+    return isRateLimited(key, rateLimitMap, limit)
+  })
+}
+
+export function handleError(e: unknown, res: ServerResponse): void {
+  // Known errors
+  if (e instanceof Error && e.name === 'TimeoutError') {
+    res.writeHead(400, { 'Content-Type': 'text/plain' })
+    res.end('Timeout')
+  } else if (e instanceof BadRequestError) {
+    res.writeHead(e.code, { 'Content-Type': 'text/plain' })
+    res.end(e.message)
+  } else {
+    // Unknown or Internal errors
+    /* c8 ignore next 9 */
+    if (e instanceof Error) {
+      process.stderr.write(styleText('red', e.stack ?? e.message) + '\n')
+    } else if (typeof e === 'string') {
+      process.stderr.write(styleText('red', e) + '\n')
+    }
+    res.writeHead(500, { 'Content-Type': 'text/plain' })
+    res.end('Internal Server Error')
+  }
+}
+
+export let processRequest = async (
+  req: IncomingMessage,
+  res: ServerResponse,
+  config: ProxyConfig,
+  url: string
+): Promise<void> => {
+  try {
+    // Remove all cookie headers so they will not be set on proxy domain
+    delete req.headers.cookie
+    delete req.headers['set-cookie']
+    delete req.headers.host
+
+    let targetResponse = await fetch(url, {
+      headers: {
+        ...(req.headers as HeadersInit),
+        'host': new URL(url).host,
+        'X-Forwarded-For': req.socket.remoteAddress!
+      },
+      method: req.method,
+      signal: AbortSignal.timeout(config.timeout)
+    })
+
+    let length = targetResponse.headers.has('content-length')
+      ? parseInt(targetResponse.headers.get('content-length')!)
+      : undefined
 
+    if (length && length > config.maxSize) {
+      throw new BadRequestError('Response too large', 413)
+    }
+
+    res.writeHead(targetResponse.status, {
+      'Access-Control-Allow-Headers': '*',
+      'Access-Control-Allow-Methods': 'OPTIONS, POST, GET, PUT, DELETE',
+      'Access-Control-Allow-Origin': req.headers.origin,
+      'Content-Type': targetResponse.headers.get('content-type') ?? 'text/plain'
+    })
+
+    if (targetResponse.body) {
+      let reader = targetResponse.body.getReader()
+      let size = 0
+      let chunk: ReadableStreamReadResult<Uint8Array>
+      do {
+        chunk = await reader.read()
+        if (chunk.value) {
+          res.write(chunk.value)
+          size += chunk.value.length
+          if (size > config.maxSize) {
+            break
+          }
+        }
+      } while (!chunk.done)
+    }
+    res.end()
+  } catch (e) {
+    handleError(e, res)
+  }
+}
+
+export let handleRequestWithDelay = async (
+  req: IncomingMessage,
+  res: ServerResponse,
+  config: ProxyConfig,
+  ip: string,
+  url: string,
+  parsedUrl: URL
+): Promise<void> => {
+  if (checkRateLimit(ip, parsedUrl.hostname)) {
+    let existingQueue = requestQueue.get(ip) || Promise.resolve()
+    let delayedRequest = existingQueue.then(() => setTimeout(1000))
+    requestQueue.set(ip, delayedRequest)
+    await delayedRequest
+  }
+
+  await processRequest(req, res, config, url)
+}
+
+export function createProxyServer(config: ProxyConfig): Server {
+  return createServer(async (req: IncomingMessage, res: ServerResponse) => {
     try {
+      let ip = req.socket.remoteAddress!
       let url = decodeURIComponent((req.url ?? '').slice(1))
 
       let parsedUrl: URL
@@ -32,6 +181,8 @@ export function createProxyServer(config: {
         throw new BadRequestError('Invalid URL')
       }
 
+      req.headers.origin = 'http://localhost:5284/' // debug
+
       // Only HTTP or HTTPS protocols are allowed
       if (!url.startsWith('http://') && !url.startsWith('https://')) {
         throw new BadRequestError('Only HTTP or HTTPS are supported')
@@ -42,7 +193,7 @@ export function createProxyServer(config: {
         throw new BadRequestError('Only GET is allowed', 405)
       }
 
-      // We only allow request from our app
+      // We only allow requests from our app
       if (
         !req.headers.origin ||
         !config.allowsFrom.some(allowed => allowed.test(req.headers.origin!))
@@ -57,77 +208,9 @@ export function createProxyServer(config: {
         throw new BadRequestError('IP addresses are not allowed')
       }
 
-      // Remove all cookie headers so they will not be set on proxy domain
-      delete req.headers.cookie
-      delete req.headers['set-cookie']
-      delete req.headers.host
-
-      let targetResponse = await fetch(url, {
-        headers: {
-          ...(req.headers as HeadersInit),
-          'host': new URL(url).host,
-          'X-Forwarded-For': req.socket.remoteAddress!
-        },
-        method: req.method,
-        signal: AbortSignal.timeout(config.timeout)
-      })
-
-      let length: number | undefined
-      if (targetResponse.headers.has('content-length')) {
-        length = parseInt(targetResponse.headers.get('content-length')!)
-      }
-      if (length && length > config.maxSize) {
-        throw new BadRequestError('Response too large', 413)
-      }
-
-      res.writeHead(targetResponse.status, {
-        'Access-Control-Allow-Headers': '*',
-        'Access-Control-Allow-Methods': 'OPTIONS, POST, GET, PUT, DELETE',
-        'Access-Control-Allow-Origin': req.headers.origin,
-        'Content-Type':
-          targetResponse.headers.get('content-type') ?? 'text/plain'
-      })
-      sent = true
-
-      let size = 0
-      if (targetResponse.body) {
-        let reader = targetResponse.body.getReader()
-        let chunk: ReadableStreamReadResult<Uint8Array>
-        do {
-          chunk = await reader.read()
-          if (chunk.value) {
-            res.write(chunk.value)
-            size += chunk.value.length
-            if (size > config.maxSize) {
-              break
-            }
-          }
-        } while (!chunk.done)
-      }
-      res.end()
+      await handleRequestWithDelay(req, res, config, ip, url, parsedUrl)
     } catch (e) {
-      // Known errors
-      if (e instanceof Error && e.name === 'TimeoutError') {
-        res.writeHead(400, { 'Content-Type': 'text/plain' })
-        res.end('Timeout')
-        return
-      } else if (e instanceof BadRequestError) {
-        res.writeHead(e.code, { 'Content-Type': 'text/plain' })
-        res.end(e.message)
-        return
-      }
-
-      // Unknown or Internal errors
-      /* c8 ignore next 9 */
-      if (e instanceof Error) {
-        process.stderr.write(styleText('red', e.stack ?? e.message) + '\n')
-      } else if (typeof e === 'string') {
-        process.stderr.write(styleText('red', e) + '\n')
-      }
-      if (!sent) {
-        res.writeHead(500, { 'Content-Type': 'text/plain' })
-        res.end('Internal Server Error')
-      }
+      handleError(e, res)
     }
   })
 }