easyimg/server/utils/rateLimit.js

import db from '../utils/db.js'

// 存储请求计数的 Map
const requestCounts = new Map()
// 存储正在上传的 IP/ApiKey
const uploadingClients = new Map()

// 清理过期的计数记录（每分钟清理一次）
setInterval(() => {
  const now = Date.now()
  for (const [key, data] of requestCounts.entries()) {
    if (now - data.startTime > 60000) {
      requestCounts.delete(key)
    }
  }
}, 60000)

/**
 * 频率限制中间件
 * @param {string} type - 'public' 或 'private'
 * @param {string} identifier - IP 地址或 ApiKey
 */
export async function checkRateLimit(type, identifier) {
  const configKey = type === 'public' ? 'publicApiConfig' : 'privateApiConfig'
  const config = await db.settings.findOne({ key: configKey })

  if (!config) {
    return { allowed: true }
  }

  const rateLimit = config.value.rateLimit
  const key = `${type}:${identifier}`
  const now = Date.now()

  let record = requestCounts.get(key)

  if (!record || now - record.startTime > 60000) {
    // 新的一分钟周期
    record = { count: 1, startTime: now }
    requestCounts.set(key, record)
    return { allowed: true, remaining: rateLimit - 1 }
  }

  if (record.count >= rateLimit) {
    return {
      allowed: false,
      remaining: 0,
      retryAfter: Math.ceil((record.startTime + 60000 - now) / 1000)
    }
  }

  record.count++
  return { allowed: true, remaining: rateLimit - record.count }
}

/**
 * 并发限制中间件
 * @param {string} type - 'public' 或 'private'
 * @param {string} identifier - IP 地址或 ApiKey
 */
export async function checkConcurrency(type, identifier) {
  const configKey = type === 'public' ? 'publicApiConfig' : 'privateApiConfig'
  const config = await db.settings.findOne({ key: configKey })

  if (!config) {
    return { allowed: true }
  }

  const key = `${type}:${identifier}`

  if (type === 'public') {
    // 公共 API 不允许并发
    if (!config.value.allowConcurrent) {
      if (uploadingClients.has(key)) {
        return { allowed: false, message: '请等待上一张图片上传完成' }
      }
    }
  } else {
    // 私有 API 限制并发数
    const maxConcurrent = config.value.maxConcurrent || 5
    const currentCount = uploadingClients.get(key) || 0

    if (currentCount >= maxConcurrent) {
      return { allowed: false, message: `并发上传数已达上限 (${maxConcurrent})` }
    }
  }

  return { allowed: true }
}

/**
 * 标记开始上传
 */
export function markUploadStart(type, identifier) {
  const key = `${type}:${identifier}`
  const current = uploadingClients.get(key) || 0
  uploadingClients.set(key, current + 1)
}

/**
 * 标记上传结束
 */
export function markUploadEnd(type, identifier) {
  const key = `${type}:${identifier}`
  const current = uploadingClients.get(key) || 0
  if (current <= 1) {
    uploadingClients.delete(key)
  } else {
    uploadingClients.set(key, current - 1)
  }
}

/**
 * 公共 API 频率限制检查
 */
export function checkPublicRateLimit(ip, rateLimit = 10) {
  const key = `public:${ip}`
  const now = Date.now()

  let record = requestCounts.get(key)

  if (!record || now - record.startTime > 60000) {
    record = { count: 1, startTime: now }
    requestCounts.set(key, record)
    return { allowed: true, remaining: rateLimit - 1 }
  }

  if (record.count >= rateLimit) {
    return {
      allowed: false,
      remaining: 0,
      retryAfter: Math.ceil((record.startTime + 60000 - now) / 1000)
    }
  }

  record.count++
  return { allowed: true, remaining: rateLimit - record.count }
}

/**
 * 私有 API 频率限制检查
 */
export function checkPrivateRateLimit(apiKey, rateLimit = 100) {
  const key = `private:${apiKey}`
  const now = Date.now()

  let record = requestCounts.get(key)

  if (!record || now - record.startTime > 60000) {
    record = { count: 1, startTime: now }
    requestCounts.set(key, record)
    return { allowed: true, remaining: rateLimit - 1 }
  }

  if (record.count >= rateLimit) {
    return {
      allowed: false,
      remaining: 0,
      retryAfter: Math.ceil((record.startTime + 60000 - now) / 1000)
    }
  }

  record.count++
  return { allowed: true, remaining: rateLimit - record.count }
}

/**
 * 公共 API 并发检查
 */
export function checkPublicConcurrency(ip) {
  const key = `public:${ip}`
  if (uploadingClients.has(key) && uploadingClients.get(key) > 0) {
    return { allowed: false }
  }
  return { allowed: true }
}

/**
 * 释放公共 API 并发锁
 */
export function releasePublicConcurrency(ip) {
  const key = `public:${ip}`
  uploadingClients.delete(key)
}

/**
 * 获取并标记公共并发
 */
export function acquirePublicConcurrency(ip) {
  const key = `public:${ip}`
  uploadingClients.set(key, 1)
}

/**
 * 私有 API 并发检查
 */
export function checkPrivateConcurrency(ip, maxConcurrent = 5) {
  const key = `private:${ip}`
  const current = uploadingClients.get(key) || 0
  if (current >= maxConcurrent) {
    return { allowed: false, current }
  }
  return { allowed: true, current }
}

/**
 * 获取私有 API 并发锁
 */
export function acquirePrivateConcurrency(ip) {
  const key = `private:${ip}`
  const current = uploadingClients.get(key) || 0
  uploadingClients.set(key, current + 1)
}

/**
 * 释放私有 API 并发锁
 */
export function releasePrivateConcurrency(ip) {
  const key = `private:${ip}`
  const current = uploadingClients.get(key) || 0
  if (current <= 1) {
    uploadingClients.delete(key)
  } else {
    uploadingClients.set(key, current - 1)
  }
}

export default {
  checkRateLimit,
  checkConcurrency,
  markUploadStart,
  markUploadEnd,
  checkPublicRateLimit,
  checkPrivateRateLimit,
  checkPublicConcurrency,
  releasePublicConcurrency,
  acquirePublicConcurrency,
  checkPrivateConcurrency,
  acquirePrivateConcurrency,
  releasePrivateConcurrency
}