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

package agent

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"sync/atomic"
	"testing"
)

// V2.1.0 probe-only tests: the mode gates both core branches of the
// config poll (no fetch, no start, no switch latch), retires a core
// that was running when the mode turns on, and self-heals back to
// the core when the mode turns off again.

// probeMaster is a fake master whose config response is assembled
// from a switchable probe_only flag plus the usual sing-box payload.
type probeMaster struct {
	probeOnly  atomic.Bool
	trafficHit atomic.Int32
	srv        *httptest.Server
}

func newProbeMaster(t *testing.T) *probeMaster {
	t.Helper()
	m := &probeMaster{}
	sum := sha256.Sum256([]byte(sbCfgA))
	mux := http.NewServeMux()
	mux.HandleFunc("/panel/api/agent/config", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		body := map[string]any{
			"singbox_config":      json.RawMessage(sbCfgA),
			"singbox_config_hash": hex.EncodeToString(sum[:]),
			"probe_tasks":         []any{},
			// The real master always sends the flag (only a pre-V2.1.0
			// master omits it); the fixture mirrors that.
			"probe_only": m.probeOnly.Load(),
		}
		_ = json.NewEncoder(w).Encode(body)
	})
	mux.HandleFunc("/panel/api/agent/traffic", func(w http.ResponseWriter, r *http.Request) {
		m.trafficHit.Add(1)
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"success":true}`))
	})
	m.srv = httptest.NewServer(mux)
	t.Cleanup(m.srv.Close)
	return m
}

func (m *probeMaster) agent(cfg *Config) *Agent {
	if cfg == nil {
		cfg = &Config{}
	}
	cfg.MasterURL = m.srv.URL
	cfg.Token = "tok"
	cfg.NodeName = "n"
	return New(cfg)
}

// The master's flag alone (no local preset) must gate the poll: the
// served sing-box payload is ignored entirely.
func TestProbeOnlyMasterFlagGatesApply(t *testing.T) {
	setupFakeSingbox(t)
	m := newProbeMaster(t)
	m.probeOnly.Store(true)
	a := m.agent(nil)
	t.Cleanup(func() { a.stopXray(); a.stopSingbox() })

	if err := a.fetchAndApplyConfig(); err != nil {
		t.Fatalf("fetchAndApplyConfig: %v", err)
	}
	if !a.probeOnlyActive() {
		t.Fatal("probe-only must be active from the master flag")
	}
	if a.sbSwitched {
		t.Fatal("switchover latched under probe-only")
	}
	if a.singboxRunning() {
		t.Fatal("sing-box started under probe-only")
	}
	if a.lastSingboxHash != "" {
		t.Fatalf("lastSingboxHash = %q, want untouched (empty)", a.lastSingboxHash)
	}
	if _, err := os.Stat(singboxConfigPath()); !os.IsNotExist(err) {
		t.Fatalf("applied config file exists under probe-only (stat err = %v)", err)
	}
}

// The local agent.json preset gates even against a master that
// still serves a payload (and predates the flag): with no core
// binary on disk, not even a core fetch may be attempted.
func TestProbeOnlyLocalPresetSkipsCoreFetch(t *testing.T) {
	binDir := t.TempDir() // deliberately no fake sing-box binary
	t.Setenv("XUI_BIN_FOLDER", binDir)

	var hits atomic.Int32
	coreSrv := fakeCoreServer(t, &hits, http.StatusOK)
	defer coreSrv.Close()
	useCoreBase(t, coreSrv.URL)

	m := newProbeMaster(t) // probe_only absent from its responses
	a := m.agent(&Config{ProbeOnly: true})
	t.Cleanup(func() { a.stopXray(); a.stopSingbox() })

	if err := a.fetchAndApplyConfig(); err != nil {
		t.Fatalf("fetchAndApplyConfig: %v", err)
	}
	if hits.Load() != 0 {
		t.Fatalf("core fetch attempted %d times under local probe-only preset", hits.Load())
	}
	if a.sbSwitched || a.singboxRunning() {
		t.Fatal("core engaged under local probe-only preset")
	}
}

// Turning the flag on while sing-box is running retires the core
// exactly once; turning it off again reactivates the core from the
// unchanged hash (the poll's not-running path), without the
// bookkeeping having been dirtied in between.
func TestProbeOnlyRetiresRunningSingboxAndHealsBack(t *testing.T) {
	setupFakeSingbox(t)
	m := newProbeMaster(t)
	a := m.agent(nil)
	t.Cleanup(func() { a.stopXray(); a.stopSingbox() })

	if err := a.fetchAndApplyConfig(); err != nil {
		t.Fatalf("initial fetch: %v", err)
	}
	if !a.singboxRunning() || !a.sbSwitched {
		t.Fatal("precondition: sing-box must be running after the normal poll")
	}
	hashBefore := a.lastSingboxHash

	m.probeOnly.Store(true)
	if err := a.fetchAndApplyConfig(); err != nil {
		t.Fatalf("probe-only fetch: %v", err)
	}
	if a.singboxRunning() {
		t.Fatal("sing-box must be retired when probe-only turns on")
	}
	if !a.sbSwitched {
		t.Fatal("sbSwitched must stay latched (bookkeeping untouched by probe-only)")
	}
	if a.lastSingboxHash != hashBefore {
		t.Fatalf("lastSingboxHash dirtied: %q -> %q", hashBefore, a.lastSingboxHash)
	}
	// A second probe-only poll is a no-op, not an error.
	if err := a.fetchAndApplyConfig(); err != nil {
		t.Fatalf("second probe-only fetch: %v", err)
	}

	m.probeOnly.Store(false)
	if err := a.fetchAndApplyConfig(); err != nil {
		t.Fatalf("heal-back fetch: %v", err)
	}
	if !a.singboxRunning() {
		t.Fatal("sing-box must restart once probe-only turns off (hash unchanged)")
	}
}

// Same retirement for a legacy xray core that predates the
// switchover: probe-only stops it and does not start sing-box.
func TestProbeOnlyRetiresLegacyXray(t *testing.T) {
	binDir := setupFakeSingbox(t)
	setupFakeXray(t, binDir)
	m := newProbeMaster(t)
	m.probeOnly.Store(true)
	a := m.agent(nil)
	t.Cleanup(func() { a.stopXray(); a.stopSingbox() })
	startLegacyXray(t, a)

	if err := a.fetchAndApplyConfig(); err != nil {
		t.Fatalf("fetchAndApplyConfig: %v", err)
	}
	if a.xrayProc != nil {
		t.Fatal("xray core must be retired under probe-only")
	}
	if a.singboxRunning() || a.sbSwitched {
		t.Fatal("sing-box must not engage under probe-only")
	}
}

// The traffic cycle is an explicit no-op under probe-only: even
// with a core process somehow alive, nothing is collected or posted.
func TestProbeOnlyTrafficCycleIsNoop(t *testing.T) {
	setupFakeSingbox(t)
	m := newProbeMaster(t)
	a := m.agent(&Config{ProbeOnly: true})
	t.Cleanup(func() { a.stopXray(); a.stopSingbox() })

	// Put a core under the agent directly (bypassing the poll gate)
	// to prove the traffic cycle itself refuses to touch it.
	if err := a.applySingboxConfig([]byte(sbCfgA)); err != nil {
		t.Fatalf("applySingboxConfig: %v", err)
	}
	if !a.singboxRunning() {
		t.Fatal("precondition: fake sing-box must be running")
	}
	if err := a.collectAndReportTraffic(); err != nil {
		t.Fatalf("collectAndReportTraffic under probe-only: %v", err)
	}
	if n := m.trafficHit.Load(); n != 0 {
		t.Fatalf("traffic posted %d times under probe-only", n)
	}
}