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

fix(database): keep the dump restore inside its own database file

A SQL dump replay only has to rebuild the tables of the database it
restores into. Run it on a single connection whose attached-database
limit is zero, so the script cannot open or create any other file.
MHSanaei 4 часов назад
Родитель
Сommit
d31465e37b
2 измененных файлов с 61 добавлено и 2 удалено
  1. 27 2
      internal/database/dump_sqlite.go
  2. 34 0
      internal/database/dump_sqlite_test.go

+ 27 - 2
internal/database/dump_sqlite.go

@@ -9,6 +9,7 @@ import (
 	"strings"
 	"unicode/utf8"
 
+	"github.com/mattn/go-sqlite3"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
 	"gorm.io/gorm/logger"
@@ -127,8 +128,7 @@ func RestoreSQLite(dumpPath, dstPath string) error {
 		return err
 	}
 
-	// mattn/go-sqlite3 executes every statement in a multi-statement string.
-	if _, err := sqlDB.ExecContext(context.Background(), string(script)); err != nil {
+	if err := replayDump(sqlDB, string(script)); err != nil {
 		sqlDB.Close()
 		os.Remove(dstPath)
 		return fmt.Errorf("restore failed: %w", err)
@@ -136,6 +136,31 @@ func RestoreSQLite(dumpPath, dstPath string) error {
 	return sqlDB.Close()
 }
 
+// replayDump runs the script on one connection that cannot open a second database
+// file: ATTACH and VACUUM INTO both attach, and a dump only rebuilds its own tables.
+func replayDump(sqlDB *sql.DB, script string) error {
+	ctx := context.Background()
+	conn, err := sqlDB.Conn(ctx)
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+	err = conn.Raw(func(driverConn any) error {
+		sc, ok := driverConn.(*sqlite3.SQLiteConn)
+		if !ok {
+			return fmt.Errorf("unexpected sqlite driver connection %T", driverConn)
+		}
+		sc.SetLimit(sqlite3.SQLITE_LIMIT_ATTACHED, 0)
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+	// mattn/go-sqlite3 executes every statement in a multi-statement string.
+	_, err = conn.ExecContext(ctx, script)
+	return err
+}
+
 // dumpTableData appends one INSERT statement per row of table to b.
 func dumpTableData(db *sql.DB, table string, b *strings.Builder) error {
 	rows, err := db.QueryContext(context.Background(), `SELECT * FROM "`+table+`"`)

+ 34 - 0
internal/database/dump_sqlite_test.go

@@ -1,8 +1,10 @@
 package database
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
@@ -135,3 +137,35 @@ func closeGorm(db *gorm.DB) {
 		s.Close()
 	}
 }
+
+func TestRestoreSQLiteWritesNoFileOutsideDestination(t *testing.T) {
+	cases := []struct {
+		name      string
+		statement string
+	}{
+		{"attach", "ATTACH DATABASE '%s' AS x; CREATE TABLE x.t (a); INSERT INTO x.t VALUES (1);"},
+		{"vacuum into", "CREATE TABLE t (a); VACUUM INTO '%s';"},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			dir := t.TempDir()
+			outside := filepath.Join(t.TempDir(), "outside.db")
+			dumpPath := filepath.Join(dir, "in.dump")
+			script := "PRAGMA foreign_keys=OFF;\n" + fmt.Sprintf(tc.statement, outside) + "\n"
+			if err := os.WriteFile(dumpPath, []byte(script), 0o600); err != nil {
+				t.Fatalf("write dump: %v", err)
+			}
+
+			err := RestoreSQLite(dumpPath, filepath.Join(dir, "rebuilt.db"))
+			if err == nil {
+				t.Fatal("RestoreSQLite accepted a dump that writes another database file")
+			}
+			if !strings.Contains(err.Error(), "too many attached databases") {
+				t.Fatalf("restore failed for another reason: %v", err)
+			}
+			if _, statErr := os.Stat(outside); !os.IsNotExist(statErr) {
+				t.Fatalf("restore created %s (stat err %v)", outside, statErr)
+			}
+		})
+	}
+}