omniai-server/src/index.js

require('dotenv').config()
const express = require('express')
const rateLimit = require('express-rate-limit')
const cors = require('cors')
const helmet = require('helmet')
const { startSettlementWorker, stopSettlementWorker } = require('./settlementWorker')
const { startProviderHealthMonitor, stopProviderHealthMonitor } = require('./providerHealthMonitor')
const {
  startStaleTaskCleanup,
  startTaskEventListener,
  startPollerRecovery,
  stopStaleTaskCleanup,
  stopTaskEventListener,
  stopPollerRecovery,
  stopAllPollers,
} = require('./aiTaskWorker')
const { ensureDatabase } = require('./dbSetup')
const { assertRuntimeSecurityConfig } = require('./securityConfig')
const { loadPriceCache } = require('./pricing')

assertRuntimeSecurityConfig()

const routes = require('./routes')

const PORT = Number(process.env.PORT) || 3600
const HOST = process.env.HOST || '0.0.0.0'
const IS_PRODUCTION = process.env.NODE_ENV === 'production'
let server = null
let staleLeaseCleanupTimer = null

// CORS: in production, require explicit allowlist; in dev, allow all with credentials
function buildCorsOptions() {
  const raw = process.env.CORS_ORIGINS || ''
  if (IS_PRODUCTION) {
    if (!raw || raw === '*') {
      console.warn('[security] CORS_ORIGINS not set in production, defaulting to same-origin only')
      return { origin: false, credentials: true }
    }
    return {
      origin: raw.split(',').map((s) => s.trim()),
      credentials: true,
    }
  }
  // Development: allow any origin for local testing
  return { origin: true, credentials: true }
}

async function main() {
  const { createdDefaultAdmin } = await ensureDatabase()
  if (createdDefaultAdmin) {
    console.log('[db] Created default admin user: admin (password sourced from configuration)')
  }

  await loadPriceCache()

  const app = express()

  // Trust single-hop Nginx reverse proxy (X-Forwarded-For, X-Real-IP)
  app.set('trust proxy', 1)

  // Security headers via helmet
  app.use(helmet({
    contentSecurityPolicy: false, // Disable CSP for now (SPA needs inline scripts in dev)
    crossOriginEmbedderPolicy: false, // Allow OSS images/videos
    crossOriginResourcePolicy: { policy: 'cross-origin' },
  }))

  // CORS
  app.use(cors(buildCorsOptions()))

  // Rate limiting: global (300 req/min per IP)
  const globalLimiter = rateLimit({
    windowMs: 60 * 1000,
    max: 300,
    standardHeaders: true,
    legacyHeaders: false,
    message: { error: '请求过于频繁，请稍后再试' },
  })
  app.use('/api', globalLimiter)

  // Rate limiting: auth endpoints (stricter — 10 req/min per IP)
  const authLimiter = rateLimit({
    windowMs: 60 * 1000,
    max: 10,
    standardHeaders: true,
    legacyHeaders: false,
    message: { error: '登录尝试过于频繁，请1分钟后再试' },
  })
  app.use('/api/auth/login', authLimiter)
  app.use('/api/auth/register', authLimiter)
  app.use('/api/auth/sms', authLimiter)

  // Rate limiting: AI generation endpoints (20 req/min per IP)
  const aiGenerationLimiter = rateLimit({
    windowMs: 60 * 1000,
    max: 20,
    standardHeaders: true,
    legacyHeaders: false,
    message: { error: 'AI生成请求过于频繁，请稍后再试' },
  })
  app.use('/api/ai/image', aiGenerationLimiter)
  app.use('/api/ai/video', aiGenerationLimiter)

  // Rate limiting: AI chat endpoint (60 req/min per IP — ecommerce flows need ~7 sequential calls)
  const aiChatLimiter = rateLimit({
    windowMs: 60 * 1000,
    max: 60,
    standardHeaders: true,
    legacyHeaders: false,
    message: { error: 'AI对话请求过于频繁，请稍后再试' },
  })
  app.use('/api/ai/chat', aiChatLimiter)


  // Skip JSON body-parser for binary upload routes (busboy handles multipart parsing)
  app.use('/api/oss/upload-binary', (req, res, next) => { req._body = true; next(); })
  app.use("/api/files/extract-text", (req, res, next) => { req._body = true; next(); })
  // JSON body limit: 5MB globally (upload routes override locally)
  app.use('/api/oss/upload', express.json({ limit: '200mb' }))
  app.use(express.json({ limit: process.env.JSON_BODY_LIMIT || '5mb' }))

  app.use('/api', routes)
  app.use('/api', (req, res) => {
    res.status(404).json({
      error: 'API route not found',
      path: req.originalUrl
    })
  })
  app.use('/api', (err, req, res, next) => {
    if (res.headersSent) return next(err)
    const status = Number(err.status || err.statusCode || 500)
    const safeStatus = status >= 400 && status < 600 ? status : 500
    // Log full error internally, but don't leak details to client
    console.error('[api] error:', err)
    const message = safeStatus >= 500
      ? '服务器内部错误，请稍后重试'
      : (err.message || '请求处理失败')
    res.status(safeStatus).json({
      error: message,
      code: err.code || undefined,
    })
  })

  // Periodic stale lease cleanup (every 5 min)
  const { cleanStaleLeases } = require('./keyManager')
  staleLeaseCleanupTimer = setInterval(() => {
    cleanStaleLeases().then((cleaned) => {
      if (cleaned > 0) console.log(`[cleanup] Released ${cleaned} stale lease(s)`)
    }).catch((err) => {
      console.error('[cleanup] error:', err)
    })
  }, 5 * 60 * 1000)
  if (staleLeaseCleanupTimer.unref) staleLeaseCleanupTimer.unref()

  startSettlementWorker()
  startProviderHealthMonitor()

  await startTaskEventListener()
  startPollerRecovery()
  startStaleTaskCleanup()

  server = app.listen(PORT, HOST, () => {
    console.log(`OmniAI Key Server running at http://${HOST}:${PORT}`)
    console.log(`Health check: http://${HOST}:${PORT}/api/health`)
  })
}

