From 17be1ad5aec88e30d467a7445e6c178a8be9a3a8 Mon Sep 17 00:00:00 2001 From: Muse Date: Tue, 20 Oct 2026 15:58:38 +0000 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E4=BA=8C=E5=BC=80=20agent=20?= =?UTF-8?q?=E6=B5=81=E9=87=8F=E5=A2=9E=E9=87=8F=E4=B8=A2=E5=A4=B1:=20?= =?UTF-8?q?=E5=9F=BA=E7=BA=BF=E6=94=B9=E4=B8=8A=E6=8A=A5=E7=A1=AE=E8=AE=A4?= =?UTF-8?q?=E5=90=8E=E6=8E=A8=E8=BF=9B=E5=B9=B6=E8=90=BD=E7=9B=98=E6=8C=81?= =?UTF-8?q?=E4=B9=85=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - xray 新增快照/确认接口 GetTrafficSnapshot/CommitTrafficSnapshot/ RestoreTrafficBaseline, 与面板 GetTraffic 的 StatsLastValues 完全隔离 - agent 改走快照: 仅主控 HTTP 200 后才推进基线, 失败轮增量由累计 计数做差在下轮自动补齐; 已确认基线原子落盘, agent 重启恢复续差 - 核心重启计数归零的从零规则保留; Onlines 仍每轮新鲜采集直报 - 新增 11 个测试覆盖断网恢复/连续失败补齐/重启恢复/计数归零/ 面板语义隔离, xray 与 agent 两包全绿 --- internal/agent/agent.go | 165 +++++++++-- internal/agent/agent_traffic_test.go | 322 +++++++++++++++++++++ internal/xray/api.go | 121 ++++++++ internal/xray/api_traffic_snapshot_test.go | 225 ++++++++++++++ main.go | 7 + 5 files changed, 818 insertions(+), 22 deletions(-) create mode 100644 internal/agent/agent_traffic_test.go create mode 100644 internal/xray/api_traffic_snapshot_test.go diff --git a/internal/agent/agent.go b/internal/agent/agent.go index bddd6ab..f8333e3 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -43,6 +43,12 @@ type Config struct { // config. Needed when agent and master share a host (e.g. testing), // otherwise both xrays fight over the default 62789. 0 = use master's. ApiPort int `json:"api_port"` + // BaselinePath is where the last master-acknowledged traffic counters + // are persisted, so an agent restart keeps diffing against them + // instead of swallowing unreported traffic into a fresh baseline. + // Empty: main derives it next to the agent config file; the agent + // falls back to the bin folder. + BaselinePath string `json:"baseline_path,omitempty"` } // Agent represents the node agent @@ -54,6 +60,7 @@ type Agent struct { xrayProc *xray.Process xrayAPI xray.XrayAPI apiPort int + baselineLoaded bool } // New creates a new agent instance @@ -331,36 +338,73 @@ type userEntry struct { Down int64 `json:"down"` } -// collectAndReportTraffic queries the local xray stats API for per-poll -// deltas and posts them to the master. The XrayAPI instance is kept across -// polls so delta baselines survive; xray restarts are handled (counters -// resetting to zero are treated as fresh). -func (a *Agent) collectAndReportTraffic() error { - if a.xrayProc == nil || !a.xrayProc.IsRunning() { - return nil - } - if a.apiPort <= 0 { - return nil - } - if err := a.xrayAPI.Init(a.apiPort); err != nil { - return err - } - defer a.xrayAPI.Close() +// statsClient is the slice of the xray stats API the traffic reporter +// uses. The snapshot/commit pair is the loss-proofing contract: snapshots +// never advance the baseline, commits do. +type statsClient interface { + GetTrafficSnapshot() (*xray.TrafficSnapshot, error) + CommitTrafficSnapshot(*xray.TrafficSnapshot) + RestoreTrafficBaseline(map[string]int64) + GetOnlineUsers() ([]xray.OnlineUser, error) +} - traffics, clientTraffics, err := a.xrayAPI.GetTraffic() +// trafficBaselineFile is the on-disk form of the last committed baseline. +type trafficBaselineFile struct { + Version int `json:"version"` + Values map[string]int64 `json:"values"` +} + +// baselinePath resolves where the acknowledged traffic baseline lives. +func (a *Agent) baselinePath() string { + if a.config.BaselinePath != "" { + return a.config.BaselinePath + } + return filepath.Join(config.GetBinFolderPath(), "agent-traffic-baseline.json") +} + +// loadTrafficBaseline reads a persisted baseline. A missing file is a +// normal first run; a corrupt one is reported so the caller can log it — +// both yield no values and the reporter falls back to a baseline-only +// first snapshot rather than crashing the agent. +func loadTrafficBaseline(path string) (map[string]int64, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, err + } + var f trafficBaselineFile + if err := json.Unmarshal(data, &f); err != nil { + return nil, fmt.Errorf("parse traffic baseline: %w", err) + } + return f.Values, nil +} + +// saveTrafficBaseline persists the baseline atomically (write-then-rename) +// so a crash mid-write can never leave a torn file behind. +func saveTrafficBaseline(path string, values map[string]int64) error { + data, err := json.Marshal(trafficBaselineFile{Version: 1, Values: values}) if err != nil { return err } - onlineUsers, _ := a.xrayAPI.GetOnlineUsers() + tmp := path + ".tmp" + if err := os.WriteFile(tmp, data, 0o600); err != nil { + return err + } + return os.Rename(tmp, path) +} +// buildTrafficReport shapes a snapshot plus the freshly polled online +// users into the wire report, dropping zero-traffic entries. Onlines are +// point-in-time data: they ride along with each report and are never part +// of the baseline bookkeeping. +func buildTrafficReport(snap *xray.TrafficSnapshot, onlineUsers []xray.OnlineUser) trafficReport { report := trafficReport{} - for _, t := range traffics { + for _, t := range snap.Traffics { if t == nil || !t.IsInbound || t.Up+t.Down <= 0 { continue } report.Inbounds = append(report.Inbounds, trafficEntry{Tag: t.Tag, Up: t.Up, Down: t.Down}) } - for _, ct := range clientTraffics { + for _, ct := range snap.ClientTraffics { if ct == nil || ct.Email == "" || ct.Up+ct.Down <= 0 { continue } @@ -369,10 +413,16 @@ func (a *Agent) collectAndReportTraffic() error { for _, u := range onlineUsers { report.Onlines = append(report.Onlines, u.Email) } - if len(report.Inbounds) == 0 && len(report.Users) == 0 && len(report.Onlines) == 0 { - return nil - } + return report +} +func (r trafficReport) empty() bool { + return len(r.Inbounds) == 0 && len(r.Users) == 0 && len(r.Onlines) == 0 +} + +// postTrafficReport sends one report to the master; only an HTTP 200 +// counts as acknowledged. +func (a *Agent) postTrafficReport(report trafficReport) error { url := fmt.Sprintf("%s/panel/api/agent/traffic", a.config.MasterURL) body, _ := json.Marshal(report) req, err := http.NewRequest("POST", url, bytes.NewReader(body)) @@ -393,6 +443,77 @@ func (a *Agent) collectAndReportTraffic() error { return nil } +// reportTrafficWith runs one snapshot → report → commit cycle. +// +// The committed baseline advances ONLY after the master acknowledges the +// report (HTTP 200). Core counters are cumulative, so when a report fails +// nothing is lost: the next snapshot diffs against the same baseline and +// re-includes everything unreported, exactly once. The acknowledged +// baseline is persisted after every commit, and restored before the first +// snapshot of a restarted agent, so an agent restart cannot swallow the +// unreported tail into a fresh baseline either. +// +// A poll with nothing to send still commits: it establishes the very +// first baseline, and keeps the persisted copy fresh. +func (a *Agent) reportTrafficWith(sc statsClient, post func(trafficReport) error) error { + if !a.baselineLoaded { + a.baselineLoaded = true + if values, err := loadTrafficBaseline(a.baselinePath()); err != nil { + if !os.IsNotExist(err) { + logger.Warningf("Traffic baseline at %s unusable (%v); starting from a fresh baseline", a.baselinePath(), err) + } + } else if len(values) > 0 { + sc.RestoreTrafficBaseline(values) + } + } + + snap, err := sc.GetTrafficSnapshot() + if err != nil { + return err + } + onlineUsers, _ := sc.GetOnlineUsers() + + report := buildTrafficReport(snap, onlineUsers) + if !report.empty() { + if err := post(report); err != nil { + // No commit: the baseline stays where the master last + // acknowledged it, and these deltas are reported again + // (grown) on the next poll. + return err + } + } + + sc.CommitTrafficSnapshot(snap) + if err := saveTrafficBaseline(a.baselinePath(), snap.BaselineValues()); err != nil { + // The in-memory commit stands — rolling it back would re-report + // traffic the master already counted. The stale file self-heals + // on the next successful poll; only a crash before that could + // re-report a single interval. + logger.Warningf("Failed to persist traffic baseline: %v", err) + return err + } + return nil +} + +// collectAndReportTraffic snapshots the local xray counters and reports +// the deltas since the last master-acknowledged baseline. The XrayAPI +// instance is kept across polls so baselines survive; xray restarts are +// handled (counters resetting to zero are treated as fresh). +func (a *Agent) collectAndReportTraffic() error { + if a.xrayProc == nil || !a.xrayProc.IsRunning() { + return nil + } + if a.apiPort <= 0 { + return nil + } + if err := a.xrayAPI.Init(a.apiPort); err != nil { + return err + } + defer a.xrayAPI.Close() + + return a.reportTrafficWith(&a.xrayAPI, a.postTrafficReport) +} + // LoadConfig loads agent config from file func LoadConfig(path string) (*Config, error) { data, err := os.ReadFile(path) diff --git a/internal/agent/agent_traffic_test.go b/internal/agent/agent_traffic_test.go new file mode 100644 index 0000000..983a0a3 --- /dev/null +++ b/internal/agent/agent_traffic_test.go @@ -0,0 +1,322 @@ +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) + } +} diff --git a/internal/xray/api.go b/internal/xray/api.go index fe0ec56..6636fa1 100644 --- a/internal/xray/api.go +++ b/internal/xray/api.go @@ -61,6 +61,19 @@ type XrayAPI struct { grpcClient *grpc.ClientConn isConnected bool StatsLastValues map[string]int64 + + // committedStats is the agent reporting path's last ACKNOWLEDGED + // baseline: the raw core counters as of the most recent snapshot the + // master confirmed receiving (HTTP 200). It is deliberately separate + // from StatsLastValues — GetTraffic (the panel's local accounting) + // advances its baseline at read time, while the agent path may only + // advance on CommitTrafficSnapshot, so a failed upload never loses + // the deltas it carried. committedReady is false until a baseline has + // been established (first snapshot committed, or one restored from + // disk); while false, snapshots report no deltas, mirroring + // GetTraffic's first-poll rule. + committedStats map[string]int64 + committedReady bool } // Init connects to the Xray API server. @@ -650,6 +663,114 @@ func (x *XrayAPI) GetTraffic() ([]*Traffic, []*ClientTraffic, error) { return mapToSlice(tagTrafficMap), mapToSlice(emailTrafficMap), nil } +// TrafficSnapshot is one read of the core counters taken for the agent +// reporting path: the deltas accrued since the last COMMITTED baseline, +// plus Values, the raw cumulative counters the deltas were computed from. +// +// Unlike GetTraffic — which advances its baseline as it reads, the right +// contract for the panel's local accounting job — a snapshot changes no +// state. The caller reports the deltas and only then calls +// CommitTrafficSnapshot; if the report fails, the next snapshot recomputes +// against the same committed baseline and the unreported traffic is +// included again, because the core counters are cumulative and are never +// reset by polling. +type TrafficSnapshot struct { + Traffics []*Traffic + ClientTraffics []*ClientTraffic + // Values holds every counter returned by the core at snapshot time, + // keyed by stat name. CommitTrafficSnapshot installs exactly this map + // as the new committed baseline. + Values map[string]int64 +} + +// BaselineValues returns a copy of the baseline CommitTrafficSnapshot +// would install for this snapshot, so callers can persist it (the agent +// writes it to disk before/while committing). +func (s *TrafficSnapshot) BaselineValues() map[string]int64 { + out := make(map[string]int64, len(s.Values)) + for k, v := range s.Values { + out[k] = v + } + return out +} + +// GetTrafficSnapshot reads the core counters and computes deltas against +// the committed baseline WITHOUT advancing it. Per-stat rules match +// GetTraffic exactly: while no baseline has been established the snapshot +// reports no deltas (it only proposes one), a counter the baseline has +// never seen counts from zero, and a counter below its baseline (the core +// restarted) counts from zero as well. +func (x *XrayAPI) GetTrafficSnapshot() (*TrafficSnapshot, error) { + if x.grpcClient == nil { + return nil, common.NewError("xray api is not initialized") + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + + stats, err := x.queryStats(ctx) + if err != nil { + logger.Debug("Failed to query Xray stats:", err) + return nil, err + } + + snap := &TrafficSnapshot{Values: make(map[string]int64, len(stats))} + tagTrafficMap := make(map[string]*Traffic) + emailTrafficMap := make(map[string]*ClientTraffic) + + for _, stat := range stats { + snap.Values[stat.Name] = stat.Value + if !x.committedReady { + continue + } + lastValue, ok := x.committedStats[stat.Name] + if !ok || stat.Value < lastValue { + lastValue = 0 + } + value := stat.Value - lastValue + if matches := trafficRegex.FindStringSubmatch(stat.Name); len(matches) == 4 { + processTraffic(matches, value, tagTrafficMap) + } else if matches := clientTrafficRegex.FindStringSubmatch(stat.Name); len(matches) == 3 { + processClientTraffic(matches, value, emailTrafficMap) + } + } + + snap.Traffics = mapToSlice(tagTrafficMap) + snap.ClientTraffics = mapToSlice(emailTrafficMap) + return snap, nil +} + +// CommitTrafficSnapshot marks a snapshot as acknowledged by the master: +// its Values become the committed baseline that later snapshots diff +// against. Installing only the counters the core currently reports also +// retires baselines for deleted inbounds/clients, the same end state +// GetTraffic's pruning converges to. +func (x *XrayAPI) CommitTrafficSnapshot(s *TrafficSnapshot) { + if s == nil { + return + } + x.committedStats = s.BaselineValues() + x.committedReady = true +} + +// RestoreTrafficBaseline installs a committed baseline recovered from +// persistent storage (the agent saves one after every acknowledged +// report), so a restarted agent keeps diffing against the last values the +// master confirmed instead of swallowing the unreported tail into a +// fresh baseline. An empty map leaves the baseline unestablished: the +// next snapshot is the baseline-only first pass, exactly like a first +// run with no state file. +func (x *XrayAPI) RestoreTrafficBaseline(values map[string]int64) { + if len(values) == 0 { + return + } + x.committedStats = make(map[string]int64, len(values)) + for k, v := range values { + x.committedStats[k] = v + } + x.committedReady = true +} + // OnlineIP is one source address of a live connection, with the unix time (seconds) // the core last dispatched a link from it. type OnlineIP struct { diff --git a/internal/xray/api_traffic_snapshot_test.go b/internal/xray/api_traffic_snapshot_test.go new file mode 100644 index 0000000..3b45d80 --- /dev/null +++ b/internal/xray/api_traffic_snapshot_test.go @@ -0,0 +1,225 @@ +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) + } +} diff --git a/main.go b/main.go index 1222d4a..a34df99 100644 --- a/main.go +++ b/main.go @@ -11,6 +11,7 @@ import ( _ "net/http/pprof" "os" "os/signal" + "path/filepath" "syscall" _ "unsafe" @@ -84,6 +85,12 @@ func runAgent(args []string) { log.Fatalf("Failed to load agent config from %s: %v", configPath, err) } + // The acknowledged traffic baseline persists next to the agent config + // (its state home, e.g. /etc/x-ui), unless the config names a path. + if cfg.BaselinePath == "" { + cfg.BaselinePath = filepath.Join(filepath.Dir(configPath), "agent-traffic-baseline.json") + } + // Initialize logger logger.InitLogger(logging.INFO) log.Printf("Starting 3x-ui agent (node: %s)", cfg.NodeName)