3x-ui-agent/internal/agent/agent_traffic_test.go

package agent

import (
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"

	"github.com/mhsanaei/3x-ui/v3/internal/xray"
)

// fakeStatsClient implements the snapshot/commit contract over scripted
// cumulative counters: snapshots diff against the committed baseline with
// the same per-stat rule as the real client (unseen or moved-backwards
// counters count from zero), and only Commit moves the baseline. The
// delta math itself is proven against the wire-level fake in the xray
// package; here the fake exists so the agent's orchestration — commit
// only after a 200, persist, restore after restart — can be driven
// exactly.
type fakeStatsClient struct {
	counters  map[string]int64
	committed map[string]int64
	ready     bool
}

func newFakeStatsClient() *fakeStatsClient {
	return &fakeStatsClient{counters: map[string]int64{}}
}

func (f *fakeStatsClient) GetTrafficSnapshot() (*xray.TrafficSnapshot, error) {
	snap := &xray.TrafficSnapshot{Values: map[string]int64{}}
	clients := map[string]*xray.ClientTraffic{}
	for name, value := range f.counters {
		snap.Values[name] = value
		if !f.ready {
			continue
		}
		last, ok := f.committed[name]
		if !ok || value < last {
			last = 0
		}
		delta := value - last
		parts := strings.Split(name, ">>>")
		if len(parts) != 4 || parts[0] != "user" {
			continue
		}
		ct, ok := clients[parts[1]]
		if !ok {
			ct = &xray.ClientTraffic{Email: parts[1]}
			clients[parts[1]] = ct
		}
		if parts[3] == "downlink" {
			ct.Down = delta
		} else {
			ct.Up = delta
		}
	}
	for _, ct := range clients {
		snap.ClientTraffics = append(snap.ClientTraffics, ct)
	}
	return snap, nil
}

func (f *fakeStatsClient) CommitTrafficSnapshot(s *xray.TrafficSnapshot) {
	f.committed = s.BaselineValues()
	f.ready = true
}

func (f *fakeStatsClient) RestoreTrafficBaseline(values map[string]int64) {
	if len(values) == 0 {
		return
	}
	f.committed = map[string]int64{}
	for k, v := range values {
		f.committed[k] = v
	}
	f.ready = true
}

func (f *fakeStatsClient) GetOnlineUsers() ([]xray.OnlineUser, error) { return nil, nil }

// fakeMaster is an HTTP master whose /panel/api/agent/traffic endpoint
// fails (500) for the first failTimes requests, then acknowledges and
// records everything it receives.
type fakeMaster struct {
	srv       *httptest.Server
	mu        sync.Mutex
	failTimes int
	posts     int
	userUp    map[string]int64
	userDown  map[string]int64
}

func newFakeMaster(t *testing.T, failTimes int) *fakeMaster {
	t.Helper()
	m := &fakeMaster{failTimes: failTimes, userUp: map[string]int64{}, userDown: map[string]int64{}}
	m.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		m.mu.Lock()
		defer m.mu.Unlock()
		m.posts++
		if m.posts <= m.failTimes {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		var report trafficReport
		if err := json.NewDecoder(r.Body).Decode(&report); err != nil {
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		for _, u := range report.Users {
			m.userUp[u.Email] += u.Up
			m.userDown[u.Email] += u.Down
		}
		w.WriteHeader(http.StatusOK)
	}))
	t.Cleanup(m.srv.Close)
	return m
}

func newTestAgent(t *testing.T, masterURL string) *Agent {
	t.Helper()
	return New(&Config{
		MasterURL:    masterURL,
		Token:        "test-token",
		NodeName:     "test-node",
		BaselinePath: filepath.Join(t.TempDir(), "agent-traffic-baseline.json"),
	})
}

const aliceUp = "user>>>alice>>>traffic>>>uplink"

// TestReportOutageThenRecoveryConservesTotals is the acceptance scenario
// for the fix: the master is down for several polls while traffic accrues,
// then recovers. Every byte the core counted must reach the master exactly
// once — the recovery report carries the whole outage window, and later
// polls carry only new traffic.
func TestReportOutageThenRecoveryConservesTotals(t *testing.T) {
	master := newFakeMaster(t, 2) // first two posts fail with 500
	a := newTestAgent(t, master.srv.URL)
	fake := newFakeStatsClient()
	fake.counters[aliceUp] = 1000 // pre-existing, unattributable

	// Poll 1 establishes the baseline: nothing is posted, the baseline
	// file appears with the raw counter.
	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
		t.Fatalf("baseline poll: %v", err)
	}
	if master.posts != 0 {
		t.Fatalf("baseline poll posted %d times, want 0", master.posts)
	}
	values, err := loadTrafficBaseline(a.baselinePath())
	if err != nil || values[aliceUp] != 1000 {
		t.Fatalf("persisted baseline = %v (err %v), want alice at 1000", values, err)
	}

	// Outage: two polls accrue traffic but the master 500s.
	fake.counters[aliceUp] = 1400
	if err := a.reportTrafficWith(fake, a.postTrafficReport); err == nil {
		t.Fatal("poll during outage returned nil, want the post error")
	}
	fake.counters[aliceUp] = 2000
	if err := a.reportTrafficWith(fake, a.postTrafficReport); err == nil {
		t.Fatal("second outage poll returned nil, want the post error")
	}
	// The failures advanced nothing: the persisted baseline still holds
	// the last acknowledged counters.
	values, err = loadTrafficBaseline(a.baselinePath())
	if err != nil || values[aliceUp] != 1000 {
		t.Fatalf("baseline after failed polls = %v (err %v), want still 1000", values, err)
	}

	// Recovery: one post carries the entire outage window (2000-1000).
	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
		t.Fatalf("recovery poll: %v", err)
	}
	if got := master.userUp["alice"]; got != 1000 {
		t.Fatalf("master received %d for alice after recovery, want exactly 1000", got)
	}

	// Steady state resumes: only fresh growth is reported.
	fake.counters[aliceUp] = 2100
	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
		t.Fatalf("steady poll: %v", err)
	}
	if got := master.userUp["alice"]; got != 1100 {
		t.Fatalf("master total = %d, want 1100 (= 2100 counted - 1000 baseline)", got)
	}
}

// TestAgentRestartRestoresBaselineFromDisk: a new agent process (fresh
// in-memory state everywhere) over the same baseline file must pick up
// diffing exactly where the old one stopped — the traffic that accrued
// while it was down is reported by its FIRST poll, and nothing is counted
// twice.
func TestAgentRestartRestoresBaselineFromDisk(t *testing.T) {
	master := newFakeMaster(t, 0)
	baselinePath := filepath.Join(t.TempDir(), "agent-traffic-baseline.json")
	newAgent := func() *Agent {
		return New(&Config{MasterURL: master.srv.URL, Token: "t", NodeName: "n", BaselinePath: baselinePath})
	}

	a1 := newAgent()
	fake := newFakeStatsClient()
	fake.counters[aliceUp] = 1000
	if err := a1.reportTrafficWith(fake, a1.postTrafficReport); err != nil {
		t.Fatalf("a1 baseline poll: %v", err)
	}
	fake.counters[aliceUp] = 1500
	if err := a1.reportTrafficWith(fake, a1.postTrafficReport); err != nil {
		t.Fatalf("a1 report poll: %v", err)
	}
	if got := master.userUp["alice"]; got != 500 {
		t.Fatalf("master received %d before restart, want 500", got)
	}

	// "Restart": brand-new agent AND a brand-new stats client (its
	// in-memory committed baseline is gone); only the file survives. The
	// core kept counting meanwhile.
	a2 := newAgent()
	fake2 := newFakeStatsClient()
	fake2.counters[aliceUp] = 1900
	if err := a2.reportTrafficWith(fake2, a2.postTrafficReport); err != nil {
		t.Fatalf("a2 first poll: %v", err)
	}
	if got := master.userUp["alice"]; got != 900 {
		t.Fatalf("master total = %d, want 900 (500 + the 400 accrued across the restart)", got)
	}
}

