Non puoi selezionare più di 25 argomenti Gli argomenti devono iniziare con una lettera o un numero, possono includere trattini ('-') e possono essere lunghi fino a 35 caratteri.
 
 
 

335 righe
9.5 KiB

  1. package main
  2. import (
  3. "bytes"
  4. "embed"
  5. "fmt"
  6. "log"
  7. "net/http"
  8. "os"
  9. "strconv"
  10. "strings"
  11. "time"
  12. "github.com/QuantumNous/new-api/common"
  13. "github.com/QuantumNous/new-api/constant"
  14. "github.com/QuantumNous/new-api/controller"
  15. "github.com/QuantumNous/new-api/i18n"
  16. "github.com/QuantumNous/new-api/logger"
  17. "github.com/QuantumNous/new-api/middleware"
  18. "github.com/QuantumNous/new-api/model"
  19. "github.com/QuantumNous/new-api/oauth"
  20. "github.com/QuantumNous/new-api/relay"
  21. "github.com/QuantumNous/new-api/router"
  22. "github.com/QuantumNous/new-api/service"
  23. "github.com/QuantumNous/new-api/service/region_sync"
  24. _ "github.com/QuantumNous/new-api/setting/performance_setting"
  25. "github.com/QuantumNous/new-api/setting/ratio_setting"
  26. "github.com/QuantumNous/new-api/setting/system_setting"
  27. "github.com/bytedance/gopkg/util/gopool"
  28. "github.com/gin-contrib/sessions"
  29. "github.com/gin-contrib/sessions/cookie"
  30. "github.com/gin-gonic/gin"
  31. "github.com/joho/godotenv"
  32. _ "net/http/pprof"
  33. )
  34. //go:embed web/dist
  35. var buildFS embed.FS
  36. //go:embed web/dist/index.html
  37. var indexPage []byte
  38. func main() {
  39. startTime := time.Now()
  40. err := InitResources()
  41. if err != nil {
  42. common.FatalLog("failed to initialize resources: " + err.Error())
  43. return
  44. }
  45. common.SysLog("New API " + common.Version + " started")
  46. if os.Getenv("GIN_MODE") != "debug" {
  47. gin.SetMode(gin.ReleaseMode)
  48. }
  49. if common.DebugEnabled {
  50. common.SysLog("running in debug mode")
  51. }
  52. defer func() {
  53. err := model.CloseDB()
  54. if err != nil {
  55. common.FatalLog("failed to close database: " + err.Error())
  56. }
  57. }()
  58. if common.RedisEnabled {
  59. // for compatibility with old versions
  60. common.MemoryCacheEnabled = true
  61. }
  62. if common.MemoryCacheEnabled {
  63. common.SysLog("memory cache enabled")
  64. common.SysLog(fmt.Sprintf("sync frequency: %d seconds", common.SyncFrequency))
  65. // Add panic recovery and retry for InitChannelCache
  66. func() {
  67. defer func() {
  68. if r := recover(); r != nil {
  69. common.SysLog(fmt.Sprintf("InitChannelCache panic: %v, retrying once", r))
  70. // Retry once
  71. _, _, fixErr := model.FixAbility()
  72. if fixErr != nil {
  73. common.FatalLog(fmt.Sprintf("InitChannelCache failed: %s", fixErr.Error()))
  74. }
  75. }
  76. }()
  77. model.InitChannelCache()
  78. }()
  79. go model.SyncChannelCache(common.SyncFrequency)
  80. }
  81. // 热更新配置
  82. go model.SyncOptions(common.SyncFrequency)
  83. // 数据看板
  84. go model.UpdateQuotaData()
  85. if os.Getenv("CHANNEL_UPDATE_FREQUENCY") != "" {
  86. frequency, err := strconv.Atoi(os.Getenv("CHANNEL_UPDATE_FREQUENCY"))
  87. if err != nil {
  88. common.FatalLog("failed to parse CHANNEL_UPDATE_FREQUENCY: " + err.Error())
  89. }
  90. go controller.AutomaticallyUpdateChannels(frequency)
  91. }
  92. go controller.AutomaticallyTestChannels()
  93. // Codex credential auto-refresh check every 10 minutes, refresh when expires within 1 day
  94. service.StartCodexCredentialAutoRefreshTask()
  95. // Subscription quota reset task (daily/weekly/monthly/custom)
  96. service.StartSubscriptionQuotaResetTask()
  97. // Wire task polling adaptor factory (breaks service -> relay import cycle)
  98. service.GetTaskAdaptorFunc = func(platform constant.TaskPlatform) service.TaskPollingAdaptor {
  99. a := relay.GetTaskAdaptor(platform)
  100. if a == nil {
  101. return nil
  102. }
  103. return a
  104. }
  105. if common.IsMasterNode && constant.UpdateTask {
  106. gopool.Go(func() {
  107. controller.UpdateMidjourneyTaskBulk()
  108. })
  109. gopool.Go(func() {
  110. controller.UpdateTaskBulk()
  111. })
  112. }
  113. if os.Getenv("BATCH_UPDATE_ENABLED") == "true" {
  114. common.BatchUpdateEnabled = true
  115. common.SysLog("batch update enabled with interval " + strconv.Itoa(common.BatchUpdateInterval) + "s")
  116. model.InitBatchUpdater()
  117. }
  118. if os.Getenv("ENABLE_PPROF") == "true" {
  119. gopool.Go(func() {
  120. log.Println(http.ListenAndServe("0.0.0.0:8005", nil))
  121. })
  122. go common.Monitor()
  123. common.SysLog("pprof enabled")
  124. }
  125. err = common.StartPyroScope()
  126. if err != nil {
  127. common.SysError(fmt.Sprintf("start pyroscope error : %v", err))
  128. }
  129. // Initialize HTTP server
  130. server := gin.New()
  131. server.Use(gin.CustomRecovery(func(c *gin.Context, err any) {
  132. common.SysLog(fmt.Sprintf("panic detected: %v", err))
  133. c.JSON(http.StatusInternalServerError, gin.H{
  134. "error": gin.H{
  135. "message": fmt.Sprintf("Panic detected, error: %v. Please submit a issue here: https://github.com/Calcium-Ion/new-api", err),
  136. "type": "new_api_panic",
  137. },
  138. })
  139. }))
  140. // This will cause SSE not to work!!!
  141. //server.Use(gzip.Gzip(gzip.DefaultCompression))
  142. server.Use(middleware.RequestId())
  143. server.Use(middleware.PoweredBy())
  144. server.Use(middleware.I18n())
  145. middleware.SetUpLogger(server)
  146. // Initialize session store
  147. store := cookie.NewStore([]byte(common.SessionSecret))
  148. store.Options(sessions.Options{
  149. Path: "/",
  150. MaxAge: 2592000, // 30 days
  151. HttpOnly: true,
  152. Secure: false,
  153. SameSite: http.SameSiteStrictMode,
  154. })
  155. // 支持通过环境变量自定义 session 名称,用于多节点部署时避免 cookie 冲突
  156. sessionName := os.Getenv("SESSION_NAME")
  157. if sessionName == "" {
  158. sessionName = "session"
  159. }
  160. server.Use(sessions.Sessions(sessionName, store))
  161. InjectUmamiAnalytics()
  162. InjectGoogleAnalytics()
  163. // 设置路由
  164. router.SetRouter(server, buildFS, indexPage)
  165. var port = os.Getenv("PORT")
  166. if port == "" {
  167. port = strconv.Itoa(*common.Port)
  168. }
  169. // Log startup success message
  170. common.LogStartupSuccess(startTime, port)
  171. err = server.Run(":" + port)
  172. if err != nil {
  173. common.FatalLog("failed to start HTTP server: " + err.Error())
  174. }
  175. }
  176. func InjectUmamiAnalytics() {
  177. analyticsInjectBuilder := &strings.Builder{}
  178. if os.Getenv("UMAMI_WEBSITE_ID") != "" {
  179. umamiSiteID := os.Getenv("UMAMI_WEBSITE_ID")
  180. umamiScriptURL := os.Getenv("UMAMI_SCRIPT_URL")
  181. if umamiScriptURL == "" {
  182. umamiScriptURL = "https://analytics.umami.is/script.js"
  183. }
  184. analyticsInjectBuilder.WriteString("<script defer src=\"")
  185. analyticsInjectBuilder.WriteString(umamiScriptURL)
  186. analyticsInjectBuilder.WriteString("\" data-website-id=\"")
  187. analyticsInjectBuilder.WriteString(umamiSiteID)
  188. analyticsInjectBuilder.WriteString("\"></script>")
  189. }
  190. analyticsInjectBuilder.WriteString("<!--Umami QuantumNous-->\n")
  191. analyticsInject := analyticsInjectBuilder.String()
  192. indexPage = bytes.ReplaceAll(indexPage, []byte("<!--umami-->\n"), []byte(analyticsInject))
  193. }
  194. func InjectGoogleAnalytics() {
  195. analyticsInjectBuilder := &strings.Builder{}
  196. if os.Getenv("GOOGLE_ANALYTICS_ID") != "" {
  197. gaID := os.Getenv("GOOGLE_ANALYTICS_ID")
  198. // Google Analytics 4 (gtag.js)
  199. analyticsInjectBuilder.WriteString("<script async src=\"https://www.googletagmanager.com/gtag/js?id=")
  200. analyticsInjectBuilder.WriteString(gaID)
  201. analyticsInjectBuilder.WriteString("\"></script>")
  202. analyticsInjectBuilder.WriteString("<script>")
  203. analyticsInjectBuilder.WriteString("window.dataLayer = window.dataLayer || [];")
  204. analyticsInjectBuilder.WriteString("function gtag(){dataLayer.push(arguments);}")
  205. analyticsInjectBuilder.WriteString("gtag('js', new Date());")
  206. analyticsInjectBuilder.WriteString("gtag('config', '")
  207. analyticsInjectBuilder.WriteString(gaID)
  208. analyticsInjectBuilder.WriteString("');")
  209. analyticsInjectBuilder.WriteString("</script>")
  210. }
  211. analyticsInjectBuilder.WriteString("<!--Google Analytics QuantumNous-->\n")
  212. analyticsInject := analyticsInjectBuilder.String()
  213. indexPage = bytes.ReplaceAll(indexPage, []byte("<!--Google Analytics-->\n"), []byte(analyticsInject))
  214. }
  215. func InitResources() error {
  216. // Initialize resources here if needed
  217. // This is a placeholder function for future resource initialization
  218. err := godotenv.Load(".env")
  219. if err != nil {
  220. if common.DebugEnabled {
  221. common.SysLog("No .env file found, using default environment variables. If needed, please create a .env file and set the relevant variables.")
  222. }
  223. }
  224. // 加载环境变量
  225. common.InitEnv()
  226. logger.SetupLogger()
  227. // Initialize model settings
  228. ratio_setting.InitRatioSettings()
  229. service.InitHttpClient()
  230. service.InitTokenEncoders()
  231. // Initialize SQL Database
  232. err = model.InitDB()
  233. if err != nil {
  234. common.FatalLog("failed to initialize database: " + err.Error())
  235. return err
  236. }
  237. model.CheckSetup()
  238. // Initialize options, should after model.InitDB()
  239. model.InitOptionMap()
  240. // 清理旧的磁盘缓存文件
  241. common.CleanupOldCacheFiles()
  242. // 初始化模型
  243. model.GetPricing()
  244. // Initialize SQL Database
  245. err = model.InitLogDB()
  246. if err != nil {
  247. return err
  248. }
  249. // Initialize Redis
  250. err = common.InitRedisClient()
  251. if err != nil {
  252. return err
  253. }
  254. // 启动系统监控
  255. common.StartSystemMonitor()
  256. // Initialize i18n
  257. err = i18n.Init()
  258. if err != nil {
  259. common.SysError("failed to initialize i18n: " + err.Error())
  260. // Don't return error, i18n is not critical
  261. } else {
  262. common.SysLog("i18n initialized with languages: " + strings.Join(i18n.SupportedLanguages(), ", "))
  263. }
  264. // Register user language loader for lazy loading
  265. i18n.SetUserLangLoader(model.GetUserLanguage)
  266. // Load custom OAuth providers from database
  267. err = oauth.LoadCustomProviders()
  268. if err != nil {
  269. common.SysError("failed to load custom OAuth providers: " + err.Error())
  270. // Don't return error, custom OAuth is not critical
  271. }
  272. // 注册余额更新回调,Master 节点更新用户余额后推送到 Slave 节点
  273. model.SetQuotaUpdateCallback(func(userId int, quota int) {
  274. if system_setting.GetRegionSyncSettings().IsMaster {
  275. region_sync.PushQuotaUpdateToSlave(userId, quota)
  276. }
  277. })
  278. // Slave 节点启动后台同步任务(批量扣费同步 + 余额定时拉取 + 清理)
  279. syncSettings := system_setting.GetRegionSyncSettings()
  280. if syncSettings.Enabled && !syncSettings.IsMaster {
  281. syncManager := region_sync.NewSyncManager()
  282. syncManager.StartSyncWorkers()
  283. }
  284. return nil
  285. }