Nie możesz wybrać więcej, niż 25 tematów Tematy muszą się zaczynać od litery lub cyfry, mogą zawierać myślniki ('-') i mogą mieć do 35 znaków.
 
 
 

356 wiersze
10 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: os.Getenv("SESSION_SECURE") != "false",
  162. SameSite: func() http.SameSite {
  163. if os.Getenv("SESSION_SAMESITE") == "strict" {
  164. return http.SameSiteStrictMode
  165. }
  166. return http.SameSiteNoneMode
  167. }(),
  168. })
  169. // 支持通过环境变量自定义 session 名称,用于多节点部署时避免 cookie 冲突
  170. sessionName := os.Getenv("SESSION_NAME")
  171. if sessionName == "" {
  172. sessionName = "session"
  173. }
  174. server.Use(sessions.Sessions(sessionName, store))
  175. InjectUmamiAnalytics()
  176. InjectGoogleAnalytics()
  177. // 设置路由
  178. router.SetRouter(server, buildFS, indexPage)
  179. var port = os.Getenv("PORT")
  180. if port == "" {
  181. port = strconv.Itoa(*common.Port)
  182. }
  183. // Log startup success message
  184. common.LogStartupSuccess(startTime, port)
  185. err = server.Run(":" + port)
  186. if err != nil {
  187. common.FatalLog("failed to start HTTP server: " + err.Error())
  188. }
  189. }
  190. func InjectUmamiAnalytics() {
  191. analyticsInjectBuilder := &strings.Builder{}
  192. if os.Getenv("UMAMI_WEBSITE_ID") != "" {
  193. umamiSiteID := os.Getenv("UMAMI_WEBSITE_ID")
  194. umamiScriptURL := os.Getenv("UMAMI_SCRIPT_URL")
  195. if umamiScriptURL == "" {
  196. umamiScriptURL = "https://analytics.umami.is/script.js"
  197. }
  198. analyticsInjectBuilder.WriteString("<script defer src=\"")
  199. analyticsInjectBuilder.WriteString(umamiScriptURL)
  200. analyticsInjectBuilder.WriteString("\" data-website-id=\"")
  201. analyticsInjectBuilder.WriteString(umamiSiteID)
  202. analyticsInjectBuilder.WriteString("\"></script>")
  203. }
  204. analyticsInjectBuilder.WriteString("<!--Umami QuantumNous-->\n")
  205. analyticsInject := analyticsInjectBuilder.String()
  206. indexPage = bytes.ReplaceAll(indexPage, []byte("<!--umami-->\n"), []byte(analyticsInject))
  207. }
  208. func InjectGoogleAnalytics() {
  209. analyticsInjectBuilder := &strings.Builder{}
  210. if os.Getenv("GOOGLE_ANALYTICS_ID") != "" {
  211. gaID := os.Getenv("GOOGLE_ANALYTICS_ID")
  212. // Google Analytics 4 (gtag.js)
  213. analyticsInjectBuilder.WriteString("<script async src=\"https://www.googletagmanager.com/gtag/js?id=")
  214. analyticsInjectBuilder.WriteString(gaID)
  215. analyticsInjectBuilder.WriteString("\"></script>")
  216. analyticsInjectBuilder.WriteString("<script>")
  217. analyticsInjectBuilder.WriteString("window.dataLayer = window.dataLayer || [];")
  218. analyticsInjectBuilder.WriteString("function gtag(){dataLayer.push(arguments);}")
  219. analyticsInjectBuilder.WriteString("gtag('js', new Date());")
  220. analyticsInjectBuilder.WriteString("gtag('config', '")
  221. analyticsInjectBuilder.WriteString(gaID)
  222. analyticsInjectBuilder.WriteString("');")
  223. analyticsInjectBuilder.WriteString("</script>")
  224. }
  225. analyticsInjectBuilder.WriteString("<!--Google Analytics QuantumNous-->\n")
  226. analyticsInject := analyticsInjectBuilder.String()
  227. indexPage = bytes.ReplaceAll(indexPage, []byte("<!--Google Analytics-->\n"), []byte(analyticsInject))
  228. }
  229. func InitResources() error {
  230. // Initialize resources here if needed
  231. // This is a placeholder function for future resource initialization
  232. err := godotenv.Load(".env")
  233. if err != nil {
  234. if common.DebugEnabled {
  235. common.SysLog("No .env file found, using default environment variables. If needed, please create a .env file and set the relevant variables.")
  236. }
  237. }
  238. // 加载环境变量
  239. common.InitEnv()
  240. logger.SetupLogger()
  241. // Initialize model settings
  242. ratio_setting.InitRatioSettings()
  243. service.InitHttpClient()
  244. service.InitTokenEncoders()
  245. // Initialize SQL Database
  246. err = model.InitDB()
  247. if err != nil {
  248. common.FatalLog("failed to initialize database: " + err.Error())
  249. return err
  250. }
  251. model.CheckSetup()
  252. // Initialize options, should after model.InitDB()
  253. model.InitOptionMap()
  254. model.LoadOptionsFromDatabase()
  255. // 清理旧的磁盘缓存文件
  256. common.CleanupOldCacheFiles()
  257. // 初始化模型
  258. model.GetPricing()
  259. // 加载抓包用户缓存
  260. if err := middleware.LoadCaptureEnabledUsers(); err != nil {
  261. common.SysError("failed to load capture enabled users: " + err.Error())
  262. }
  263. // Initialize SQL Database
  264. err = model.InitLogDB()
  265. if err != nil {
  266. return err
  267. }
  268. // Initialize Redis
  269. err = common.InitRedisClient()
  270. if err != nil {
  271. return err
  272. }
  273. // 启动系统监控
  274. common.StartSystemMonitor()
  275. // Initialize i18n
  276. err = i18n.Init()
  277. if err != nil {
  278. common.SysError("failed to initialize i18n: " + err.Error())
  279. // Don't return error, i18n is not critical
  280. } else {
  281. common.SysLog("i18n initialized with languages: " + strings.Join(i18n.SupportedLanguages(), ", "))
  282. }
  283. // Register user language loader for lazy loading
  284. i18n.SetUserLangLoader(model.GetUserLanguage)
  285. // Load custom OAuth providers from database
  286. err = oauth.LoadCustomProviders()
  287. if err != nil {
  288. common.SysError("failed to load custom OAuth providers: " + err.Error())
  289. // Don't return error, custom OAuth is not critical
  290. }
  291. // 注册余额更新回调,Master 节点更新用户余额后推送到 Slave 节点
  292. model.SetQuotaUpdateCallback(func(userId int, quota int) {
  293. if system_setting.GetRegionSyncSettings().IsMaster {
  294. region_sync.PushQuotaUpdateToSlave(userId, quota)
  295. }
  296. })
  297. // Slave 节点启动后台同步任务(批量扣费同步 + 余额定时拉取 + 清理)
  298. syncSettings := system_setting.GetRegionSyncSettings()
  299. if syncSettings.Enabled && !syncSettings.IsMaster {
  300. syncManager := region_sync.NewSyncManager()
  301. syncManager.StartSyncWorkers()
  302. }
  303. return nil
  304. }