3x-ui/internal/web/service/tx_pool_starvation_test.go

package service

import (
	"testing"
	"time"

	"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/web/runtime"
	"github.com/mhsanaei/3x-ui/v3/internal/xray"
)

// capPoolAtOneConn models a pool whose other connections are all parked on the
// transaction's row locks: any query outside tx then waits forever.
func capPoolAtOneConn(t *testing.T) {
	t.Helper()
	sqlDB, err := database.GetDB().DB()
	if err != nil {
		t.Fatalf("sql.DB: %v", err)
	}
	sqlDB.SetMaxOpenConns(1)
}

func runWithoutSecondConn(t *testing.T, name string, fn func() error) error {
	t.Helper()
	done := make(chan error, 1)
	go func() { done <- fn() }()
	select {
	case err := <-done:
		return err
	case <-time.After(5 * time.Second):
		// The DB cleanup's Close fails the parked pool request, which unwinds fn.
		t.Fatalf("%s blocked waiting for a second pool connection while its transaction is open", name)
		return nil
	}
}

func createNodeInboundWithSettings(t *testing.T, db *gorm.DB, nodeID int, tag string, port int, settings string) {
	t.Helper()
	nid := nodeID
	ib := &model.Inbound{UserId: 1, Tag: tag, Enable: true, Port: port, Protocol: model.VLESS, NodeID: &nid, Settings: settings}
	if err := db.Create(ib).Error; err != nil {
		t.Fatalf("create node inbound %q: %v", tag, err)
	}
}

func syncNodeSnapshot(t *testing.T, svc *InboundService, nodeID int, inbounds ...*model.Inbound) error {
	t.Helper()
	_, err := svc.setRemoteTrafficLocked(nodeID, &runtime.TrafficSnapshot{Inbounds: inbounds}, false, false)
	return err
}

func assertNoTraffic(t *testing.T, email string) {
	t.Helper()
	var n int64
	if err := database.GetDB().Model(xray.ClientTraffic{}).Where("email = ?", email).Count(&n).Error; err != nil {
		t.Fatalf("count traffic %q: %v", email, err)
	}
	if n != 0 {
		t.Fatalf("traffic rows for %q = %d, want 0", email, n)
	}
}

func TestNodeSyncInboundRemovalStaysOnOneConnection(t *testing.T) {
	db := initTrafficTestDB(t)
	goneSettings := `{"clients": [{"email": "gone", "enable": true}]}`
	keepSettings := `{"clients": [{"email": "keep", "enable": true}]}`
	createNodeInboundWithSettings(t, db, 1, "n1-gone", 41001, goneSettings)
	createNodeInboundWithSettings(t, db, 1, "n1-keep", 41002, keepSettings)
	svc := &InboundService{}
	keep := &model.Inbound{Tag: "n1-keep", Enable: true, Settings: keepSettings, ClientStats: []xray.ClientTraffic{{Email: "keep", Up: 10, Down: 10, Enable: true}}}
	if err := syncNodeSnapshot(t, svc, 1,
		&model.Inbound{Tag: "n1-gone", Enable: true, Settings: goneSettings, ClientStats: []xray.ClientTraffic{{Email: "gone", Up: 10, Down: 10, Enable: true}}},
		keep); err != nil {
		t.Fatalf("seed sync: %v", err)
	}
	readTraffic(t, db, "gone")

	capPoolAtOneConn(t)
	err := runWithoutSecondConn(t, "node sync inbound-removal sweep", func() error {
		return syncNodeSnapshot(t, svc, 1, keep)
	})
	if err != nil {
		t.Fatalf("sync node 1 without n1-gone: %v", err)
	}
	assertNoTraffic(t, "gone")
}

