services/utils/ai-utils.js

const crypto = require('crypto')

const KEYWORD_MAX = 200
const URL_MAX = 2048
const BASE64_MAX = 5 * 1024 * 1024

const BANNED_PATTERNS = [
  /ignore\s+(all\s+)?(previous|prior|above)\s+(instructions?|prompts?)/i,
  /system\s*(prompt|message|role)/i,
  /disregard/i,
  /forget\s+(everything|all)/i,
  /act\s+as\s+(a\s+)?(developer|admin|root)/i,
  /reveal\s+(your|the)\s+(prompt|system|secret|key)/i,
  /api[_-]?key/i,
  /password|secret/i
]

function sanitizeKeyword(input) {
  if (input === undefined || input === null) return ''
  const text = String(input).trim()
  if (text.length === 0) return ''
  if (text.length > KEYWORD_MAX) {
    return { error: `关键词不能超过 ${KEYWORD_MAX} 个字符` }
  }
  for (const re of BANNED_PATTERNS) {
    if (re.test(text)) return { error: '关键词包含非法内容' }
  }
  return { value: text }
}

function sanitizeImageUrl(input) {
  if (!input) return { value: '' }
  const text = String(input).trim()
  if (text.length === 0) return { value: '' }
  if (text.length > URL_MAX) return { error: '图片 URL 过长' }
  if (/^https?:\/\//i.test(text)) {
    let parsed
    try {
      parsed = new URL(text)
    } catch (e) {
      return { error: '图片 URL 格式不合法' }
    }
    if (parsed.protocol !== 'http:' && parsed.protocol !== 'https:') {
      return { error: '图片 URL 协议必须是 http/https' }
    }
    const host = parsed.hostname
    if (!host || host.length === 0) return { error: '图片 URL 缺少主机' }
    if (host.includes('\u0000') || /[\s<>"'`\\]/.test(host)) {
      return { error: '图片 URL 主机名非法' }
    }
    const labels = host.split('.')
    for (const label of labels) {
      if (label.length === 0 || label.length > 63) return { error: '图片 URL 主机名非法' }
      if (!/^[a-z0-9](?:[a-z0-9-]*[a-z0-9])?$/i.test(label)) {
        return { error: '图片 URL 主机名非法' }
      }
    }
    if (text.includes('..') || text.includes('\u0000')) return { error: '图片 URL 格式不合法' }
  } else if (!text.startsWith('/uploads/')) {
    return { error: '图片 URL 必须是 https/http 链接或 /uploads/ 相对路径' }
  }
  return { value: text }
}

function sanitizeImageBase64(input) {
  if (!input) return { error: '缺少图片数据' }
  const text = String(input)
  if (text.length > BASE64_MAX) return { error: `图片不能超过 ${(BASE64_MAX / 1024 / 1024).toFixed(1)}MB` }
  if (!/^[A-Za-z0-9+/=\r\n]+$/.test(text)) return { error: '图片 base64 格式不合法' }
  return { value: text }
}

function makeCacheKey(prefix, payload) {
  const json = JSON.stringify(payload, Object.keys(payload).sort())
  return prefix + ':' + crypto.createHash('sha1').update(json).digest('hex')
}

class LRU {
  constructor(max = 100, ttlMs = 5 * 60 * 1000) {
    this.max = max
    this.ttlMs = ttlMs
    this.map = new Map()
  }
  get(key) {
    const item = this.map.get(key)
    if (!item) return undefined
    if (Date.now() > item.expire) {
      this.map.delete(key)
      return undefined
    }
    this.map.delete(key)
    this.map.set(key, item)
    return item.value
  }
  set(key, value) {
    if (this.map.has(key)) this.map.delete(key)
    this.map.set(key, { value, expire: Date.now() + this.ttlMs })
    if (this.map.size > this.max) {
      const first = this.map.keys().next().value
      this.map.delete(first)
    }
  }
}

class TokenBucket {
  constructor(capacity, refillPerSec) {
    this.capacity = capacity
    this.tokens = capacity
    this.refillPerSec = refillPerSec
    this.lastRefill = Date.now()
    this.queue = []
  }
  take() {
    this._refill()
    if (this.tokens >= 1) {
      this.tokens -= 1
      return true
    }
    return false
  }
  _refill() {
    const now = Date.now()
    const elapsed = (now - this.lastRefill) / 1000
    if (elapsed > 0) {
      this.tokens = Math.min(this.capacity, this.tokens + elapsed * this.refillPerSec)
      this.lastRefill = now
    }
  }
}

module.exports = {
  KEYWORD_MAX,
  URL_MAX,
  BASE64_MAX,
  sanitizeKeyword,
  sanitizeImageUrl,
  sanitizeImageBase64,
  makeCacheKey,
  LRU,
  TokenBucket
}