Просмотр исходного кода

feat(server): sniff sqlite panel restore uploads and keep the fallback on failure

The SQLite panel's Restore now detects the upload by content like the
PostgreSQL panel does: migration dumps are rebuilt with RestoreSQLite,
pg_dump archives get a clear error instead of 'Invalid db file format',
and every upload passes the panel-schema pre-flight before Xray stops.
The .backup fallback survives a failed Xray start and is named in the
error, the DB pool is reopened on every error path after CloseDB, and a
failed InitDB closes the imported file before restoring the fallback so
the rename cannot hit a Windows sharing violation.
MHSanaei 17 часов назад
Родитель
Сommit
77dffe9a85

+ 86 - 65
internal/web/service/server.go

@@ -1426,44 +1426,26 @@ func (s *ServerService) ImportDB(file multipart.File) error {
 	if database.IsPostgres() {
 		return s.importPostgresDB(file)
 	}
-	// Check if the file is a SQLite database
-	isValidDb, err := database.IsSQLiteDB(file)
+	kind, err := sniffUploadKind(file)
 	if err != nil {
-		return common.NewErrorf("Error checking db file format: %v", err)
+		return common.NewErrorf("Error reading uploaded file: %v", err)
 	}
-	if !isValidDb {
-		return common.NewError("Invalid db file format")
-	}
-
-	// Reset the file reader to the beginning
-	_, err = file.Seek(0, 0)
-	if err != nil {
-		return common.NewErrorf("Error resetting file reader: %v", err)
+	switch kind {
+	case importKindSQLiteDB, importKindSQLiteDump:
+	case importKindPgDump:
+		return common.NewError("This file is a PostgreSQL backup; it can only be restored on a panel running PostgreSQL")
+	default:
+		return common.NewError("Invalid file: expected a SQLite database (.db) from Back Up or a SQLite migration dump (.dump)")
 	}
 
-	// Save the file as a temporary file
 	tempPath := fmt.Sprintf("%s.temp", config.GetDBPath())
 
-	// Remove the existing temporary file (if any)
 	if _, err := os.Stat(tempPath); err == nil {
 		if errRemove := os.Remove(tempPath); 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)
-	}
-
-	// Robust deferred cleanup for the temporary file
 	defer func() {
-		if tempFile != nil {
-			if cerr := tempFile.Close(); cerr != nil {
-				logger.Warningf("Warning: failed to close temp file: %v", cerr)
-			}
-		}
 		if _, err := os.Stat(tempPath); err == nil {
 			if rerr := os.Remove(tempPath); rerr != nil {
 				logger.Warningf("Warning: failed to remove temp file: %v", rerr)
@@ -1471,21 +1453,16 @@ func (s *ServerService) ImportDB(file multipart.File) error {
 		}
 	}()
 
-	// Save uploaded file to temporary file
-	if _, err = io.Copy(tempFile, file); err != nil {
-		return common.NewErrorf("Error saving db: %v", err)
-	}
-
-	// Close temp file before opening via sqlite
-	if err = tempFile.Close(); err != nil {
-		return common.NewErrorf("Error closing temporary db file: %v", err)
+	if err := stageSQLiteUpload(file, kind, tempPath); err != nil {
+		return err
 	}
-	tempFile = nil
 
-	// Validate integrity (no migrations / side effects)
 	if err = database.ValidateSQLiteDB(tempPath); err != nil {
 		return common.NewErrorf("Invalid or corrupt db file: %v", err)
 	}
+	if err = database.PrepareSQLiteForMigration(tempPath); err != nil {
+		return common.NewErrorf("This file cannot be imported: %v", err)
+	}
 
 	xrayStopped := true
 	defer func() {
@@ -1503,6 +1480,19 @@ func (s *ServerService) ImportDB(file multipart.File) error {
 		logger.Warningf("Failed to close existing DB before replacement: %v", errClose)
 	}
 
+	// Registered after the xray-restart defer so it runs first (LIFO): every
+	// error return below leaves a database file at the configured path, and the
+	// restart needs an open pool to build the xray config from it.
+	dbReopened := false
+	defer func() {
+		if dbReopened {
+			return
+		}
+		if errReopen := database.InitDB(config.GetDBPath()); errReopen != nil {
+			logger.Warningf("Failed to reopen the database after import error: %v", errReopen)
+		}
+	}()
+
 	// Backup the current database for fallback
 	fallbackPath := fmt.Sprintf("%s.backup", config.GetDBPath())
 
@@ -1518,15 +1508,6 @@ func (s *ServerService) ImportDB(file multipart.File) error {
 		return common.NewErrorf("Error backing up current db file: %v", err)
 	}
 
-	// Defer fallback cleanup ONLY if everything goes well
-	defer func() {
-		if _, err := os.Stat(fallbackPath); err == nil {
-			if rerr := os.Remove(fallbackPath); rerr != nil {
-				logger.Warningf("Warning: failed to remove fallback file: %v", rerr)
-			}
-		}
-	}()
-
 	// Move temp to DB path
 	if err = os.Rename(tempPath, config.GetDBPath()); err != nil {
 		// Restore from fallback
@@ -1538,19 +1519,30 @@ func (s *ServerService) ImportDB(file multipart.File) error {
 
 	// Open & migrate new DB
 	if err = database.InitDB(config.GetDBPath()); err != nil {
+		// A failed InitDB still holds the imported file open; close before the
+		// rename or Windows refuses to replace it.
+		if errClose := database.CloseDB(); errClose != nil {
+			logger.Warningf("Failed to close the imported DB before restoring fallback: %v", errClose)
+		}
 		if errRename := os.Rename(fallbackPath, config.GetDBPath()); errRename != nil {
 			return common.NewErrorf("Error migrating db and restoring fallback: %v", errRename)
 		}
 		return common.NewErrorf("Error migrating db: %v", err)
 	}
+	dbReopened = true
 
 	s.inboundService.MigrateDB()
 
 	xrayStopped = false
 	if err = s.RestartXrayService(); err != nil {
-		return common.NewErrorf("Imported DB but failed to start Xray: %v", err)
+		return common.NewErrorf("Imported DB but failed to start Xray: %v; the previous database was kept at %s", err, fallbackPath)
 	}
 
+	if _, err := os.Stat(fallbackPath); err == nil {
+		if rerr := os.Remove(fallbackPath); rerr != nil {
+			logger.Warningf("Warning: failed to remove fallback file: %v", rerr)
+		}
+	}
 	return nil
 }
 
@@ -1659,46 +1651,54 @@ func parsePgToolVersion(versionOutput string) string {
 }
 
 const (
-	pgImportUnknown = iota
-	pgImportPgDump
-	pgImportSQLiteDB
-	pgImportSQLiteDump
+	importKindUnknown = iota
+	importKindPgDump
+	importKindSQLiteDB
+	importKindSQLiteDump
 )
 
-// sniffPgImportKind classifies an uploaded restore file by its leading bytes:
+// sniffImportKind classifies an uploaded restore file by its leading bytes:
 // a pg_dump custom archive, a raw SQLite database, or a SQLite SQL text dump.
-func sniffPgImportKind(header []byte) int {
+func sniffImportKind(header []byte) int {
 	if bytes.HasPrefix(header, []byte("PGDMP")) {
-		return pgImportPgDump
+		return importKindPgDump
 	}
 	if bytes.HasPrefix(header, []byte("SQLite format 3\x00")) {
-		return pgImportSQLiteDB
+		return importKindSQLiteDB
 	}
 	text := bytes.TrimLeft(bytes.TrimPrefix(header, []byte("\xef\xbb\xbf")), " \t\r\n")
 	if bytes.HasPrefix(text, []byte("PRAGMA")) || bytes.HasPrefix(text, []byte("BEGIN TRANSACTION")) {
-		return pgImportSQLiteDump
+		return importKindSQLiteDump
 	}
-	return pgImportUnknown
+	return importKindUnknown
 }
 
-func (s *ServerService) importPostgresDB(file multipart.File) error {
+func sniffUploadKind(file multipart.File) (int, error) {
 	header := make([]byte, 64)
 	n, err := file.ReadAt(header, 0)
 	if err != nil && !errors.Is(err, io.EOF) {
-		return common.NewErrorf("Error reading dump file: %v", err)
+		return importKindUnknown, err
 	}
 	if _, err := file.Seek(0, 0); err != nil {
-		return common.NewErrorf("Error resetting file reader: %v", err)
+		return importKindUnknown, err
+	}
+	return sniffImportKind(header[:n]), nil
+}
+
+func (s *ServerService) importPostgresDB(file multipart.File) error {
+	kind, err := sniffUploadKind(file)
+	if err != nil {
+		return common.NewErrorf("Error reading uploaded file: %v", err)
 	}
-	switch sniffPgImportKind(header[:n]) {
-	case pgImportPgDump:
+	switch kind {
+	case importKindPgDump:
 		return s.restorePostgresDump(file)
-	case pgImportSQLiteDB:
+	case importKindSQLiteDB:
 		return s.migrateSQLiteIntoPostgres(file, false)
-	case pgImportSQLiteDump:
+	case importKindSQLiteDump:
 		return s.migrateSQLiteIntoPostgres(file, true)
 	default:
-		return common.NewError("Invalid file: expected a PostgreSQL custom-format dump (.dump) from this panel's Back Up, a SQLite database (.db), or a SQLite migration dump from Download Migration")
+		return common.NewError("Invalid file: expected a PostgreSQL custom-format dump (.dump) from this panel's Back Up, a SQLite database (.db), or a SQLite migration dump")
 	}
 }
 
@@ -1796,6 +1796,9 @@ func (s *ServerService) migrateSQLiteIntoPostgres(file multipart.File, isSQLDump
 	if err := database.ValidateSQLiteDB(dbPath); err != nil {
 		return common.NewErrorf("Invalid or corrupt db file: %v", err)
 	}
+	if err := database.PrepareSQLiteForMigration(dbPath); err != nil {
+		return common.NewErrorf("This file cannot be imported: %v", err)
+	}
 
 	xrayStopped := true
 	defer func() {
@@ -1821,7 +1824,7 @@ func (s *ServerService) migrateSQLiteIntoPostgres(file multipart.File, isSQLDump
 	s.inboundService.MigrateDB()
 
 	if migrateErr != nil {
-		return common.NewErrorf("Importing the SQLite data into PostgreSQL failed: %v; the destination tables are cleared on every attempt, so fixing the issue and retrying is safe", migrateErr)
+		return common.NewErrorf("Importing the SQLite data into PostgreSQL failed: %v; the import runs in a single transaction, so the database was left unchanged", migrateErr)
 	}
 
 	xrayStopped = false
@@ -1843,6 +1846,24 @@ func saveUploadedFile(file multipart.File, dstPath string) error {
 	return dst.Close()
 }
 
+func stageSQLiteUpload(file multipart.File, kind int, tempPath string) error {
+	if kind == importKindSQLiteDump {
+		dumpPath := tempPath + ".dump"
+		defer os.Remove(dumpPath)
+		if err := saveUploadedFile(file, dumpPath); err != nil {
+			return common.NewErrorf("Error saving migration dump: %v", err)
+		}
+		if err := database.RestoreSQLite(dumpPath, tempPath); err != nil {
+			return common.NewErrorf("Error rebuilding a SQLite database from the migration dump: %v", err)
+		}
+		return nil
+	}
+	if err := saveUploadedFile(file, tempPath); err != nil {
+		return common.NewErrorf("Error saving db: %v", err)
+	}
+	return nil
+}
+
 // IsValidGeofileName validates that the filename is safe for geofile operations.
 // It checks for path traversal attempts and ensures the filename contains only safe characters.
 func (s *ServerService) IsValidGeofileName(filename string) bool {

+ 11 - 11
internal/web/service/server_pg_import_sniff_test.go → internal/web/service/server_import_sniff_test.go

@@ -7,23 +7,23 @@ import (
 	"github.com/mhsanaei/3x-ui/v3/internal/database"
 )
 
-func TestSniffPgImportKind(t *testing.T) {
+func TestSniffImportKind(t *testing.T) {
 	cases := []struct {
 		name   string
 		header []byte
 		want   int
 	}{
-		{"pg custom archive", []byte("PGDMP\x01\x10\x04"), pgImportPgDump},
-		{"raw sqlite database", []byte("SQLite format 3\x00rest of header"), pgImportSQLiteDB},
-		{"sqlite cli dump without pragma", []byte("BEGIN TRANSACTION;\nCREATE TABLE t(i);"), pgImportSQLiteDump},
-		{"bom and whitespace before pragma", []byte("\xef\xbb\xbf\r\n PRAGMA foreign_keys=OFF;"), pgImportSQLiteDump},
-		{"plain-format postgres dump", []byte("--\n-- PostgreSQL database dump\n--"), pgImportUnknown},
-		{"empty file", nil, pgImportUnknown},
+		{"pg custom archive", []byte("PGDMP\x01\x10\x04"), importKindPgDump},
+		{"raw sqlite database", []byte("SQLite format 3\x00rest of header"), importKindSQLiteDB},
+		{"sqlite cli dump without pragma", []byte("BEGIN TRANSACTION;\nCREATE TABLE t(i);"), importKindSQLiteDump},
+		{"bom and whitespace before pragma", []byte("\xef\xbb\xbf\r\n PRAGMA foreign_keys=OFF;"), importKindSQLiteDump},
+		{"plain-format postgres dump", []byte("--\n-- PostgreSQL database dump\n--"), importKindUnknown},
+		{"empty file", nil, importKindUnknown},
 	}
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			if got := sniffPgImportKind(tc.header); got != tc.want {
-				t.Errorf("sniffPgImportKind(%q) = %d, want %d", tc.header, got, tc.want)
+			if got := sniffImportKind(tc.header); got != tc.want {
+				t.Errorf("sniffImportKind(%q) = %d, want %d", tc.header, got, tc.want)
 			}
 		})
 	}
@@ -38,8 +38,8 @@ func TestSniffPgImportKind(t *testing.T) {
 		if err != nil {
 			t.Fatalf("DumpSQLiteToBytes: %v", err)
 		}
-		if got := sniffPgImportKind(dump[:64]); got != pgImportSQLiteDump {
-			t.Errorf("sniffPgImportKind(real migration dump) = %d, want %d", got, pgImportSQLiteDump)
+		if got := sniffImportKind(dump[:64]); got != importKindSQLiteDump {
+			t.Errorf("sniffImportKind(real migration dump) = %d, want %d", got, importKindSQLiteDump)
 		}
 	})
 }

+ 46 - 0
internal/web/service/server_import_stage_test.go

@@ -0,0 +1,46 @@
+package service
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/mhsanaei/3x-ui/v3/internal/database"
+)
+
+func TestStageSQLiteUploadRebuildsFromDump(t *testing.T) {
+	dir := t.TempDir()
+	dbPath := filepath.Join(dir, "x-ui.db")
+	if err := database.InitDB(dbPath); err != nil {
+		t.Fatalf("InitDB: %v", err)
+	}
+	t.Cleanup(func() { _ = database.CloseDB() })
+	dump, err := database.DumpSQLiteToBytes(dbPath)
+	if err != nil {
+		t.Fatalf("DumpSQLiteToBytes: %v", err)
+	}
+
+	uploadPath := filepath.Join(dir, "upload.dump")
+	if err := os.WriteFile(uploadPath, dump, 0o644); err != nil {
+		t.Fatalf("write upload: %v", err)
+	}
+	upload, err := os.Open(uploadPath)
+	if err != nil {
+		t.Fatalf("open upload: %v", err)
+	}
+	defer upload.Close()
+
+	staged := filepath.Join(dir, "x-ui.db.temp")
+	if err := stageSQLiteUpload(upload, importKindSQLiteDump, staged); err != nil {
+		t.Fatalf("stageSQLiteUpload: %v", err)
+	}
+	if _, err := os.Stat(staged + ".dump"); !os.IsNotExist(err) {
+		t.Errorf("intermediate dump file %s.dump was not cleaned up", staged)
+	}
+	if err := database.ValidateSQLiteDB(staged); err != nil {
+		t.Errorf("staged database fails integrity check: %v", err)
+	}
+	if err := database.PrepareSQLiteForMigration(staged); err != nil {
+		t.Errorf("staged database fails the panel pre-flight: %v", err)
+	}
+}