package job import ( "encoding/json" "fmt" "io" "net/http" "net/http/httptest" "strconv" "strings" "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/web/runtime" "github.com/mhsanaei/3x-ui/v3/internal/xray" ) // resetGate holds every node reset open until released, counting how many nodes // the job reaches at once. type resetGate struct { entered atomic.Int32 release chan struct{} once sync.Once } func (g *resetGate) open() { g.once.Do(func() { close(g.release) }) } func (g *resetGate) waitAll(t *testing.T, want int32) { t.Helper() deadline := time.Now().Add(3 * time.Second) for g.entered.Load() < want { if time.Now().After(deadline) { t.Fatalf("periodic reset reached %d of %d hanging nodes, want all of them at once", g.entered.Load(), want) } time.Sleep(10 * time.Millisecond) } } // resetNode is a node whose every traffic reset hangs until the gate opens. func resetNode(t *testing.T, gate *resetGate, name string) int { t.Helper() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.Copy(io.Discard, r.Body) if strings.Contains(r.URL.Path, "resetTraffic") { gate.entered.Add(1) select { case <-r.Context().Done(): case <-gate.release: } } w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"success":true}`)) })) t.Cleanup(srv.Close) host, port, _ := strings.Cut(strings.TrimPrefix(srv.URL, "http://"), ":") portNum, _ := strconv.Atoi(port) node := &model.Node{ Name: name, Scheme: "http", Address: host, Port: portNum, BasePath: "/", ApiToken: "tok", Enable: true, Status: "online", AllowPrivateAddress: true, TlsVerifyMode: "verify", } if err := database.GetDB().Create(node).Error; err != nil { t.Fatalf("create node: %v", err) } return node.Id } func runResetJobAgainstGate(t *testing.T, gate *resetGate, want int32) { t.Helper() done := make(chan struct{}) go func() { defer close(done) NewPeriodicTrafficResetJob("daily", time.UTC).Run() }() t.Cleanup(func() { gate.open(); <-done }) gate.waitAll(t, want) } func newResetFleet(t *testing.T) *resetGate { t.Helper() initResetJobDB(t) runtime.SetManager(runtime.NewManager(runtime.LocalDeps{APIPort: func() int { return 0 }, SetNeedRestart: func() {}})) t.Cleanup(func() { runtime.SetManager(nil) }) return &resetGate{release: make(chan struct{})} } // The job reset due clients and inbounds one by one, each waiting on its node, // so a few hanging nodes stretched one run across hours. func TestPeriodicResetReachesClientNodesConcurrently(t *testing.T) { gate := newResetFleet(t) db := database.GetDB() for i := range 3 { nodeID := resetNode(t, gate, fmt.Sprintf("client-node-%d", i)) email := fmt.Sprintf("cycle-%d@node", i) client := model.Client{Email: email, ID: fmt.Sprintf("00000000-0000-4000-8000-00000000000%d", i), Enable: true, TrafficReset: "daily"} settings, _ := json.Marshal(map[string]any{"clients": []model.Client{client}}) ib := model.Inbound{ UserId: 1, Enable: true, Port: 47000 + i, Protocol: model.VLESS, NodeID: &nodeID, Tag: "reset-client-" + strconv.Itoa(i), TrafficReset: "never", Settings: string(settings), } if err := db.Create(&ib).Error; err != nil { t.Fatalf("create inbound: %v", err) } rec := model.ClientRecord{Email: email, UUID: client.ID, Enable: true, TrafficReset: "daily"} if err := db.Create(&rec).Error; err != nil { t.Fatalf("create client record: %v", err) } if err := db.Create(&model.ClientInbound{ClientId: rec.Id, InboundId: ib.Id}).Error; err != nil { t.Fatalf("link client: %v", err) } if err := db.Create(&xray.ClientTraffic{InboundId: ib.Id, Email: email, Enable: true, Up: 500, Down: 700}).Error; err != nil { t.Fatalf("create traffic: %v", err) } } runResetJobAgainstGate(t, gate, 3) } func TestPeriodicResetReachesInboundNodesConcurrently(t *testing.T) { gate := newResetFleet(t) for i := range 3 { nodeID := resetNode(t, gate, fmt.Sprintf("inbound-node-%d", i)) ib := model.Inbound{ UserId: 1, Enable: true, Port: 47100 + i, Protocol: model.VLESS, NodeID: &nodeID, Tag: "reset-inbound-" + strconv.Itoa(i), TrafficReset: "daily", Settings: `{"clients":[]}`, } if err := database.GetDB().Create(&ib).Error; err != nil { t.Fatalf("create inbound: %v", err) } } runResetJobAgainstGate(t, gate, 3) }