fix(database): restore working InitDB/GetDB with automigrate and admin seed

database/db.go

@@ -1,198 +1,73 @@
-// Package database provides database initialization, migration, and management utilities
-// for the 3x-ui panel using GORM with SQLite.
 package database
 
 import (
-	"bytes"
-	"io"
 	"io/fs"
-	"log"
 	"os"
 	"path"
-	"slices"
 
-	"github.com/mhsanaei/3x-ui/v2/config"
 	"github.com/mhsanaei/3x-ui/v2/database/model"
-	"github.com/mhsanaei/3x-ui/v2/util/crypto"
-	"github.com/mhsanaei/3x-ui/v2/xray"
-
+	"golang.org/x/crypto/bcrypt"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
-	"gorm.io/gorm/logger"
 )
 
 var db *gorm.DB
 
-const (
-	defaultUsername = "admin"
-	defaultPassword = "admin"
-)
-
-func initModels() error {
-	models := []any{
-		&model.User{},
-		&model.Inbound{},
-		&model.OutboundTraffics{},
-		&model.Setting{},
-		&model.InboundClientIps{},
-		&xray.ClientTraffic{},
-		&model.HistoryOfSeeders{},
-	}
-	for _, model := range models {
-		if err := db.AutoMigrate(model); err != nil {
-			log.Printf("Error auto migrating model: %v", err)
-			return err
-		}
-	}
-	return nil
-}
-
-// initUser creates a default admin user if the users table is empty.
-func initUser() error {
-	empty, err := isTableEmpty("users")
-	if err != nil {
-		log.Printf("Error checking if users table is empty: %v", err)
-		return err
-	}
-	if empty {
-		hashedPassword, err := crypto.HashPasswordAsBcrypt(defaultPassword)
-
-		if err != nil {
-			log.Printf("Error hashing default password: %v", err)
-			return err
-		}
-
-		user := &model.User{
-			Username: defaultUsername,
-			Password: hashedPassword,
-		}
-		return db.Create(user).Error
-	}
-	return nil
-}
-
-// runSeeders migrates user passwords to bcrypt and records seeder execution to prevent re-running.
-func runSeeders(isUsersEmpty bool) error {
-	empty, err := isTableEmpty("history_of_seeders")
-	if err != nil {
-		log.Printf("Error checking if users table is empty: %v", err)
-		return err
-	}
-
-	if empty && isUsersEmpty {
-		hashSeeder := &model.HistoryOfSeeders{
-			SeederName: "UserPasswordHash",
-		}
-		return db.Create(hashSeeder).Error
-	} else {
-		var seedersHistory []string
-		db.Model(&model.HistoryOfSeeders{}).Pluck("seeder_name", &seedersHistory)
-
-		if !slices.Contains(seedersHistory, "UserPasswordHash") && !isUsersEmpty {
-			var users []model.User
-			db.Find(&users)
-
-			for _, user := range users {
-				hashedPassword, err := crypto.HashPasswordAsBcrypt(user.Password)
-				if err != nil {
-					log.Printf("Error hashing password for user '%s': %v", user.Username, err)
-					return err
-				}
-				db.Model(&user).Update("password", hashedPassword)
-			}
-
-			hashSeeder := &model.HistoryOfSeeders{
-				SeederName: "UserPasswordHash",
-			}
-			return db.Create(hashSeeder).Error
-		}
-	}
-
-	return nil
-}
-
-// isTableEmpty returns true if the named table contains zero rows.
-func isTableEmpty(tableName string) (bool, error) {
-	var count int64
-	err := db.Table(tableName).Count(&count).Error
-	return count == 0, err
-}
-
-// InitDB sets up the database connection, migrates models, and runs seeders.
-func InitDB(dbPath string) error {
-    dir := path.Dir(dbPath)
-    if err := os.MkdirAll(dir, fs.ModePerm); err != nil { return err }
-    // ...
-}
-
-	var gormLogger logger.Interface
-
-	if config.IsDebug() {
-		gormLogger = logger.Default
-	} else {
-		gormLogger = logger.Discard
-	}
-
-	c := &gorm.Config{
-		Logger: gormLogger,
-	}
-	db, err = gorm.Open(sqlite.Open(dbPath), c)
-	if err != nil {
-		return err
-	}
-
-	if err := initModels(); err != nil {
-		return err
-	}
-
-	isUsersEmpty, err := isTableEmpty("users")
-	if err != nil {
-		return err
-	}
-
-	if err := initUser(); err != nil {
-		return err
-	}
-	return runSeeders(isUsersEmpty)
-}
-
-// CloseDB closes the database connection if it exists.
-func CloseDB() error {
-	if db != nil {
-		sqlDB, err := db.DB()
-		if err != nil {
-			return err
-		}
-		return sqlDB.Close()
-	}
-	return nil
-}
-
 // GetDB returns the global GORM database instance.
 func GetDB() *gorm.DB { return db }
 
-// IsNotFound checks if the given error is a GORM record not found error.
-func IsNotFound(err error) bool {
-	return err == gorm.ErrRecordNotFound
-}
-
-// IsSQLiteDB checks if the given file is a valid SQLite database by reading its signature.
-func IsSQLiteDB(file io.ReaderAt) (bool, error) {
-	signature := []byte("SQLite format 3\x00")
-	buf := make([]byte, len(signature))
-	_, err := file.ReadAt(buf, 0)
-	if err != nil {
-		return false, err
+// InitDB sets up the database connection, migrates models, and runs seeders.
+func InitDB(dbPath string) error {
+	// ensure dir exists
+	dir := path.Dir(dbPath)
+	if err := os.MkdirAll(dir, fs.ModePerm); err != nil {
+		return err
 	}
-	return bytes.Equal(buf, signature), nil
-}
 
-// Checkpoint performs a WAL checkpoint on the SQLite database to ensure data consistency.
-func Checkpoint() error {
-	// Update WAL
-	err := db.Exec("PRAGMA wal_checkpoint;").Error
+	// open SQLite (dev)
+	database, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{})
 	if err != nil {
 		return err
 	}
+	db = database
+
+	// migrations
+	if err := AutoMigrate(); err != nil {
+		return err
+	}
+
+	// seed admin
+	if err := SeedAdmin(); err != nil {
+		return err
+	}
+
 	return nil
 }
+
+// AutoMigrate applies schema migrations.
+func AutoMigrate() error {
+	return db.AutoMigrate(
+		&model.User{}, // User{ Id, Username, PasswordHash, Role }
+	)
+}
+
+// SeedAdmin creates a default admin if it doesn't exist.
+func SeedAdmin() error {
+	var count int64
+	if err := db.Model(&model.User{}).
+		Where("username = ?", "admin@local.test").
+		Count(&count).Error; err != nil {
+		return err
+	}
+	if count > 0 {
+		return nil
+	}
+
+	hash, _ := bcrypt.GenerateFromPassword([]byte("Admin12345!"), 12)
+	admin := model.User{
+		Username:     "admin@local.test",
+		PasswordHash: string(hash),
+		Role:         "admin",
+	}
+	return db.Create(&admin).Error
+}