order_site/apps/backend/src/middleware/rate-limit.ts

import type { NextFunction, Request, Response } from 'express'

import { buildErrorPayload, createHttpError } from '../utils/http.js'
import { logWarn } from '../utils/logger.js'

type RateLimitKeyResolver = (req: Request) => string

type RateLimitExceededContext = {
  scope: string
  retryAfterSeconds: number
}

type RateLimitOptions = {
  scope: string
  windowMs: number
  max: number
  key?: RateLimitKeyResolver
  onLimit?: (req: Request, res: Response, context: RateLimitExceededContext) => void
}

type RateLimitBucket = {
  count: number
  resetAt: number
}

const buckets = new Map<string, RateLimitBucket>()
let lastCleanupAt = 0
/** 限流桶数量硬上限：key 含请求体字段时可被构造成高基数，必须兜底防止内存无限增长。 */
const MAX_RATE_LIMIT_BUCKETS = 50_000

export function createRateLimitMiddleware({
  scope,
  windowMs,
  max,
  key = defaultRateLimitKey,
  onLimit,
}: RateLimitOptions) {
  const normalizedScope = String(scope || 'default').trim() || 'default'
  const normalizedWindowMs = Math.max(1000, Number(windowMs || 0))
  const normalizedMax = Math.max(1, Number(max || 0))

  return (req: Request, res: Response, next: NextFunction): void => {
    const now = Date.now()
    cleanupExpiredBuckets(now)

    const bucketKey = `${normalizedScope}:${key(req)}`
    const current = buckets.get(bucketKey)
    const bucket =
      current && current.resetAt > now ? current : { count: 0, resetAt: now + normalizedWindowMs }

    bucket.count += 1
    buckets.set(bucketKey, bucket)

    if (bucket.count <= normalizedMax) {
      next()
      return
    }

    const retryAfterSeconds = Math.max(1, Math.ceil((bucket.resetAt - now) / 1000))
    res.setHeader('Retry-After', String(retryAfterSeconds))

    logWarn('[rate-limit]', '请求触发限流', {
      scope: normalizedScope,
      ip: req.ip,
      originalUrl: req.originalUrl,
      retryAfterSeconds,
    })

    if (onLimit) {
      onLimit(req, res, {
        scope: normalizedScope,
        retryAfterSeconds,
      })
      return
    }

    const error = createHttpError('请求过于频繁，请稍后再试', {
      statusCode: 429,
      errorCode: 'rate_limited',
    })
    res.status(429).json(buildErrorPayload(error, '请求过于频繁，请稍后再试'))
  }
}

export function resetRateLimitBucketsForTest(): void {
  buckets.clear()
  lastCleanupAt = 0
}

export function getBodyFieldRateLimitKey(fieldName: string): RateLimitKeyResolver {
  return (req) => [defaultRateLimitKey(req), normalizeKeyPart(req.body?.[fieldName])].join(':')
}

export function getParamRateLimitKey(paramName: string): RateLimitKeyResolver {
  return (req) => [defaultRateLimitKey(req), normalizeKeyPart(req.params?.[paramName])].join(':')
}

function defaultRateLimitKey(req: Request): string {
  return normalizeKeyPart(
    req.ip ||
      String(req.headers['x-forwarded-for'] || '').split(',')[0] ||
      req.socket.remoteAddress ||
      'unknown',
  )
}

function normalizeKeyPart(value: unknown): string {
  const normalized = String(value || '')
    .trim()
    .toLowerCase()
  return normalized || 'unknown'
}

function cleanupExpiredBuckets(now: number): void {
  if (now - lastCleanupAt < 60_000) {
    return
  }

  lastCleanupAt = now
  for (const [key, bucket] of buckets.entries()) {
    if (bucket.resetAt <= now) {
      buckets.delete(key)
    }
  }

  // 兜底：限流 key 里含请求体字段（如 collect 路由的 orderNo），是可被构造的高基数键，
  // 只按"过期"清理挡不住短时间内灌进来的海量新键。这里再加一道硬上限，
  // 超限时按插入顺序淘汰最老的桶（Map 保持插入序），保证内存不随扫描流量无限涨。
  if (buckets.size > MAX_RATE_LIMIT_BUCKETS) {
    const overflow = buckets.size - MAX_RATE_LIMIT_BUCKETS
    let removed = 0
    for (const key of buckets.keys()) {
      buckets.delete(key)
      removed += 1
      if (removed >= overflow) break
    }
  }
}