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

package agent

import (
	"context"
	"errors"
	"testing"

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

	"google.golang.org/protobuf/encoding/protowire"
)

// buildStatsResponse hand-encodes a QueryStatsResponse carrying the
// given counters, the same wire shape the real core returns.
func buildStatsResponse(counters map[string]int64) []byte {
	var out []byte
	for name, value := range counters {
		var stat []byte
		stat = protowire.AppendTag(stat, 1, protowire.BytesType)
		stat = protowire.AppendString(stat, name)
		stat = protowire.AppendTag(stat, 2, protowire.VarintType)
		stat = protowire.AppendVarint(stat, uint64(value))
		out = protowire.AppendTag(out, 1, protowire.BytesType)
		out = protowire.AppendBytes(out, stat)
	}
	return out
}

func TestParseSingboxStatsResponse(t *testing.T) {
	raw := buildStatsResponse(map[string]int64{
		"inbound>>>ss22>>>traffic>>>uplink":           78,
		"user>>>bob@example.com>>>traffic>>>downlink": 573,
	})
	entries, err := parseSingboxStatsResponse(raw)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	if len(entries) != 2 {
		t.Fatalf("entries = %d, want 2", len(entries))
	}
	got := map[string]int64{}
	for _, e := range entries {
		got[e.Name] = e.Value
	}
	if got["inbound>>>ss22>>>traffic>>>uplink"] != 78 || got["user>>>bob@example.com>>>traffic>>>downlink"] != 573 {
		t.Fatalf("counters = %v", got)
	}
	if _, err := parseSingboxStatsResponse([]byte{0xff, 0xff}); err == nil {
		t.Fatal("garbage input must error")
	}
}

func TestEncodeSingboxQuery(t *testing.T) {
	b := encodeSingboxQuery()
	var patterns []string
	for len(b) > 0 {
		num, typ, n := protowire.ConsumeTag(b)
		if n < 0 || num != 3 || typ != protowire.BytesType {
			t.Fatalf("unexpected field num=%d typ=%d", num, typ)
		}
		b = b[n:]
		v, n := protowire.ConsumeString(b)
		if n < 0 {
			t.Fatal("bad string field")
		}
		patterns = append(patterns, v)
		b = b[n:]
	}
	if len(patterns) != 3 {
		t.Fatalf("patterns = %v", patterns)
	}
}

// scriptedStats drives the tracker with scripted counters and a
// scripted epoch, following the fakeStatsClient pattern of the xray
// path: the delta math contract is what's under test here; the wire
// itself is covered by the parse test and the real-binary E2E.
type scriptedStats struct {
	counters map[string]int64
	epoch    string
	err      error
}

func (s *scriptedStats) client() *singboxStatsClient {
	c := newSingboxStatsClient(func() string { return s.epoch })
	c.queryFn = func(context.Context) ([]singboxStatEntry, error) {
		if s.err != nil {
			return nil, s.err
		}
		out := make([]singboxStatEntry, 0, len(s.counters))
		for name, value := range s.counters {
			out = append(out, singboxStatEntry{Name: name, Value: value})
		}
		return out, nil
	}
	return c
}

const (
	sbUserUp   = "user>>>bob@example.com>>>traffic>>>uplink"
	sbUserDown = "user>>>bob@example.com>>>traffic>>>downlink"
	sbInUp     = "inbound>>>ss22>>>traffic>>>uplink"
)

func snapUserDown(snap *xray.TrafficSnapshot, email string) int64 {
	for _, ct := range snap.ClientTraffics {
		if ct.Email == email {
			return ct.Down
		}
	}
	return 0
}

func TestSingboxSnapshotBaselineThenDeltas(t *testing.T) {
	s := &scriptedStats{epoch: "sing-box:1:100", counters: map[string]int64{sbUserDown: 500, sbInUp: 40}}
	c := s.client()

	snap, err := c.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("snapshot 1: %v", err)
	}
	if len(snap.ClientTraffics) != 0 || len(snap.Traffics) != 0 {
		t.Fatal("first snapshot of an epoch must be baseline-only")
	}
	c.CommitTrafficSnapshot(snap)

	s.counters[sbUserDown] = 900
	s.counters[sbInUp] = 55
	snap, err = c.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("snapshot 2: %v", err)
	}
	if got := snapUserDown(snap, "bob@example.com"); got != 400 {
		t.Fatalf("user delta = %d, want 400", got)
	}
	c.CommitTrafficSnapshot(snap)

	// No movement: zero deltas are still shaped but empty of content.
	snap, err = c.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("snapshot 3: %v", err)
	}
	if got := snapUserDown(snap, "bob@example.com"); got != 0 {
		t.Fatalf("idle delta = %d, want 0", got)
	}
}