main().catch((err) => {
  console.error('[startup] Fatal error:', err)
  process.exit(1)
})

process.on('unhandledRejection', (reason) => {
  console.error('[fatal] Unhandled promise rejection:', reason)
  process.exitCode = 1
  setTimeout(() => process.exit(1), 5000).unref()
})

process.on('uncaughtException', (err) => {
  console.error('[fatal] Uncaught exception:', err)
  process.exit(1)
})


// ── Graceful shutdown ───────────────────────────────────────────────────
let shuttingDown = false

async function shutdownRuntimeState() {
  if (staleLeaseCleanupTimer) {
    clearInterval(staleLeaseCleanupTimer)
    staleLeaseCleanupTimer = null
  }
  stopSettlementWorker()
  stopProviderHealthMonitor()
  stopPollerRecovery()
  stopStaleTaskCleanup()
  await Promise.allSettled([stopTaskEventListener(), stopAllPollers()])
}

function closeServer() {
  if (!server || !server.listening) return Promise.resolve()
  return new Promise((resolve) => {
    server.close(() => {
      console.log('[shutdown] Server closed, cleaning up...')
      resolve()
    })
  })
}

async function gracefulShutdown(signal) {
  if (shuttingDown) return
  shuttingDown = true
  console.log('[shutdown] Received ' + signal + ', draining connections...')

  setTimeout(() => {
    console.error('[shutdown] Forced exit after timeout')
    process.exit(1)
  }, 15000).unref()

  try {
    await shutdownRuntimeState()
    await closeServer()
    const { pool } = require('./db')
    await pool.end()
    console.log('[shutdown] Database pool closed')
    process.exit(0)
  } catch (err) {
    console.error('[shutdown] error:', err)
    process.exit(0)
  }
}

process.on('SIGTERM', () => gracefulShutdown('SIGTERM'))
process.on('SIGINT', () => gracefulShutdown('SIGINT'))