1
0

db.go 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138
  1. package database
  2. import (
  3. "bytes"
  4. "io"
  5. "io/fs"
  6. "log"
  7. "os"
  8. "path"
  9. "x-ui/config"
  10. "x-ui/database/model"
  11. "x-ui/xray"
  12. "gorm.io/driver/sqlite"
  13. "gorm.io/gorm"
  14. "gorm.io/gorm/logger"
  15. )
  16. var db *gorm.DB
  17. const (
  18. defaultUsername = "admin"
  19. defaultPassword = "admin"
  20. defaultSecret = ""
  21. )
  22. func initModels() error {
  23. models := []interface{}{
  24. &model.User{},
  25. &model.Inbound{},
  26. &model.OutboundTraffics{},
  27. &model.Setting{},
  28. &model.InboundClientIps{},
  29. &xray.ClientTraffic{},
  30. }
  31. for _, model := range models {
  32. if err := db.AutoMigrate(model); err != nil {
  33. log.Printf("Error auto migrating model: %v", err)
  34. return err
  35. }
  36. }
  37. return nil
  38. }
  39. func initUser() error {
  40. empty, err := isTableEmpty("users")
  41. if err != nil {
  42. log.Printf("Error checking if users table is empty: %v", err)
  43. return err
  44. }
  45. if empty {
  46. user := &model.User{
  47. Username: defaultUsername,
  48. Password: defaultPassword,
  49. LoginSecret: defaultSecret,
  50. }
  51. return db.Create(user).Error
  52. }
  53. return nil
  54. }
  55. func isTableEmpty(tableName string) (bool, error) {
  56. var count int64
  57. err := db.Table(tableName).Count(&count).Error
  58. return count == 0, err
  59. }
  60. func InitDB(dbPath string) error {
  61. dir := path.Dir(dbPath)
  62. err := os.MkdirAll(dir, fs.ModePerm)
  63. if err != nil {
  64. return err
  65. }
  66. var gormLogger logger.Interface
  67. if config.IsDebug() {
  68. gormLogger = logger.Default
  69. } else {
  70. gormLogger = logger.Discard
  71. }
  72. c := &gorm.Config{
  73. Logger: gormLogger,
  74. }
  75. db, err = gorm.Open(sqlite.Open(dbPath), c)
  76. if err != nil {
  77. return err
  78. }
  79. if err := initModels(); err != nil {
  80. return err
  81. }
  82. if err := initUser(); err != nil {
  83. return err
  84. }
  85. return nil
  86. }
  87. func CloseDB() error {
  88. if db != nil {
  89. sqlDB, err := db.DB()
  90. if err != nil {
  91. return err
  92. }
  93. return sqlDB.Close()
  94. }
  95. return nil
  96. }
  97. func GetDB() *gorm.DB {
  98. return db
  99. }
  100. func IsNotFound(err error) bool {
  101. return err == gorm.ErrRecordNotFound
  102. }
  103. func IsSQLiteDB(file io.ReaderAt) (bool, error) {
  104. signature := []byte("SQLite format 3\x00")
  105. buf := make([]byte, len(signature))
  106. _, err := file.ReadAt(buf, 0)
  107. if err != nil {
  108. return false, err
  109. }
  110. return bytes.Equal(buf, signature), nil
  111. }
  112. func Checkpoint() error {
  113. // Update WAL
  114. err := db.Exec("PRAGMA wal_checkpoint;").Error
  115. if err != nil {
  116. return err
  117. }
  118. return nil
  119. }