1
0
Эх сурвалжийг харах

fix(clients): zero traffic before re-enabling a client on reset

Invariant: a traffic reset leaves a quota-disabled client enabled
everywhere. ResetTrafficByEmail and BulkResetTraffic enabled the client
first and zeroed its counters afterwards. A traffic tick landing between
the two still saw the client depleted and switched it off again in
client_traffics, the record and settings. Update's direct record write
then set the record back to enabled and the reset zeroed the counters,
leaving the settings entry disabled: the client showed enabled with zero
usage but was dropped from the runtime. The periodic reset job goes through
ResetTrafficByEmail for every depleted client on its cycle, and a node push
inside Update widens the window to seconds.

Zero first, then enable: with the counters at zero the depletion predicate
no longer matches, so the tick has nothing to undo. UpdateInboundClient
re-adds the enabled user to the runtime itself.
MHSanaei 15 цаг өмнө
parent
commit
feb8451bd1

+ 34 - 28
internal/web/service/client_traffic.go

@@ -26,6 +26,26 @@ func (s *ClientService) ResetTrafficByEmail(inboundSvc *InboundService, email st
 	}
 
 	needRestart := false
+	if len(inboundIds) == 0 {
+		if rErr := inboundSvc.ResetClientTrafficByEmail(email); rErr != nil {
+			return false, rErr
+		}
+	} else {
+		applies := make([]inboundApply, 0, len(inboundIds))
+		for _, ibId := range inboundIds {
+			applies = append(applies, inboundApply{id: ibId, run: func() (bool, error) {
+				return inboundSvc.ResetClientTraffic(ibId, email)
+			}})
+		}
+		nr, applyErr := fanoutInboundApplies(applies)
+		if applyErr != nil {
+			return nr, applyErr
+		}
+		needRestart = nr
+	}
+
+	// Enable only once the counters are zero: a still-depleted client enabled
+	// first is switched off again by the next traffic tick.
 	if !rec.Enable {
 		updated := rec.ToClient()
 		updated.Enable = true
@@ -37,22 +57,7 @@ func (s *ClientService) ResetTrafficByEmail(inboundSvc *InboundService, email st
 			needRestart = true
 		}
 	}
