Bläddra i källkod

fix(migration): stop a half-applied startup migration from committing silently (#6182)

* fix(traffic): check maintenance commits and IP-limit errors

* fix(migrations): propagate transactional failures

---------

Co-authored-by: n0ctal <[email protected]>
n0ctal 18 timmar sedan
förälder
incheckning
b56b087254

+ 43 - 23
internal/web/service/inbound_migration.go

@@ -33,13 +33,15 @@ func (s *InboundService) MigrationRemoveOrphanedTraffics() {
 	}
 }
 
-func (s *InboundService) MigrationRequirements() {
+func (s *InboundService) MigrationRequirements() (err error) {
 	db := database.GetDB()
 	tx := db.Begin()
-	var err error
 	defer func() {
 		if err == nil {
-			tx.Commit()
+			if commitErr := tx.Commit().Error; commitErr != nil {
+				err = commitErr
+				return
+			}
 			if !database.IsPostgres() {
 				if dbErr := db.Exec(`VACUUM "main"`).Error; dbErr != nil {
 					logger.Warningf("VACUUM failed: %v", dbErr)
@@ -76,8 +78,8 @@ func (s *InboundService) MigrationRequirements() {
 	// SQLite (no PG :: casts).
 	if database.IsPostgres() {
 		// Use DO block so it is idempotent and doesn't fail if already boolean.
-		normalizeBool := func(table, col string) {
-			tx.Exec(fmt.Sprintf(`
+		normalizeBool := func(table, col string) error {
+			return tx.Exec(fmt.Sprintf(`
 				DO $$
 				BEGIN
 					IF EXISTS (
@@ -88,14 +90,13 @@ func (s *InboundService) MigrationRequirements() {
 						ALTER TABLE %s ALTER COLUMN %s
 							TYPE boolean USING (CASE WHEN %s::text IN ('1','true','t','yes') THEN true ELSE false END);
 					END IF;
-				END $$;`, table, col, table, col, col))
+				END $$;`, table, col, table, col, col)).Error
+		}
+		for _, column := range [][2]string{{"inbounds", "enable"}, {"client_traffics", "enable"}, {"nodes", "enable"}, {"clients", "enable"}, {"api_tokens", "enabled"}, {"outbound_subscriptions", "enabled"}} {
+			if err = normalizeBool(column[0], column[1]); err != nil {
+				return
+			}
 		}
-		normalizeBool("inbounds", "enable")
-		normalizeBool("client_traffics", "enable")
-		normalizeBool("nodes", "enable")
-		normalizeBool("clients", "enable")
-		normalizeBool("api_tokens", "enabled")
-		normalizeBool("outbound_subscriptions", "enabled")
 	}
 
 	// Fix inbounds based problems
@@ -160,7 +161,8 @@ func (s *InboundService) MigrationRequirements() {
 				delete(settings, "testseed")
 			}
 
-			modifiedSettings, err := json.MarshalIndent(settings, "", "  ")
+			var modifiedSettings []byte
+			modifiedSettings, err = json.MarshalIndent(settings, "", "  ")
 			if err != nil {
 				return
 			}
@@ -169,30 +171,39 @@ func (s *InboundService) MigrationRequirements() {
 		}
 
 		// Add client traffic row for all clients which has email
-		modelClients, err := s.GetClients(inbounds[inbound_index])
+		var modelClients []model.Client
+		modelClients, err = s.GetClients(inbounds[inbound_index])
 		if err != nil {
 			return
 		}
 		for _, modelClient := range modelClients {
 			if len(modelClient.Email) > 0 {
 				var count int64
-				tx.Model(xray.ClientTraffic{}).Where("email = ?", modelClient.Email).Count(&count)
+				if err = tx.Model(xray.ClientTraffic{}).Where("email = ?", modelClient.Email).Count(&count).Error; err != nil {
+					return
+				}
 				if count == 0 {
-					_ = s.AddClientStat(tx, inbounds[inbound_index].Id, &modelClient)
+					if err = s.AddClientStat(tx, inbounds[inbound_index].Id, &modelClient); err != nil {
+						return
+					}
 				}
 			}
 		}
 
 		// Heal clients table for installs where the one-shot seeder
 		// skipped clients due to a tgId-string unmarshal error.
-		if syncErr := s.clientService.SyncInbound(tx, inbounds[inbound_index].Id, modelClients); syncErr != nil {
-			logger.Warning("MigrationRequirements sync clients failed:", syncErr)
+		if err = s.clientService.SyncInbound(tx, inbounds[inbound_index].Id, modelClients); err != nil {
+			return
 		}
 	}
-	tx.Save(inbounds)
+	if err = tx.Save(inbounds).Error; err != nil {
+		return
+	}
 
 	// Remove orphaned traffics
-	tx.Where("inbound_id = 0").Delete(xray.ClientTraffic{})
+	if err = tx.Where("inbound_id = 0").Delete(xray.ClientTraffic{}).Error; err != nil {
+		return
+	}
 
 	// Migrate old MultiDomain to External Proxy
 	var externalProxy []struct {
@@ -238,8 +249,14 @@ func (s *InboundService) MigrationRequirements() {
 			}
 		}
 		stream["externalProxy"] = reverses
-		newStream, _ := json.MarshalIndent(stream, " ", "  ")
-		tx.Model(model.Inbound{}).Where("id = ?", ep.Id).Update("stream_settings", newStream)
+		newStream, marshalErr := json.MarshalIndent(stream, " ", "  ")
+		if marshalErr != nil {
+			err = marshalErr
+			return
+		}
+		if err = tx.Model(model.Inbound{}).Where("id = ?", ep.Id).Update("stream_settings", newStream).Error; err != nil {
+			return
+		}
 	}
 
 	// Legacy tag cleanup for old auto-generated tags (e.g. "0.0.0.0:443-...").
@@ -256,10 +273,13 @@ func (s *InboundService) MigrationRequirements() {
 	if err != nil {
 		return
 	}
+	return err
 }
 
 func (s *InboundService) MigrateDB() {
-	s.MigrationRequirements()
+	if err := s.MigrationRequirements(); err != nil {
+		logger.Errorf("MigrationRequirements failed: %v", err)
+	}
 	s.MigrationRemoveOrphanedTraffics()
 	s.MigrationRestoreVisionFlow()
 }

+ 38 - 0
internal/web/service/inbound_migration_test.go

@@ -1,10 +1,13 @@
 package service
 
 import (
+	"errors"
 	"path/filepath"
 	"strings"
 	"testing"
 
+	"gorm.io/gorm"
+
 	"github.com/mhsanaei/3x-ui/v3/internal/database"
 	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
 	"github.com/mhsanaei/3x-ui/v3/internal/xray"
@@ -90,6 +93,41 @@ func TestMigrationRequirements_BackfillsClientTrafficsWithMultiDomainInbound(t *
 	}
 }
 
+func TestMigrationRequirementsReturnsAddClientStatFailure(t *testing.T) {
+	dbDir := t.TempDir()
+	t.Setenv("XUI_DB_FOLDER", dbDir)
+	if err := database.InitDB(filepath.Join(dbDir, "x-ui.db")); err != nil {
+		t.Fatalf("InitDB: %v", err)
+	}
+	t.Cleanup(func() { _ = database.CloseDB() })
+	db := database.GetDB()
+	first := &model.Inbound{UserId: 1, Tag: "first", Port: 31001, Protocol: model.VLESS, Settings: `{"clients":[{"email":"[email protected]","id":"id-1"}]}`, StreamSettings: `{}`}
+	if err := db.Create(first).Error; err != nil {
+		t.Fatalf("create first: %v", err)
+	}
+	const injected = "injected AddClientStat failure"
+	failSave := func(tx *gorm.DB) {
+		tx.AddError(errors.New(injected))
+	}
+	if err := db.Callback().Update().Before("gorm:update").Register("test:fail-migration-inbound-save", failSave); err != nil {
+		t.Fatalf("register update callback: %v", err)
+	}
+	if err := db.Callback().Create().Before("gorm:create").Register("test:fail-migration-inbound-save", failSave); err != nil {
+		t.Fatalf("register create callback: %v", err)
+	}
+	err := (&InboundService{}).MigrationRequirements()
+	if err == nil || err.Error() != injected {
+		t.Fatalf("MigrationRequirements error = %v, want %q", err, injected)
+	}
+	var count int64
+	if err := db.Model(&xray.ClientTraffic{}).Where("email = ?", "[email protected]").Count(&count).Error; err != nil {
+		t.Fatalf("count rolled-back traffic: %v", err)
+	}
+	if count != 0 {
+		t.Fatalf("earlier traffic write committed after save failure: count=%d", count)
+	}
+}
+
 // TestMigrationRequirements_CleansLegacyZeroAddrTag guards the legacy tag cleanup that
 // strips the auto-generated "0.0.0.0:" prefix. The inbound is MultiDomain TLS so the
 // externalProxy detection query returns rows and the cleanup is reached (it early-returns

+ 4 - 18
internal/web/service/outbound/outbound.go

@@ -22,24 +22,10 @@ import (
 type OutboundService struct{}
 
 func (s *OutboundService) AddTraffic(traffics []*xray.Traffic, clientTraffics []*xray.ClientTraffic) (error, bool) {
-	var err error
-	db := database.GetDB()
-	tx := db.Begin()
-
-	defer func() {
-		if err != nil {
-			tx.Rollback()
-		} else {
-			tx.Commit()
-		}
-	}()
-
-	err = s.addOutboundTraffic(tx, traffics)
-	if err != nil {
-		return err, false
-	}
-
-	return nil, false
+	err := database.GetDB().Transaction(func(tx *gorm.DB) error {
+		return s.addOutboundTraffic(tx, traffics)
+	})
+	return err, false
 }
 
 // saturatingAdd caps counters at database.TrafficMax: unlike the SQL paths,

+ 54 - 0
internal/web/service/outbound/outbound_commit_postgres_test.go

@@ -0,0 +1,54 @@
+package outbound
+
+import (
+	"os"
+	"strings"
+	"testing"
+
+	"gorm.io/gorm"
+
+	"github.com/mhsanaei/3x-ui/v3/internal/database"
+	"github.com/mhsanaei/3x-ui/v3/internal/xray"
+)
+
+func TestAddTrafficReturnsDeferredCommitFailure(t *testing.T) {
+	if os.Getenv("XUI_DB_TYPE") != "postgres" || strings.TrimSpace(os.Getenv("XUI_DB_DSN")) == "" {
+		t.Skip("set XUI_DB_TYPE=postgres and XUI_DB_DSN to run commit-failure injection")
+	}
+	if err := database.InitDB(""); err != nil {
+		t.Fatalf("InitDB: %v", err)
+	}
+	t.Cleanup(func() { _ = database.CloseDB() })
+	db := database.GetDB()
+	const parent = "outbound_commit_parent"
+	const child = "outbound_commit_child"
+	_ = db.Exec("DROP TABLE IF EXISTS " + child).Error
+	_ = db.Exec("DROP TABLE IF EXISTS " + parent).Error
+	if err := db.Exec("CREATE TABLE " + parent + " (id bigint PRIMARY KEY)").Error; err != nil {
+		t.Fatal(err)
+	}
+	if err := db.Exec("CREATE TABLE " + child + " (id bigint PRIMARY KEY, parent_id bigint REFERENCES " + parent + "(id) DEFERRABLE INITIALLY DEFERRED)").Error; err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		_ = db.Exec("DROP TABLE IF EXISTS " + child).Error
+		_ = db.Exec("DROP TABLE IF EXISTS " + parent).Error
+	})
+	const callback = "test:outbound-deferred-commit"
+	if err := db.Callback().Create().After("gorm:create").Register(callback, func(tx *gorm.DB) {
+		if tx.Statement == nil || tx.Statement.Table != "outbound_traffics" {
+			return
+		}
+		if result := tx.Session(&gorm.Session{NewDB: true}).Exec("INSERT INTO " + child + " (id, parent_id) VALUES (1, 999999)"); result.Error != nil {
+			tx.AddError(result.Error)
+		}
+	}); err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = db.Callback().Create().Remove(callback) })
+
+	err, _ := (&OutboundService{}).AddTraffic([]*xray.Traffic{{Tag: "commit-test", IsOutbound: true, Up: 1}}, nil)
+	if err == nil || !strings.Contains(strings.ToLower(err.Error()), "foreign key") {
+		t.Fatalf("AddTraffic error = %v, want deferred foreign-key commit failure", err)
+	}
+}