// TestCounterResetAfterCoreRestart: the core restarts while the agent
// keeps running; its counters fall below the committed baseline. The
// report must carry the new counter as fresh traffic — never a negative
// delta — and then resume normal diffing in the new epoch.
func TestCounterResetAfterCoreRestart(t *testing.T) {
	master := newFakeMaster(t, 0)
	a := newTestAgent(t, master.srv.URL)
	fake := newFakeStatsClient()
	fake.counters[aliceUp] = 5000
	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
		t.Fatalf("baseline poll: %v", err)
	}

	fake.counters[aliceUp] = 300 // core restarted; counters began again
	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
		t.Fatalf("post-reset poll: %v", err)
	}
	if got := master.userUp["alice"]; got != 300 {
		t.Fatalf("master received %d after core restart, want the fresh 300", got)
	}

	fake.counters[aliceUp] = 450
	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
		t.Fatalf("new-epoch poll: %v", err)
	}
	if got := master.userUp["alice"]; got != 450 {
		t.Fatalf("master total = %d, want 450 (300 + 150 of new-epoch growth)", got)
	}
}

// TestCorruptBaselineFileTolerated: a torn or foreign baseline file must
// not wedge the agent — it logs and falls back to a baseline-only first
// snapshot, then reports normally.
func TestCorruptBaselineFileTolerated(t *testing.T) {
	master := newFakeMaster(t, 0)
	a := newTestAgent(t, master.srv.URL)
	if err := os.WriteFile(a.baselinePath(), []byte("{not json"), 0o600); err != nil {
		t.Fatal(err)
	}
	fake := newFakeStatsClient()
	fake.counters[aliceUp] = 800
	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
		t.Fatalf("poll with corrupt baseline: %v", err)
	}
	if master.posts != 0 {
		t.Fatalf("corrupt-baseline first poll posted %d times, want baseline-only (0)", master.posts)
	}
	fake.counters[aliceUp] = 900
	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
		t.Fatalf("poll after corrupt baseline: %v", err)
	}
	if got := master.userUp["alice"]; got != 100 {
		t.Fatalf("master received %d, want only the post-baseline 100", got)
	}
}

// TestSaveLoadBaselineRoundTrip pins the persistence helpers: atomic save,
// version field, exact values back.
func TestSaveLoadBaselineRoundTrip(t *testing.T) {
	path := filepath.Join(t.TempDir(), "baseline.json")
	want := map[string]int64{aliceUp: 42, "inbound>>>in-443>>>traffic>>>downlink": 7}
	if err := saveTrafficBaseline(path, want); err != nil {
		t.Fatalf("save: %v", err)
	}
	got, err := loadTrafficBaseline(path)
	if err != nil {
		t.Fatalf("load: %v", err)
	}
	if len(got) != len(want) || got[aliceUp] != 42 || got["inbound>>>in-443>>>traffic>>>downlink"] != 7 {
		t.Fatalf("round trip = %v, want %v", got, want)
	}
	if _, err := os.Stat(path + ".tmp"); !os.IsNotExist(err) {
		t.Fatal("temp file left behind after atomic save")
	}
	if _, err := loadTrafficBaseline(filepath.Join(t.TempDir(), "missing.json")); !os.IsNotExist(err) {
		t.Fatalf("missing file error = %v, want a not-exist error", err)
	}
	var parseErr error
	_, parseErr = loadTrafficBaseline(path)
	if parseErr != nil {
		t.Fatalf("valid file failed to load: %v", parseErr)
	}
	bad := filepath.Join(t.TempDir(), "bad.json")
	_ = os.WriteFile(bad, []byte("[]"), 0o600)
	if _, err := loadTrafficBaseline(bad); err == nil || errors.Is(err, os.ErrNotExist) {
		t.Fatalf("corrupt file error = %v, want a parse error", err)
	}
}