-
-	if len(inboundIds) == 0 {
-		if rErr := inboundSvc.ResetClientTrafficByEmail(email); rErr != nil {
-			return false, rErr
-		}
-		return needRestart, nil
-	}
-
-	applies := make([]inboundApply, 0, len(inboundIds))
-	for _, ibId := range inboundIds {
-		applies = append(applies, inboundApply{id: ibId, run: func() (bool, error) {
-			return inboundSvc.ResetClientTraffic(ibId, email)
-		}})
-	}
-	nr, applyErr := fanoutInboundApplies(applies)
-	return needRestart || nr, applyErr
+	return needRestart, nil
 }
 
 func (s *ClientService) BulkResetTraffic(inboundSvc *InboundService, emails []string) (int, error) {
@@ -68,18 +73,6 @@ func (s *ClientService) BulkResetTraffic(inboundSvc *InboundService, emails []st
 	if err != nil {
 		return 0, err
 	}
-	for _, e := range cleanEmails {
-		rec := recordsByEmail[e]
-		if rec == nil || rec.Enable {
-			continue
-		}
-		updated := rec.ToClient()
-		updated.Enable = true
-		if _, uErr := s.Update(inboundSvc, rec.Id, *updated, rec.LimitHwid); uErr != nil {
-			logger.Warning("Failed to auto-enable client during bulk traffic reset:", uErr)
-		}
-	}
-
 	affected := 0
 	err = submitTrafficWrite(func() error {
 		db := database.GetDB()
@@ -110,6 +103,19 @@ func (s *ClientService) BulkResetTraffic(inboundSvc *InboundService, emails []st
 	if err != nil {
 		return 0, err
 	}
+	// After the zeroing, as in ResetTrafficByEmail: enabling a still-depleted
+	// client first lets the next traffic tick switch it off again.
+	for _, e := range cleanEmails {
+		rec := recordsByEmail[e]
+		if rec == nil || rec.Enable {
+			continue
+		}
+		updated := rec.ToClient()
+		updated.Enable = true
+		if _, uErr := s.Update(inboundSvc, rec.Id, *updated, rec.LimitHwid); uErr != nil {
+			logger.Warning("Failed to auto-enable client during bulk traffic reset:", uErr)
+		}
+	}
 	return affected, nil
 }
 

+ 91 - 0
internal/web/service/client_traffic_reset_order_test.go

@@ -0,0 +1,91 @@
+package service
+
+import (
+	"testing"
+	"time"
+
+	"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"
+)
+
+// A reset of a quota-disabled client must end enabled even when a traffic tick
+// runs between its steps: enabling before zeroing lets the tick re-disable it.
+func TestResetTrafficOfDepletedClientSurvivesTickMidReset(t *testing.T) {
+	resets := []struct {
+		name string
+		run  func() error
+	}{
+		{"single", func() error {
+			_, err := (&ClientService{}).ResetTrafficByEmail(&InboundService{}, "d@stale")
+			return err
+		}},
+		{"bulk", func() error {
+			_, err := (&ClientService{}).BulkResetTraffic(&InboundService{}, []string{"d@stale"})
+			return err
+		}},
+	}
+	for _, reset := range resets {
+		t.Run(reset.name, func(t *testing.T) {
+			requireResetSurvivesTick(t, reset.run)
+		})
+	}
+}
+
+func requireResetSurvivesTick(t *testing.T, reset func() error) {
+	t.Helper()
+	setupBulkDB(t)
+	clients := []model.Client{{Email: "d@stale", ID: "aaaaaaaa-0000-0000-0000-00000000000d", SubID: "sub-d", Enable: false, TotalGB: 1000}}
+	ib := mkInbound(t, 23140, model.VLESS, clientsSettings(t, clients))
+	if err := (&ClientService{}).SyncInbound(nil, ib.Id, clients); err != nil {
+		t.Fatalf("SyncInbound: %v", err)
+	}
+	row := xray.ClientTraffic{InboundId: ib.Id, Email: "d@stale", Enable: false, Up: 600, Down: 400, Total: 1000}
+	if err := database.GetDB().Create(&row).Error; err != nil {
+		t.Fatalf("seed client_traffics: %v", err)
+	}
+
+	resetTrafficWriterForTest(t)
+	StartTrafficWriter()
+	parked := make(chan struct{})
+	release := make(chan struct{})
+	go func() {
+		_ = submitTrafficWrite(func() error {
+			close(parked)
+			<-release
+			return nil
+		})
+	}()
+	<-parked
+
+	resetDone := make(chan error, 1)
+	go func() { resetDone <- reset() }()
+	waitTrafficWriterQueued(t)
+	tickDone := make(chan error, 1)
+	go func() {
+		_, _, err := (&InboundService{}).AddTraffic(nil, nil)
+		tickDone <- err
+	}()
+	deadline := time.Now().Add(5 * time.Second)
+	for len(twQueue) < 2 && time.Now().Before(deadline) {
+		time.Sleep(time.Millisecond)
+	}
+	close(release)
+	if err := <-resetDone; err != nil {
+		t.Fatalf("reset: %v", err)
+	}
+	if err := <-tickDone; err != nil {
+		t.Fatalf("AddTraffic: %v", err)
+	}
+
+	var after xray.ClientTraffic
+	if err := database.GetDB().Where("email = ?", "d@stale").First(&after).Error; err != nil {
+		t.Fatalf("read client_traffics: %v", err)
+	}
+	c, _ := settingsClient(t, ib.Id, "d@stale")
+	rec := lookupClientRecord(t, "d@stale")
+	if !after.Enable || !c.Enable || !rec.Enable || after.Up+after.Down != 0 {
+		t.Fatalf("after reset: traffic enable=%v used=%d, settings enable=%v, record enable=%v; want all enabled at 0",
+			after.Enable, after.Up+after.Down, c.Enable, rec.Enable)
+	}
+}