func TestNodeSyncClientGoneStaysOnOneConnection(t *testing.T) {
	db := initTrafficTestDB(t)
	bothSettings := `{"clients": [{"email": "gone", "enable": true}, {"email": "keep", "enable": true}]}`
	createNodeInboundWithSettings(t, db, 1, "n1-in", 41001, bothSettings)
	svc := &InboundService{}
	if err := syncNodeSnapshot(t, svc, 1, &model.Inbound{Tag: "n1-in", Enable: true, Settings: bothSettings, ClientStats: []xray.ClientTraffic{
		{Email: "gone", Up: 10, Down: 10, Enable: true},
		{Email: "keep", Up: 10, Down: 10, Enable: true},
	}}); err != nil {
		t.Fatalf("seed sync: %v", err)
	}
	readTraffic(t, db, "gone")

	capPoolAtOneConn(t)
	err := runWithoutSecondConn(t, "node sync per-client sweep", func() error {
		return syncNodeSnapshot(t, svc, 1, &model.Inbound{
			Tag:         "n1-in",
			Enable:      true,
			Settings:    `{"clients": [{"email": "keep", "enable": true}]}`,
			ClientStats: []xray.ClientTraffic{{Email: "keep", Up: 20, Down: 20, Enable: true}},
		})
	})
	if err != nil {
		t.Fatalf("sync node 1 without the gone client: %v", err)
	}
	var baselines int64
	if err := db.Model(&model.NodeClientTraffic{}).Where("node_id = ? AND email = ?", 1, "gone").Count(&baselines).Error; err != nil {
		t.Fatalf("count node baselines: %v", err)
	}
	if baselines != 0 {
		t.Fatalf("node baseline for the gone client = %d rows, want 0", baselines)
	}
}

func TestUpdateClientTrafficsStaysOnOneConnection(t *testing.T) {
	setupBulkDB(t)
	svc := &InboundService{}
	clients := []model.Client{{Email: "gone@x", ID: "aaaaaaaa-0000-0000-0000-000000000101", Enable: true}}
	oldIb := mkInbound(t, 22101, model.VLESS, clientsSettings(t, clients))
	if err := database.GetDB().Create(&xray.ClientTraffic{InboundId: oldIb.Id, Email: "gone@x", Enable: true}).Error; err != nil {
		t.Fatalf("seed traffic: %v", err)
	}
	newIb := *oldIb
	newIb.Settings = `{"clients": []}`

	capPoolAtOneConn(t)
	err := runWithoutSecondConn(t, "inbound edit dropping a client", func() error {
		return database.GetDB().Transaction(func(tx *gorm.DB) error {
			return svc.updateClientTraffics(tx, oldIb, &newIb)
		})
	})
	if err != nil {
		t.Fatalf("updateClientTraffics: %v", err)
	}
	var n int64
	database.GetDB().Model(xray.ClientTraffic{}).Where("email = ?", "gone@x").Count(&n)
	if n != 0 {
		t.Fatalf("traffic rows for dropped client = %d, want 0", n)
	}
}

func TestUpdateInboundClientRenameOntoExistingStatStaysOnOneConnection(t *testing.T) {
	setupBulkDB(t)
	svc := &ClientService{}
	inboundSvc := &InboundService{}
	source := []model.Client{{Email: "old@x", ID: "aaaaaaaa-0000-0000-0000-000000000102", SubID: "sub-102", Enable: true}}
	ib := mkInbound(t, 22102, model.VLESS, clientsSettings(t, source))
	if err := svc.SyncInbound(nil, ib.Id, source); err != nil {
		t.Fatalf("seed linkage: %v", err)
	}
	if err := database.GetDB().Create(&xray.ClientTraffic{InboundId: ib.Id, Email: "new@x", Enable: true}).Error; err != nil {
		t.Fatalf("seed target traffic: %v", err)
	}

	renamed := source
	renamed[0].Email = "new@x"
	capPoolAtOneConn(t)
	err := runWithoutSecondConn(t, "client rename onto an existing stats row", func() error {
		_, err := svc.UpdateInboundClient(inboundSvc, &model.Inbound{Id: ib.Id, Settings: clientsSettings(t, renamed)}, "old@x")
		return err
	})
	if err != nil {
		t.Fatalf("UpdateInboundClient: %v", err)
	}
	var n int64
	database.GetDB().Model(xray.ClientTraffic{}).Where("email = ?", "old@x").Count(&n)
	if n != 0 {
		t.Fatalf("traffic rows for renamed-away email = %d, want 0", n)
	}
}