3x-ui-agent/internal/xray/api_traffic_snapshot_test.go

package xray

import "testing"

// The snapshot/commit contract is what keeps agent traffic reporting from
// losing deltas when the master is unreachable: a snapshot never advances
// the baseline, only CommitTrafficSnapshot does. These tests pin that
// contract against the scripted stats server, plus the recovery paths
// (baseline restore after an agent restart, counter reset after a core
// restart) and the isolation from GetTraffic's read-advances semantics.

func clientUp(t *testing.T, traffics []*ClientTraffic, email string) (int64, bool) {
	t.Helper()
	for _, ct := range traffics {
		if ct.Email == email {
			return ct.Up, true
		}
	}
	return 0, false
}

// TestTrafficSnapshotFirstPassIsBaselineOnly mirrors the GetTraffic rule:
// counters may already hold unattributable traffic when reporting starts,
// so the first snapshot proposes a baseline and reports nothing — and
// committing it is what arms delta reporting.
func TestTrafficSnapshotFirstPassIsBaselineOnly(t *testing.T) {
	api := startFakeStats(t, [][]statEntry{
		{stat("user>>>alice>>>traffic>>>uplink", 5000)},
		{stat("user>>>alice>>>traffic>>>uplink", 5300)},
	})

	snap, err := api.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("GetTrafficSnapshot: %v", err)
	}
	if len(snap.ClientTraffics) != 0 {
		t.Fatalf("first snapshot reported %+v, want baseline only", snap.ClientTraffics[0])
	}
	if got := snap.Values["user>>>alice>>>traffic>>>uplink"]; got != 5000 {
		t.Fatalf("snapshot values = %d, want the raw counter 5000", got)
	}

	api.CommitTrafficSnapshot(snap)

	snap, err = api.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("GetTrafficSnapshot after commit: %v", err)
	}
	up, ok := clientUp(t, snap.ClientTraffics, "alice")
	if !ok || up != 300 {
		t.Fatalf("after baseline commit, delta = %d (found %v), want 300", up, ok)
	}
}

// TestTrafficSnapshotWithoutCommitReReports is the core of the outage fix:
// while reports fail, the baseline stays put and every later snapshot
// re-includes the unreported traffic, so a recovery report carries the
// whole accumulated amount — exactly once.
func TestTrafficSnapshotWithoutCommitReReports(t *testing.T) {
	api := startFakeStats(t, [][]statEntry{
		{stat("user>>>alice>>>traffic>>>uplink", 100)},
		{stat("user>>>alice>>>traffic>>>uplink", 260)},
		{stat("user>>>alice>>>traffic>>>uplink", 700)},
		{stat("user>>>alice>>>traffic>>>uplink", 750)},
	})

	base, err := api.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("GetTrafficSnapshot (baseline): %v", err)
	}
	api.CommitTrafficSnapshot(base)

	// Two polls whose uploads "fail": no commit happens between them.
	snap1, err := api.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("GetTrafficSnapshot (failed round 1): %v", err)
	}
	if up, _ := clientUp(t, snap1.ClientTraffics, "alice"); up != 160 {
		t.Fatalf("first unreported delta = %d, want 160", up)
	}
	snap2, err := api.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("GetTrafficSnapshot (failed round 2): %v", err)
	}
	// Not 440+160 and not just 440: the delta is recomputed from the
	// committed baseline, covering both failed rounds in one amount.
	if up, _ := clientUp(t, snap2.ClientTraffics, "alice"); up != 600 {
		t.Fatalf("accumulated delta after two failed rounds = %d, want 600", up)
	}

	// The master recovers and acknowledges this report.
	api.CommitTrafficSnapshot(snap2)

	snap3, err := api.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("GetTrafficSnapshot (after recovery): %v", err)
	}
	if up, _ := clientUp(t, snap3.ClientTraffics, "alice"); up != 50 {
		t.Fatalf("delta after recovery commit = %d, want only the new 50", up)
	}
}

// TestTrafficSnapshotRestoreBaselineSkipsBaselinePass covers an agent
// restart with a persisted baseline: a fresh XrayAPI with the restored
// values must diff against them immediately — the traffic that accrued
// while the agent was down is reported, not folded into a new baseline.
func TestTrafficSnapshotRestoreBaselineSkipsBaselinePass(t *testing.T) {
	api := startFakeStats(t, [][]statEntry{
		{stat("user>>>alice>>>traffic>>>uplink", 1250)},
	})
	api.RestoreTrafficBaseline(map[string]int64{
		"user>>>alice>>>traffic>>>uplink": 1000,
	})

	snap, err := api.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("GetTrafficSnapshot: %v", err)
	}
	up, ok := clientUp(t, snap.ClientTraffics, "alice")
	if !ok || up != 250 {
		t.Fatalf("restored-baseline delta = %d (found %v), want 250 with no baseline-only pass", up, ok)
	}
}

// TestTrafficSnapshotCountsAfterCounterReset: the core restarts and its
// counters start over below the committed baseline. The whole new counter
// is new traffic — never a negative delta — and after committing, normal
// diffing resumes in the new epoch.
func TestTrafficSnapshotCountsAfterCounterReset(t *testing.T) {
	api := startFakeStats(t, [][]statEntry{
		{stat("user>>>alice>>>traffic>>>uplink", 250)},
		{stat("user>>>alice>>>traffic>>>uplink", 300)},
	})
	api.RestoreTrafficBaseline(map[string]int64{
		"user>>>alice>>>traffic>>>uplink": 900,
	})

	snap, err := api.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("GetTrafficSnapshot (after reset): %v", err)
	}
	if up, _ := clientUp(t, snap.ClientTraffics, "alice"); up != 250 {
		t.Fatalf("delta after counter reset = %d, want 250 counted from zero", up)
	}
	api.CommitTrafficSnapshot(snap)

	snap, err = api.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("GetTrafficSnapshot (new epoch delta): %v", err)
	}
	if up, _ := clientUp(t, snap.ClientTraffics, "alice"); up != 50 {
		t.Fatalf("delta in new epoch = %d, want 50", up)
	}
}

// TestTrafficSnapshotNewStatCountsFromZero: a counter that appears after
// the baseline was committed (client's first use) is entirely new traffic.
func TestTrafficSnapshotNewStatCountsFromZero(t *testing.T) {
	api := startFakeStats(t, [][]statEntry{
		{
			stat("user>>>alice>>>traffic>>>uplink", 180),
			stat("user>>>bob>>>traffic>>>uplink", 4096),
			stat("inbound>>>in-443>>>traffic>>>downlink", 7000),
		},
	})
	api.RestoreTrafficBaseline(map[string]int64{
		"user>>>alice>>>traffic>>>uplink": 100,
	})

	snap, err := api.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("GetTrafficSnapshot: %v", err)
	}
	if up, _ := clientUp(t, snap.ClientTraffics, "alice"); up != 80 {
		t.Fatalf("alice delta = %d, want 80", up)
	}
	if up, _ := clientUp(t, snap.ClientTraffics, "bob"); up != 4096 {
		t.Fatalf("bob delta = %d, want his full counter 4096", up)
	}
	if len(snap.Traffics) != 1 || snap.Traffics[0].Tag != "in-443" || snap.Traffics[0].Down != 7000 {
		t.Fatalf("inbound deltas = %+v, want in-443 down 7000 from zero", snap.Traffics)
	}
}

// TestSnapshotDoesNotDisturbGetTraffic pins the range constraint: the
// panel path must keep its exact read-advances-baseline semantics even
// when snapshot calls interleave on the same client — the two baselines
// are fully independent.
func TestSnapshotDoesNotDisturbGetTraffic(t *testing.T) {
	api := startFakeStats(t, [][]statEntry{
		{stat("user>>>alice>>>traffic>>>uplink", 100)},
		{stat("user>>>alice>>>traffic>>>uplink", 160)},
		{stat("user>>>alice>>>traffic>>>uplink", 400)},
		{stat("user>>>alice>>>traffic>>>uplink", 460)},
	})

	// Panel poll 1: baseline only.
	if _, _, err := api.GetTraffic(); err != nil {
		t.Fatalf("GetTraffic (baseline): %v", err)
	}
	// Agent snapshot in between: must not move the panel's baseline.
	if _, err := api.GetTrafficSnapshot(); err != nil {
		t.Fatalf("GetTrafficSnapshot: %v", err)
	}
	// Panel poll 2: sees 400-100=300, including what accrued while the
	// snapshot was taken — proof the snapshot consumed nothing.
	_, clients, err := api.GetTraffic()
	if err != nil {
		t.Fatalf("GetTraffic: %v", err)
	}
	if up, _ := clientUp(t, clients, "alice"); up != 300 {
		t.Fatalf("panel delta after interleaved snapshot = %d, want 300", up)
	}
	if got := api.StatsLastValues["user>>>alice>>>traffic>>>uplink"]; got != 400 {
		t.Fatalf("panel baseline = %d, want 400 (advanced at read, as before)", got)
	}
	// Panel poll 3 advances per poll, exactly as it always has.
	_, clients, err = api.GetTraffic()
	if err != nil {
		t.Fatalf("GetTraffic: %v", err)
	}
	if up, _ := clientUp(t, clients, "alice"); up != 60 {
		t.Fatalf("panel delta = %d, want 60", up)
	}
}