func TestSingboxSnapshotEpochChangeCountsInFull(t *testing.T) {
	s := &scriptedStats{epoch: "sing-box:1:100", counters: map[string]int64{sbUserDown: 500}}
	c := s.client()
	snap, err := c.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("snapshot 1: %v", err)
	}
	c.CommitTrafficSnapshot(snap)

	// Core restarted: counters reset and grow again. The new
	// instance's counters started at zero, so they count in full
	// (NexusGate's ApplyDelta rule) — never as a diff against the
	// dead instance's baseline, and never dropped.
	s.epoch = "sing-box:2:200"
	s.counters[sbUserDown] = 120
	snap, err = c.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("snapshot after restart: %v", err)
	}
	if got := snapUserDown(snap, "bob@example.com"); got != 120 {
		t.Fatalf("post-restart delta = %d, want 120 (counted in full)", got)
	}
	c.CommitTrafficSnapshot(snap)

	s.counters[sbUserDown] = 150
	snap, err = c.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("snapshot 3: %v", err)
	}
	if got := snapUserDown(snap, "bob@example.com"); got != 30 {
		t.Fatalf("post-rebase delta = %d, want 30", got)
	}
}

func TestSingboxSnapshotEpochUnstableFails(t *testing.T) {
	epochs := []string{"sing-box:1:100", "sing-box:2:200"}
	calls := 0
	c := newSingboxStatsClient(func() string {
		e := epochs[min(calls, len(epochs)-1)]
		calls++
		return e
	})
	c.queryFn = func(context.Context) ([]singboxStatEntry, error) {
		return []singboxStatEntry{{Name: sbUserDown, Value: 10}}, nil
	}
	if _, err := c.GetTrafficSnapshot(); !errors.Is(err, errSingboxEpochChanged) {
		t.Fatalf("err = %v, want errSingboxEpochChanged", err)
	}
}

func TestSingboxSnapshotCounterResetCountsFromZero(t *testing.T) {
	s := &scriptedStats{epoch: "sing-box:1:100", counters: map[string]int64{sbUserDown: 500, sbUserUp: 10}}
	c := s.client()
	snap, err := c.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("snapshot 1: %v", err)
	}
	c.CommitTrafficSnapshot(snap)

	// Same epoch but a counter moved backwards (core reset its stats
	// without a process restart): counts from zero, like xray.
	s.counters[sbUserDown] = 60
	s.counters[sbUserUp] = 4 // unseen-before direction appears too
	snap, err = c.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("snapshot 2: %v", err)
	}
	if got := snapUserDown(snap, "bob@example.com"); got != 60 {
		t.Fatalf("reset counter delta = %d, want 60", got)
	}
}

func TestSingboxRestoreBaselineEpochGate(t *testing.T) {
	s := &scriptedStats{epoch: "sing-box:7:700", counters: map[string]int64{sbUserDown: 800}}
	c := s.client()

	// Baseline from a dead instance: ignored, snapshot rebases.
	c.RestoreTrafficBaseline(map[string]int64{sbUserDown: 300}, "sing-box:6:600")
	snap, err := c.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("snapshot: %v", err)
	}
	if len(snap.ClientTraffics) != 0 {
		t.Fatal("stale-epoch baseline must not produce deltas")
	}

	// Baseline from the live instance: deltas continue from it.
	c2 := s.client()
	c2.RestoreTrafficBaseline(map[string]int64{sbUserDown: 300}, "sing-box:7:700")
	snap, err = c2.GetTrafficSnapshot()
	if err != nil {
		t.Fatalf("snapshot 2: %v", err)
	}
	if got := snapUserDown(snap, "bob@example.com"); got != 500 {
		t.Fatalf("restored delta = %d, want 500", got)
	}
}