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

This commit is contained in:
Dikiy13371 2025-10-07 23:39:24 +03:00
parent d457c5a9d0
commit f0eca194e2

View file

@ -1,198 +1,73 @@
// Package database provides database initialization, migration, and management utilities
// for the 3x-ui panel using GORM with SQLite.
package database package database
import ( import (
"bytes"
"io"
"io/fs" "io/fs"
"log"
"os" "os"
"path" "path"
"slices"
"github.com/mhsanaei/3x-ui/v2/config"
"github.com/mhsanaei/3x-ui/v2/database/model" "github.com/mhsanaei/3x-ui/v2/database/model"
"github.com/mhsanaei/3x-ui/v2/util/crypto" "golang.org/x/crypto/bcrypt"
"github.com/mhsanaei/3x-ui/v2/xray"
"gorm.io/driver/sqlite" "gorm.io/driver/sqlite"
"gorm.io/gorm" "gorm.io/gorm"
"gorm.io/gorm/logger"
) )
var db *gorm.DB 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. // GetDB returns the global GORM database instance.
func GetDB() *gorm.DB { return db } func GetDB() *gorm.DB { return db }
// IsNotFound checks if the given error is a GORM record not found error. // InitDB sets up the database connection, migrates models, and runs seeders.
func IsNotFound(err error) bool { func InitDB(dbPath string) error {
return err == gorm.ErrRecordNotFound // ensure dir exists
} dir := path.Dir(dbPath)
if err := os.MkdirAll(dir, fs.ModePerm); err != nil {
// IsSQLiteDB checks if the given file is a valid SQLite database by reading its signature. return err
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
} }
return bytes.Equal(buf, signature), nil
}
// Checkpoint performs a WAL checkpoint on the SQLite database to ensure data consistency. // open SQLite (dev)
func Checkpoint() error { database, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{})
// Update WAL
err := db.Exec("PRAGMA wal_checkpoint;").Error
if err != nil { if err != nil {
return err return err
} }
db = database
// migrations
if err := AutoMigrate(); err != nil {
return err
}
// seed admin
if err := SeedAdmin(); err != nil {
return err
}
return nil 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
}