From 83c853ffb6b896c8a6d1eef4e0354ba1201ebf13 Mon Sep 17 00:00:00 2001
From: Hamidreza Ghavami <hamid.r.gh.1998@gmail.com>
Date: Sat, 6 May 2023 04:47:57 +0430
Subject: [PATCH] update ImportDB and enhancement

---
 web/controller/server.go |  4 ++-
 web/service/server.go    | 60 ++++++++++++++++++++++++++++------------
 2 files changed, 46 insertions(+), 18 deletions(-)

diff --git a/web/controller/server.go b/web/controller/server.go
index b3dc0ad6..9e649e6c 100644
--- a/web/controller/server.go
+++ b/web/controller/server.go
@@ -154,13 +154,15 @@ func (a *ServerController) importDB(c *gin.Context) {
 	defer file.Close()
 	// Always restart Xray before return
 	defer a.serverService.RestartXrayService()
+	defer func() {
+		a.lastGetStatusTime = time.Now()
+	}()
 	// Import it
 	err = a.serverService.ImportDB(file)
 	if err != nil {
 		jsonMsg(c, "", err)
 		return
 	}
-	a.lastGetStatusTime = time.Now()
 	jsonObj(c, "Import DB", nil)
 }
 
diff --git a/web/service/server.go b/web/service/server.go
index a9e10ea2..d8a2239b 100644
--- a/web/service/server.go
+++ b/web/service/server.go
@@ -409,23 +409,33 @@ func (s *ServerService) ImportDB(file multipart.File) error {
 		return common.NewError("Invalid db file format")
 	}
 
-	// Save the file as temporary file
-	tempPath := fmt.Sprintf("%s.temp", config.GetDBPath())
-	// remove temp file before return
-	defer os.Remove(tempPath)
-	tempFile, err := os.Create(tempPath)
-	if err != nil {
-		return common.NewErrorf("Error creating temporary db file: %v", err)
-	}
-	defer tempFile.Close()
-
 	// Reset the file reader to the beginning
 	_, err = file.Seek(0, 0)
 	if err != nil {
 		return common.NewErrorf("Error resetting file reader: %v", err)
 	}
 
-	// Save temp file
+	// Save the file as temporary file
+	tempPath := fmt.Sprintf("%s.temp", config.GetDBPath())
+	// Remove the existing fallback file (if any) before creating one
+	_, err = os.Stat(tempPath)
+	if err == nil {
+		errRemove := os.Remove(tempPath)
+		if errRemove != nil {
+			return common.NewErrorf("Error removing existing temporary db file: %v", errRemove)
+		}
+	}
+	// Create the temporary file
+	tempFile, err := os.Create(tempPath)
+	if err != nil {
+		return common.NewErrorf("Error creating temporary db file: %v", err)
+	}
+	defer tempFile.Close()
+
+	// Remove temp file before returning
+	defer os.Remove(tempPath)
+
+	// Save uploaded file to temporary file
 	_, err = io.Copy(tempFile, file)
 	if err != nil {
 		return common.NewErrorf("Error saving db: %v", err)
@@ -440,26 +450,42 @@ func (s *ServerService) ImportDB(file multipart.File) error {
 	// Stop Xray
 	s.StopXrayService()
 
-	// Backup db for fallback
+	// Backup the current database for fallback
 	fallbackPath := fmt.Sprintf("%s.backup", config.GetDBPath())
-	// remove fallback file before return
-	defer os.Remove(fallbackPath)
+	// Remove the existing fallback file (if any)
+	_, err = os.Stat(fallbackPath)
+	if err == nil {
+		errRemove := os.Remove(fallbackPath)
+		if errRemove != nil {
+			return common.NewErrorf("Error removing existing fallback db file: %v", errRemove)
+		}
+	}
+	// Move the current database to the fallback location
 	err = os.Rename(config.GetDBPath(), fallbackPath)
 	if err != nil {
-		return common.NewErrorf("Error backup temporary db file: %v", err)
+		return common.NewErrorf("Error backing up temporary db file: %v", err)
 	}
 
+	// Remove the temporary file before returning
+	defer os.Remove(fallbackPath)
+
 	// Move temp to DB path
 	err = os.Rename(tempPath, config.GetDBPath())
 	if err != nil {
-		os.Rename(fallbackPath, config.GetDBPath())
+		errRename := os.Rename(fallbackPath, config.GetDBPath())
+		if errRename != nil {
+			return common.NewErrorf("Error moving db file and restoring fallback: %v", errRename)
+		}
 		return common.NewErrorf("Error moving db file: %v", err)
 	}
 
 	// Migrate DB
 	err = database.InitDB(config.GetDBPath())
 	if err != nil {
-		os.Rename(fallbackPath, config.GetDBPath())
+		errRename := os.Rename(fallbackPath, config.GetDBPath())
+		if errRename != nil {
+			return common.NewErrorf("Error migrating db and restoring fallback: %v", errRename)
+		}
 		return common.NewErrorf("Error migrating db: %v", err)
 	}
 	s.inboundService.MigrateDB()