|
|
@@ -0,0 +1,275 @@
|
|
|
+package service
|
|
|
+
|
|
|
+import (
|
|
|
+ "context"
|
|
|
+ "errors"
|
|
|
+ "fmt"
|
|
|
+ "slices"
|
|
|
+ "sync"
|
|
|
+ "sync/atomic"
|
|
|
+ "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"
|
|
|
+
|
|
|
+ "gorm.io/gorm"
|
|
|
+)
|
|
|
+
|
|
|
+const (
|
|
|
+ resetLostOn = `{"clients":[{"email":"reset-lost","totalGB":100,"enable":true}]}`
|
|
|
+ resetLostOff = `{"clients":[{"email":"reset-lost","totalGB":100,"enable":false}]}`
|
|
|
+)
|
|
|
+
|
|
|
+// seedLatchedNodeClient leaves reset-lost depleted and latched off on the
|
|
|
+// master by its node's own usage, as a real node sync does.
|
|
|
+func seedLatchedNodeClient(t *testing.T, svc *InboundService) (*gorm.DB, *model.Inbound) {
|
|
|
+ t.Helper()
|
|
|
+ db := initTrafficTestDB(t)
|
|
|
+ createNodeInboundWithClient(t, db, 1, "n1-in", 41901, "reset-lost")
|
|
|
+ syncNodeWithSettings(t, svc, 1, "n1-in", resetLostOn, xray.ClientTraffic{Email: "reset-lost", Up: 10, Down: 10, Total: 100, Enable: true})
|
|
|
+ syncNodeWithSettings(t, svc, 1, "n1-in", resetLostOff, xray.ClientTraffic{Email: "reset-lost", Up: 60, Down: 60, Total: 100, Enable: false})
|
|
|
+ if got := readTraffic(t, db, "reset-lost"); got.Enable {
|
|
|
+ t.Fatal("setup: the depleted client should be latched off")
|
|
|
+ }
|
|
|
+ var ib model.Inbound
|
|
|
+ if err := db.Where("tag = ?", "n1-in").First(&ib).Error; err != nil {
|
|
|
+ t.Fatalf("load inbound: %v", err)
|
|
|
+ }
|
|
|
+ return db, &ib
|
|
|
+}
|
|
|
+
|
|
|
+// A reset the node never received leaves its old counters, so the node keeps
|
|
|
+// switching the client off; the master must not adopt that verdict.
|
|
|
+func TestNodeResetNotDeliveredDoesNotRedisableClient(t *testing.T) {
|
|
|
+ resets := []struct {
|
|
|
+ name string
|
|
|
+ run func(svc *InboundService, ib *model.Inbound) error
|
|
|
+ }{
|
|
|
+ {"single", func(svc *InboundService, ib *model.Inbound) error {
|
|
|
+ _, err := svc.ResetClientTraffic(ib.Id, "reset-lost")
|
|
|
+ return err
|
|
|
+ }},
|
|
|
+ {"bulk", func(svc *InboundService, _ *model.Inbound) error {
|
|
|
+ _, err := (&ClientService{}).BulkResetTraffic(svc, []string{"reset-lost"})
|
|
|
+ return err
|
|
|
+ }},
|
|
|
+ {"inbound", func(svc *InboundService, ib *model.Inbound) error {
|
|
|
+ return (&ClientService{}).ResetAllClientTraffics(svc, ib.Id)
|
|
|
+ }},
|
|
|
+ {"all", func(*InboundService, *model.Inbound) error {
|
|
|
+ _, err := (&ClientService{}).ResetAllTraffics()
|
|
|
+ return err
|
|
|
+ }},
|
|
|
+ }
|
|
|
+ for _, reset := range resets {
|
|
|
+ t.Run(reset.name, func(t *testing.T) {
|
|
|
+ svc := &InboundService{}
|
|
|
+ db, ib := seedLatchedNodeClient(t, svc)
|
|
|
+ if err := reset.run(svc, ib); err != nil {
|
|
|
+ t.Fatalf("reset: %v", err)
|
|
|
+ }
|
|
|
+ syncNodeWithSettings(t, svc, 1, "n1-in", resetLostOff, xray.ClientTraffic{Email: "reset-lost", Up: 60, Down: 60, Total: 100, Enable: false})
|
|
|
+ got := readTraffic(t, db, "reset-lost")
|
|
|
+ if !got.Enable || got.Up+got.Down != 0 {
|
|
|
+ t.Fatalf("after reset: enable=%v used=%d, want enabled at 0 — the undelivered reset re-disabled it", got.Enable, got.Up+got.Down)
|
|
|
+ }
|
|
|
+ })
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+// resetRecordingRuntime is a node that accepts per-client resets unless failing.
|
|
|
+type resetRecordingRuntime struct {
|
|
|
+ fakeNodeRuntime
|
|
|
+ mu sync.Mutex
|
|
|
+ fail bool
|
|
|
+ got []string
|
|
|
+}
|
|
|
+
|
|
|
+func (r *resetRecordingRuntime) ResetClientTraffic(_ context.Context, _ *model.Inbound, email string) error {
|
|
|
+ r.mu.Lock()
|
|
|
+ defer r.mu.Unlock()
|
|
|
+ if r.fail {
|
|
|
+ return errors.New("node unreachable")
|
|
|
+ }
|
|
|
+ r.got = append(r.got, email)
|
|
|
+ return nil
|
|
|
+}
|
|
|
+
|
|
|
+func (r *resetRecordingRuntime) delivered() []string {
|
|
|
+ r.mu.Lock()
|
|
|
+ defer r.mu.Unlock()
|
|
|
+ return slices.Clone(r.got)
|
|
|
+}
|
|
|
+
|
|
|
+func pendingResetEmails(t *testing.T, nodeID int) []string {
|
|
|
+ t.Helper()
|
|
|
+ var emails []string
|
|
|
+ if err := database.GetDB().Model(&model.NodePendingReset{}).Where("node_id = ?", nodeID).
|
|
|
+ Order("email").Pluck("email", &emails).Error; err != nil {
|
|
|
+ t.Fatalf("read pending resets: %v", err)
|
|
|
+ }
|
|
|
+ return emails
|
|
|
+}
|
|
|
+
|
|
|
+func setupRecordingNode(t *testing.T, fail bool) (int, *resetRecordingRuntime, *model.Inbound) {
|
|
|
+ t.Helper()
|
|
|
+ setupBulkDB(t)
|
|
|
+ mgr := useTestRuntimeManager(t)
|
|
|
+ node := &model.Node{Name: "reset-node", Address: "127.0.0.1", Port: 2096, ApiToken: "tok", Enable: true, Status: "online"}
|
|
|
+ if err := database.GetDB().Create(node).Error; err != nil {
|
|
|
+ t.Fatalf("create node: %v", err)
|
|
|
+ }
|
|
|
+ rec := &resetRecordingRuntime{fail: fail}
|
|
|
+ mgr.SetRuntimeOverride(node.Id, rec)
|
|
|
+ ib := nodeInbound(t, node.Id, 41911, []model.Client{{Email: "reset-lost", ID: "11111111-1111-1111-1111-1111111111aa", Enable: true}})
|
|
|
+ if err := (&InboundService{}).AddClientStat(database.GetDB(), ib.Id, &model.Client{Email: "reset-lost", Enable: true}); err != nil {
|
|
|
+ t.Fatalf("AddClientStat: %v", err)
|
|
|
+ }
|
|
|
+ return node.Id, rec, ib
|
|
|
+}
|
|
|
+
|
|
|
+// A reachable node gets the reset right after the master commits it.
|
|
|
+func TestNodeResetDeliveredRightAway(t *testing.T) {
|
|
|
+ resets := []struct {
|
|
|
+ name string
|
|
|
+ run func(svc *InboundService, ib *model.Inbound) error
|
|
|
+ }{
|
|
|
+ {"single", func(svc *InboundService, ib *model.Inbound) error {
|
|
|
+ _, err := svc.ResetClientTraffic(ib.Id, "reset-lost")
|
|
|
+ return err
|
|
|
+ }},
|
|
|
+ {"bulk", func(svc *InboundService, _ *model.Inbound) error {
|
|
|
+ _, err := (&ClientService{}).BulkResetTraffic(svc, []string{"reset-lost"})
|
|
|
+ return err
|
|
|
+ }},
|
|
|
+ {"inbound", func(svc *InboundService, ib *model.Inbound) error {
|
|
|
+ return (&ClientService{}).ResetAllClientTraffics(svc, ib.Id)
|
|
|
+ }},
|
|
|
+ {"all", func(*InboundService, *model.Inbound) error {
|
|
|
+ _, err := (&ClientService{}).ResetAllTraffics()
|
|
|
+ return err
|
|
|
+ }},
|
|
|
+ }
|
|
|
+ for _, reset := range resets {
|
|
|
+ t.Run(reset.name, func(t *testing.T) {
|
|
|
+ nodeID, rec, ib := setupRecordingNode(t, false)
|
|
|
+ if err := reset.run(&InboundService{}, ib); err != nil {
|
|
|
+ t.Fatalf("reset: %v", err)
|
|
|
+ }
|
|
|
+ if got := rec.delivered(); !slices.Equal(got, []string{"reset-lost"}) {
|
|
|
+ t.Fatalf("node received resets %v, want [reset-lost]", got)
|
|
|
+ }
|
|
|
+ if left := pendingResetEmails(t, nodeID); len(left) != 0 {
|
|
|
+ t.Fatalf("delivered reset still queued: %v", left)
|
|
|
+ }
|
|
|
+ })
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+// bulkResetRuntime also takes a batch in one call.
|
|
|
+type bulkResetRuntime struct {
|
|
|
+ resetRecordingRuntime
|
|
|
+ batches [][]string
|
|
|
+}
|
|
|
+
|
|
|
+func (b *bulkResetRuntime) ResetClientTraffics(_ context.Context, emails []string) error {
|
|
|
+ b.mu.Lock()
|
|
|
+ defer b.mu.Unlock()
|
|
|
+ b.batches = append(b.batches, slices.Clone(emails))
|
|
|
+ return nil
|
|
|
+}
|
|
|
+
|
|
|
+// Above the per-client push threshold a backlog goes out as one bulk request,
|
|
|
+// not one round-trip per client.
|
|
|
+func TestNodeResetBacklogUsesBulkRequest(t *testing.T) {
|
|
|
+ setupBulkDB(t)
|
|
|
+ const nodeID = 7
|
|
|
+ rows := make([]model.NodePendingReset, nodeBulkPushThreshold+1)
|
|
|
+ for i := range rows {
|
|
|
+ rows[i] = model.NodePendingReset{NodeId: nodeID, Email: fmt.Sprintf("owed-%02d", i), QueuedAt: 1}
|
|
|
+ }
|
|
|
+ if err := database.GetDB().Create(&rows).Error; err != nil {
|
|
|
+ t.Fatalf("seed pending resets: %v", err)
|
|
|
+ }
|
|
|
+ rt := &bulkResetRuntime{}
|
|
|
+ if err := (&InboundService{}).DeliverNodeResets(context.Background(), nodeID, rt); err != nil {
|
|
|
+ t.Fatalf("DeliverNodeResets: %v", err)
|
|
|
+ }
|
|
|
+ if len(rt.batches) != 1 || len(rt.batches[0]) != len(rows) || len(rt.delivered()) != 0 {
|
|
|
+ t.Fatalf("bulk batches %d (first %d emails), per-client calls %d; want one batch of %d",
|
|
|
+ len(rt.batches), len(rt.batches[0]), len(rt.delivered()), len(rows))
|
|
|
+ }
|
|
|
+ if left := pendingResetEmails(t, nodeID); len(left) != 0 {
|
|
|
+ t.Fatalf("delivered backlog still queued: %d rows", len(left))
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+// An unreachable node keeps the reset queued until a later delivery lands.
|
|
|
+func TestNodeResetReplayedAfterFailure(t *testing.T) {
|
|
|
+ nodeID, rec, ib := setupRecordingNode(t, true)
|
|
|
+ if _, err := (&InboundService{}).ResetClientTraffic(ib.Id, "reset-lost"); err != nil {
|
|
|
+ t.Fatalf("ResetClientTraffic: %v", err)
|
|
|
+ }
|
|
|
+ if left := pendingResetEmails(t, nodeID); !slices.Equal(left, []string{"reset-lost"}) {
|
|
|
+ t.Fatalf("pending after failed delivery = %v, want [reset-lost]", left)
|
|
|
+ }
|
|
|
+
|
|
|
+ rec.mu.Lock()
|
|
|
+ rec.fail = false
|
|
|
+ rec.mu.Unlock()
|
|
|
+ if err := (&InboundService{}).DeliverNodeResets(context.Background(), nodeID, rec); err != nil {
|
|
|
+ t.Fatalf("DeliverNodeResets: %v", err)
|
|
|
+ }
|
|
|
+ if got := rec.delivered(); !slices.Equal(got, []string{"reset-lost"}) {
|
|
|
+ t.Fatalf("node received resets %v, want [reset-lost]", got)
|
|
|
+ }
|
|
|
+ if left := pendingResetEmails(t, nodeID); len(left) != 0 {
|
|
|
+ t.Fatalf("delivered reset still queued: %v", left)
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+// slowResetRuntime holds each reset until a second one arrives or a short
|
|
|
+// timeout passes, so two unserialized deliveries both reach the node.
|
|
|
+type slowResetRuntime struct {
|
|
|
+ resetRecordingRuntime
|
|
|
+ calls atomic.Int32
|
|
|
+ both chan struct{}
|
|
|
+}
|
|
|
+
|
|
|
+func (r *slowResetRuntime) ResetClientTraffic(ctx context.Context, ib *model.Inbound, email string) error {
|
|
|
+ if r.calls.Add(1) == 2 {
|
|
|
+ close(r.both)
|
|
|
+ }
|
|
|
+ select {
|
|
|
+ case <-r.both:
|
|
|
+ case <-time.After(300 * time.Millisecond):
|
|
|
+ }
|
|
|
+ return r.resetRecordingRuntime.ResetClientTraffic(ctx, ib, email)
|
|
|
+}
|
|
|
+
|
|
|
+// The sync job and a reset's own delivery can run at once; the node must still
|
|
|
+// get each owed reset once, or usage made in between is wiped a second time.
|
|
|
+func TestConcurrentNodeResetDeliveriesSendOnce(t *testing.T) {
|
|
|
+ setupBulkDB(t)
|
|
|
+ const nodeID = 9
|
|
|
+ if err := database.GetDB().Create(&model.NodePendingReset{NodeId: nodeID, Email: "once", QueuedAt: 1}).Error; err != nil {
|
|
|
+ t.Fatalf("seed pending reset: %v", err)
|
|
|
+ }
|
|
|
+ rt := &slowResetRuntime{both: make(chan struct{})}
|
|
|
+ var wg sync.WaitGroup
|
|
|
+ for range 2 {
|
|
|
+ wg.Add(1)
|
|
|
+ go func() {
|
|
|
+ defer wg.Done()
|
|
|
+ if err := (&InboundService{}).DeliverNodeResets(context.Background(), nodeID, rt); err != nil {
|
|
|
+ t.Errorf("DeliverNodeResets: %v", err)
|
|
|
+ }
|
|
|
+ }()
|
|
|
+ }
|
|
|
+ wg.Wait()
|
|
|
+ if got := rt.delivered(); !slices.Equal(got, []string{"once"}) {
|
|
|
+ t.Fatalf("node received resets %v, want exactly [once]", got)
|
|
|
+ }
|
|
|
+}
|
MHSanaei