mirror of
https://github.com/MHSanaei/3x-ui.git
synced 2025-12-23 06:42:41 +00:00
fix(database): restore working InitDB/GetDB with automigrate and admin seed
This commit is contained in:
parent
d457c5a9d0
commit
f0eca194e2
1 changed files with 49 additions and 174 deletions
223
database/db.go
223
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
|
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
|
||||||
|
}
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue