Merge pull request #479 from varbhat/fix/admin_user_init

fix: admin user not init after switch database
varbhat · May 21, 2024 · 5f6f2e7 · 5f6f2e7
2 parents 98c4634 + 7dad61f
commit 5f6f2e7
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 21 deletions.
diff --git a/internal/core/engine.go b/internal/core/engine.go
@@ -611,3 +611,7 @@ func AddTrackerstoTorrent(User string, infohash metainfo.Hash, announcelist [][]
 	Info.Println("Trackers added to torrent ", infohash, "by ", User)
 	MainHub.SendMsgU(User, "resp", infohash.HexString(), "success", "Trackers added")
 }
+
+func CheckUserExists(username string) bool {
+	return Engine.UDb.CheckUserExists(username)
+}
diff --git a/internal/core/init.go b/internal/core/init.go
@@ -250,28 +250,21 @@ func Initialize() {
 		sqliteSetup(tc)
 	}
 
-	_, err = os.Stat(filepath.Join(Dirconfig.DataDir, ".adminadded"))
-	if errors.Is(err, os.ErrNotExist) {
+	if !CheckUserExists(auser) {
+		var password = "adminpassword"
 		if pw {
-			Info.Println(`Adding Admin user with username "` + auser + `" and custom password`)
-			er := Engine.UDb.Add(auser, os.Getenv("EXAPASSWORD"), 1)
-			if er != nil {
-				Err.Fatalln("Unable to add admin user to adminless exatorrent instance :", er)
-			}
-			_, er = os.Create(filepath.Join(Dirconfig.DataDir, ".adminadded"))
-			if er != nil {
-				Err.Fatalln(er)
-			}
-		} else {
-			Info.Println(`Adding Admin user with username "` + auser + `" and password "adminpassword"`)
-			er := Engine.UDb.Add(auser, "adminpassword", 1)
-			if er != nil {
-				Err.Fatalln("Unable to add admin user to adminless exatorrent instance :", er)
-			}
-			_, er = os.Create(filepath.Join(Dirconfig.DataDir, ".adminadded"))
-			if er != nil {
-				Err.Fatalln(er)
-			}
+			password = os.Getenv("EXAPASSWORD")
+		}
+
+		Info.Printf("Adding Admin user with username %s and password %s\n", auser, password)
+		er := Engine.UDb.Add(auser, password, 1)
+		if er != nil {
+			Err.Fatalln("Unable to add admin user to adminless exatorrent instance :", er)
+		}
+		// keep for backward compatibility
+		_, er = os.Create(filepath.Join(Dirconfig.DataDir, ".adminadded"))
+		if er != nil {
+			Err.Fatalln(er)
 		}
 	}
 

diff --git a/internal/db/db.go b/internal/db/db.go
@@ -79,6 +79,7 @@ type UserDb interface {
 	Validate(string, string) (int, bool)
 	ValidateToken(string) (string, int, error)
 	SetToken(string, string) error
+	CheckUserExists(string) bool
 }
 
 type TorrentUserDb interface {

diff --git a/internal/db/psqluserdb.go b/internal/db/psqluserdb.go
@@ -137,3 +137,15 @@ func (db *PsqlUserDb) ValidateToken(Token string) (user string, ut int, err erro
 	}
 	return
 }
+
+func (db *PsqlUserDb) CheckUserExists(username string) bool {
+	var exists bool
+	var err = db.Db.
+		QueryRow(context.Background(), `select exists(select * from userdb where username = $1);`, username).
+		Scan(&exists)
+	if err != nil {
+		DbL.Printf("fail to check username exists: %s, err: %v", username, err)
+		return false
+	}
+	return exists
+}
diff --git a/internal/db/sqlite3userdb.go b/internal/db/sqlite3userdb.go
@@ -189,3 +189,19 @@ func (db *Sqlite3UserDb) SetToken(Username string, Token string) (err error) {
 	err = sqlitex.Exec(db.Db, `update userdb set token=? where username=?;`, nil, Token, Username)
 	return
 }
+
+func (db *Sqlite3UserDb) CheckUserExists(username string) bool {
+	var exists bool
+	var err = sqlitex.Exec(
+		db.Db,
+		`select exists(select * from userdb where username = ?);`,
+		func(stmt *sqlite.Stmt) error {
+			exists = stmt.ColumnInt(0) != 0
+			return nil
+		}, username)
+	if err != nil {
+		DbL.Printf("fail to check username exists: %s, err: %v", username, err)
+		return false
+	}
+	return exists
+}