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

package service

import (
	"errors"
	"fmt"
	"path/filepath"
	"testing"
	"time"

	"github.com/mhsanaei/3x-ui/v3/internal/database"
	"github.com/mhsanaei/3x-ui/v3/internal/database/model"

	"gorm.io/gorm"
)

var errInjectedHwidDelete = errors.New("injected client_hwids delete failure")

func failHwidDeletes(t *testing.T, db *gorm.DB) {
	t.Helper()
	if err := db.Callback().Delete().Before("gorm:delete").Register("t:hwid:fail", func(tx *gorm.DB) {
		if tx.Statement != nil && tx.Statement.Table == "client_hwids" {
			_ = tx.AddError(errInjectedHwidDelete)
		}
	}); err != nil {
		t.Fatalf("register callback: %v", err)
	}
	t.Cleanup(func() {
		if err := db.Callback().Delete().Remove("t:hwid:fail"); err != nil {
			t.Fatalf("remove callback: %v", err)
		}
	})
}

func seedHwids(t *testing.T, db *gorm.DB, subID string, n int) {
	t.Helper()
	now := time.Now().UnixMilli()
	rows := make([]model.ClientHwid, 0, n)
	for i := range n {
		rows = append(rows, model.ClientHwid{
			SubID: subID, HwidHash: fmt.Sprintf("%s-hash-%d", subID, i),
			FirstSeen: now, LastSeen: now + int64(i),
		})
	}
	if err := db.Create(&rows).Error; err != nil {
		t.Fatalf("seed client_hwids for %q: %v", subID, err)
	}
}

func assertHwidState(t *testing.T, db *gorm.DB, email string, limit, devices int) {
	t.Helper()
	var rec model.ClientRecord
	if err := db.Where("email = ?", email).First(&rec).Error; err != nil {
		t.Fatalf("reload client: %v", err)
	}
	if rec.LimitHwid != limit {
		t.Fatalf("limit_hwid = %d, want %d", rec.LimitHwid, limit)
	}
	var n int64
	if err := db.Model(&model.ClientHwid{}).Where("sub_id = ?", rec.SubID).Count(&n).Error; err != nil {
		t.Fatalf("count client_hwids: %v", err)
	}
	if n != int64(devices) {
		t.Fatalf("client_hwids = %d, want %d", n, devices)
	}
}

func TestSetClientLimitHwidRollsBackFailedTrim(t *testing.T) {
	initClientHwidTestDB(t)
	db := database.GetDB()
	rec := seedHwidClient(t, 5)
	seedHwids(t, db, rec.SubID, 3)
	failHwidDeletes(t, db)

	err := (&ClientService{}).setClientLimitHwidByEmail(rec.Email, 1)
	if !errors.Is(err, errInjectedHwidDelete) {
		t.Fatalf("want errInjectedHwidDelete, got: %v", err)
	}
	assertHwidState(t, db, rec.Email, 5, 3)
}

func TestClientHwidTxRejectsUnserializedHandle(t *testing.T) {
	initClientHwidTestDB(t)
	db := database.GetDB()
	rec := seedHwidClient(t, 5)
	svc := &ClientService{}

	if err := svc.setClientLimitHwidByEmailTx(db, rec.Email, 1); !errors.Is(err, errClientHwidWriteNotSerialized) {
		t.Fatalf("bare handle error = %v, want errClientHwidWriteNotSerialized", err)
	}
	if err := runSerializedTx(func(tx *gorm.DB) error {
		return svc.setClientLimitHwidByEmailTx(tx, rec.Email, 1)
	}); err != nil {
		t.Fatalf("serialized update: %v", err)
	}
	assertHwidState(t, db, rec.Email, 1, 0)
}

func TestBulkAdjustHwidRollsBackFailedTrim(t *testing.T) {
	setupBulkDB(t)
	db := database.GetDB()
	rec := &model.ClientRecord{Email: "bulk-hwid@x", SubID: "bulk-sub", Enable: true, LimitHwid: 5}
	if err := db.Create(rec).Error; err != nil {
		t.Fatalf("seed client: %v", err)
	}
	seedHwids(t, db, rec.SubID, 3)
	failHwidDeletes(t, db)

	limit := 1
	res, _, err := (&ClientService{}).BulkAdjust(&InboundService{}, []string{rec.Email}, 0, 0, "", &limit, "")
	if err != nil {
		t.Fatalf("BulkAdjust: %v", err)
	}
	if len(res.Skipped) != 1 || res.Skipped[0].Reason != errInjectedHwidDelete.Error() {
		t.Fatalf("skipped = %+v, want injected failure", res.Skipped)
	}
	assertHwidState(t, db, rec.Email, 5, 3)
}

