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 { _, err := (&ClientService{}).ResetAllClientTraffics(svc, ib.Id) return err }}, {"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 { _, err := (&ClientService{}).ResetAllClientTraffics(svc, ib.Id) return err }}, {"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) } } // selectiveResetRuntime refuses the resets of the emails in refuse. type selectiveResetRuntime struct { resetRecordingRuntime refuse map[string]bool } func (r *selectiveResetRuntime) ResetClientTraffic(ctx context.Context, ib *model.Inbound, email string) error { if r.refuse[email] { return errors.New("refused") } return r.resetRecordingRuntime.ResetClientTraffic(ctx, ib, email) } // One reset the node never accepts must not hold back the ones queued after it, nor // have the ones before it re-sent (each re-send wipes the usage made in between). func TestNodeResetsSettleEachClientOnItsOwn(t *testing.T) { nodeID, _, _ := setupRecordingNode(t, false) now := time.Now().UnixNano() for i, email := range []string{"a", "stuck", "c"} { if err := database.GetDB().Create(&model.NodePendingReset{NodeId: nodeID, Email: email, QueuedAt: now + int64(i)}).Error; err != nil { t.Fatalf("queue %s: %v", email, err) } } rt := &selectiveResetRuntime{refuse: map[string]bool{"stuck": true}} for range 2 { _ = (&InboundService{}).DeliverNodeResets(context.Background(), nodeID, rt) } if got := rt.delivered(); !slices.Equal(got, []string{"a", "c"}) { t.Fatalf("node received %v, want a and c once each", got) } if left := pendingResetEmails(t, nodeID); !slices.Equal(left, []string{"stuck"}) { t.Fatalf("still queued: %v, want only the refused reset", left) } } // stallingResetRuntime never answers a reset until its context ends. type stallingResetRuntime struct{ fakeNodeRuntime } func (*stallingResetRuntime) ResetClientTraffic(ctx context.Context, _ *model.Inbound, _ string) error { <-ctx.Done() return ctx.Err() } // A node relays its master's reset to its own sub-nodes inline; one silent sub-node must // not hold that request until the master times out and sends the reset again. func TestNodeResetDeliveryDoesNotHoldTheRequestBehindASilentNode(t *testing.T) { setupBulkDB(t) mgr := useTestRuntimeManager(t) node := &model.Node{Name: "silent", 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) } mgr.SetRuntimeOverride(node.Id, &stallingResetRuntime{}) ib := nodeInbound(t, node.Id, 41912, []model.Client{{Email: "held", ID: "11111111-1111-1111-1111-1111111111ab", Enable: true}}) if err := (&InboundService{}).AddClientStat(database.GetDB(), ib.Id, &model.Client{Email: "held", Enable: true}); err != nil { t.Fatalf("AddClientStat: %v", err) } start := time.Now() if _, err := (&InboundService{}).ResetClientTraffic(ib.Id, "held"); err != nil { t.Fatalf("ResetClientTraffic: %v", err) } if took := time.Since(start); took > 2*time.Second { t.Fatalf("the reset request took %v behind a silent node", took) } } // stallingUpdateRuntime never answers an inbound update until its context ends. type stallingUpdateRuntime struct{ fakeNodeRuntime } func (*stallingUpdateRuntime) UpdateInbound(ctx context.Context, _, _ *model.Inbound) error { <-ctx.Done() return ctx.Err() } // The sync job waits on these pushes before it serves any other node, so a hanging node // gets one push budget in all, not one per inbound it chose to report. func TestTrafficRemotePlansWaitOneBudgetPerNode(t *testing.T) { setupBulkDB(t) mgr := useTestRuntimeManager(t) node := &model.Node{Name: "hanging", 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) } mgr.SetRuntimeOverride(node.Id, &stallingUpdateRuntime{}) plans := make([]trafficInboundUpdatePlan, 9) for i := range plans { ib := model.Inbound{Id: i + 1, Tag: fmt.Sprintf("in-%d", i), NodeID: &node.Id} plans[i] = trafficInboundUpdatePlan{oldInbound: ib, newInbound: ib} } start := time.Now() (&InboundService{}).applyTrafficRemotePlans(plans) if took := time.Since(start); took > nodeClientPushTimeout+2*time.Second { t.Fatalf("pushing %d inbounds to one hanging node took %v", len(plans), took) } }