Não pode escolher mais do que 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.
 
 
 

345 linhas
9.8 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. logoFilePath := os.Getenv("LOGO_FILE_PATH")
  119. if logoFilePath != "" {
  120. if _, err := os.Stat(logoFilePath); err != nil {
  121. common.SysLog("LOGO_FILE_PATH file not found: " + logoFilePath + ", falling back to default")
  122. } else {
  123. common.LogoFilePath = logoFilePath
  124. common.SysLog("custom logo file: " + common.LogoFilePath)
  125. }
  126. }
  127. if os.Getenv("ENABLE_PPROF") == "true" {
  128. gopool.Go(func() {
  129. log.Println(http.ListenAndServe("0.0.0.0:8005", nil))
  130. })
  131. go common.Monitor()
  132. common.SysLog("pprof enabled")
  133. }
  134. err = common.StartPyroScope()
  135. if err != nil {
  136. common.SysError(fmt.Sprintf("start pyroscope error : %v", err))
  137. }
  138. // Initialize HTTP server
  139. server := gin.New()
  140. server.Use(gin.CustomRecovery(func(c *gin.Context, err any) {
  141. common.SysLog(fmt.Sprintf("panic detected: %v", err))
  142. c.JSON(http.StatusInternalServerError, gin.H{
  143. "error": gin.H{
  144. "message": fmt.Sprintf("Panic detected, error: %v. Please submit a issue here: https://github.com/Calcium-Ion/new-api", err),
  145. "type": "new_api_panic",
  146. },
  147. })
  148. }))
  149. // This will cause SSE not to work!!!
  150. //server.Use(gzip.Gzip(gzip.DefaultCompression))
  151. server.Use(middleware.RequestId())
  152. server.Use(middleware.PoweredBy())
  153. server.Use(middleware.I18n())
  154. middleware.SetUpLogger(server)
  155. // Initialize session store
  156. store := cookie.NewStore([]byte(common.SessionSecret))
  157. store.Options(sessions.Options{
  158. Path: "/",
  159. MaxAge: 2592000, // 30 days
  160. HttpOnly: true,
  161. Secure: false,
  162. SameSite: http.SameSiteStrictMode,
  163. })
  164. // 支持通过环境变量自定义 session 名称,用于多节点部署时避免 cookie 冲突
  165. sessionName := os.Getenv("SESSION_NAME")
  166. if sessionName == "" {
  167. sessionName = "session"
  168. }
  169. server.Use(sessions.Sessions(sessionName, store))
  170. InjectUmamiAnalytics()
  171. InjectGoogleAnalytics()
  172. // 设置路由
  173. router.SetRouter(server, buildFS, indexPage)
  174. var port = os.Getenv("PORT")
  175. if port == "" {
  176. port = strconv.Itoa(*common.Port)
  177. }
  178. // Log startup success message
  179. common.LogStartupSuccess(startTime, port)
  180. err = server.Run(":" + port)
  181. if err != nil {
  182. common.FatalLog("failed to start HTTP server: " + err.Error())
  183. }
  184. }
  185. func InjectUmamiAnalytics() {
  186. analyticsInjectBuilder := &strings.Builder{}
  187. if os.Getenv("UMAMI_WEBSITE_ID") != "" {
  188. umamiSiteID := os.Getenv("UMAMI_WEBSITE_ID")
  189. umamiScriptURL := os.Getenv("UMAMI_SCRIPT_URL")
  190. if umamiScriptURL == "" {
  191. umamiScriptURL = "https://analytics.umami.is/script.js"
  192. }
  193. analyticsInjectBuilder.WriteString("<script defer src=\"")
  194. analyticsInjectBuilder.WriteString(umamiScriptURL)
  195. analyticsInjectBuilder.WriteString("\" data-website-id=\"")
  196. analyticsInjectBuilder.WriteString(umamiSiteID)
  197. analyticsInjectBuilder.WriteString("\"></script>")
  198. }
  199. analyticsInjectBuilder.WriteString("<!--Umami QuantumNous-->\n")
  200. analyticsInject := analyticsInjectBuilder.String()
  201. indexPage = bytes.ReplaceAll(indexPage, []byte("<!--umami-->\n"), []byte(analyticsInject))
  202. }
  203. func InjectGoogleAnalytics() {
  204. analyticsInjectBuilder := &strings.Builder{}
  205. if os.Getenv("GOOGLE_ANALYTICS_ID") != "" {
  206. gaID := os.Getenv("GOOGLE_ANALYTICS_ID")
  207. // Google Analytics 4 (gtag.js)
  208. analyticsInjectBuilder.WriteString("<script async src=\"https://www.googletagmanager.com/gtag/js?id=")
  209. analyticsInjectBuilder.WriteString(gaID)
  210. analyticsInjectBuilder.WriteString("\"></script>")
  211. analyticsInjectBuilder.WriteString("<script>")
  212. analyticsInjectBuilder.WriteString("window.dataLayer = window.dataLayer || [];")
  213. analyticsInjectBuilder.WriteString("function gtag(){dataLayer.push(arguments);}")
  214. analyticsInjectBuilder.WriteString("gtag('js', new Date());")
  215. analyticsInjectBuilder.WriteString("gtag('config', '")
  216. analyticsInjectBuilder.WriteString(gaID)
  217. analyticsInjectBuilder.WriteString("');")
  218. analyticsInjectBuilder.WriteString("</script>")
  219. }
  220. analyticsInjectBuilder.WriteString("<!--Google Analytics QuantumNous-->\n")
  221. analyticsInject := analyticsInjectBuilder.String()
  222. indexPage = bytes.ReplaceAll(indexPage, []byte("<!--Google Analytics-->\n"), []byte(analyticsInject))
  223. }
  224. func InitResources() error {
  225. // Initialize resources here if needed
  226. // This is a placeholder function for future resource initialization
  227. err := godotenv.Load(".env")
  228. if err != nil {
  229. if common.DebugEnabled {
  230. common.SysLog("No .env file found, using default environment variables. If needed, please create a .env file and set the relevant variables.")
  231. }
  232. }
  233. // 加载环境变量
  234. common.InitEnv()
  235. logger.SetupLogger()
  236. // Initialize model settings
  237. ratio_setting.InitRatioSettings()
  238. service.InitHttpClient()
  239. service.InitTokenEncoders()
  240. // Initialize SQL Database
  241. err = model.InitDB()
  242. if err != nil {
  243. common.FatalLog("failed to initialize database: " + err.Error())
  244. return err
  245. }
  246. model.CheckSetup()
  247. // Initialize options, should after model.InitDB()
  248. model.InitOptionMap()
  249. // 清理旧的磁盘缓存文件
  250. common.CleanupOldCacheFiles()
  251. // 初始化模型
  252. model.GetPricing()
  253. // Initialize SQL Database
  254. err = model.InitLogDB()
  255. if err != nil {
  256. return err
  257. }
  258. // Initialize Redis
  259. err = common.InitRedisClient()
  260. if err != nil {
  261. return err
  262. }
  263. // 启动系统监控
  264. common.StartSystemMonitor()
  265. // Initialize i18n
  266. err = i18n.Init()
  267. if err != nil {
  268. common.SysError("failed to initialize i18n: " + err.Error())
  269. // Don't return error, i18n is not critical
  270. } else {
  271. common.SysLog("i18n initialized with languages: " + strings.Join(i18n.SupportedLanguages(), ", "))
  272. }
  273. // Register user language loader for lazy loading
  274. i18n.SetUserLangLoader(model.GetUserLanguage)
  275. // Load custom OAuth providers from database
  276. err = oauth.LoadCustomProviders()
  277. if err != nil {
  278. common.SysError("failed to load custom OAuth providers: " + err.Error())
  279. // Don't return error, custom OAuth is not critical
  280. }
  281. // 注册余额更新回调,Master 节点更新用户余额后推送到 Slave 节点
  282. model.SetQuotaUpdateCallback(func(userId int, quota int) {
  283. if system_setting.GetRegionSyncSettings().IsMaster {
  284. region_sync.PushQuotaUpdateToSlave(userId, quota)
  285. }
  286. })
  287. // Slave 节点启动后台同步任务(批量扣费同步 + 余额定时拉取 + 清理)
  288. syncSettings := system_setting.GetRegionSyncSettings()
  289. if syncSettings.Enabled && !syncSettings.IsMaster {
  290. syncManager := region_sync.NewSyncManager()
  291. syncManager.StartSyncWorkers()
  292. }
  293. return nil
  294. }