func TestBulkCreateWithdrawsTombstoneWhenHwidTrimFails(t *testing.T) {
	setupBulkDB(t)
	StartTrafficWriter()
	t.Cleanup(StopTrafficWriter)
	db := database.GetDB()
	const email = "reborn-bulk@x"
	const subID = "reborn-bulk-sub"
	tombstoneClientEmail(email)
	t.Cleanup(func() { withdrawClientTombstones(email) })
	seedHwids(t, db, subID, 3)
	failHwidDeletes(t, db)
	ib := mkInbound(t, 30441, model.VLESS, `{"clients":[]}`)

	res, _, err := (&ClientService{}).BulkCreate(&InboundService{}, []ClientCreatePayload{{
		Client: model.Client{
			Email: email, SubID: subID, ID: "aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee", Enable: true,
		},
		InboundIds: []int{ib.Id}, LimitHwid: 1,
	}})
	if err != nil {
		t.Fatalf("BulkCreate: %v", err)
	}
	if len(res.Skipped) != 1 || res.Skipped[0].Reason != errInjectedHwidDelete.Error() {
		t.Fatalf("skipped = %+v, want injected HWID failure", res.Skipped)
	}
	if isClientEmailTombstoned(email) {
		t.Fatal("live bulk-created client retained a delete tombstone")
	}
}

func TestSetClientLimitHwidIsSerializedWithSyncInbound(t *testing.T) {
	db := durablePostgresDB(t)
	if err := db.Exec("TRUNCATE client_hwids, clients RESTART IDENTITY CASCADE").Error; err != nil {
		t.Fatalf("reset tables: %v", err)
	}
	rec := seedHwidClient(t, 5)
	seedHwids(t, db, rec.SubID, 3)
	StartTrafficWriter()
	t.Cleanup(StopTrafficWriter)

	read := make(chan struct{})
	release := make(chan struct{})
	staleDone := make(chan error, 1)
	go func() {
		staleDone <- runSerializedTx(func(tx *gorm.DB) error {
			var stale model.ClientRecord
			if err := tx.Where("email = ?", rec.Email).First(&stale).Error; err != nil {
				return err
			}
			close(read)
			<-release
			return tx.Save(&stale).Error
		})
	}()
	<-read

	limitDone := make(chan error, 1)
	go func() { limitDone <- (&ClientService{}).setClientLimitHwidByEmail(rec.Email, 1) }()
	time.Sleep(100 * time.Millisecond)
	close(release)
	if err := <-staleDone; err != nil {
		t.Fatalf("stale SyncInbound write: %v", err)
	}
	if err := <-limitDone; err != nil {
		t.Fatalf("set limit: %v", err)
	}
	assertHwidState(t, db, rec.Email, 1, 1)
}

func BenchmarkSetClientLimitHwidSerialized(b *testing.B) {
	dbDir := b.TempDir()
	b.Setenv("XUI_DB_FOLDER", dbDir)
	if err := database.InitDB(filepath.Join(dbDir, "x-ui.db")); err != nil {
		b.Fatalf("InitDB: %v", err)
	}
	b.Cleanup(func() { _ = database.CloseDB() })
	StartTrafficWriter()
	b.Cleanup(StopTrafficWriter)
	db := database.GetDB()
	emails := make([]string, 100)
	for i := range emails {
		emails[i] = fmt.Sprintf("bench-%03d@x", i)
		rec := &model.ClientRecord{Email: emails[i], SubID: fmt.Sprintf("bench-sub-%03d", i), Enable: true}
		if err := db.Create(rec).Error; err != nil {
			b.Fatalf("seed client: %v", err)
		}
	}
	svc := &ClientService{}
	for _, count := range []int{1, 100} {
		b.Run(fmt.Sprintf("clients_%d", count), func(b *testing.B) {
			for range b.N {
				for i := range count {
					if err := svc.setClientLimitHwidByEmail(emails[i], 2); err != nil {
						b.Fatal(err)
					}
				}
			}
		})
	}
}