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

// Package agent implements the lightweight 3x-ui node agent.
// In agent mode, the node does not run the web panel or database.
// It registers with the master via token, receives xray config, manages the
// xray process, and reports traffic back to the master.
package agent

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"time"

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

const (
	// Default poll interval for config updates
	defaultPollInterval = 30 * time.Second
	// Heartbeat interval
	defaultHeartbeatInterval = 60 * time.Second
	// Traffic report interval (piggybacks the heartbeat tick)
	defaultTrafficInterval = 60 * time.Second
)

// DefaultBinFolder is used when XUI_BIN_FOLDER is unset. It matches the
// panel's directory layout: xray binary, geo files and config.json live here.
const DefaultBinFolder = "/usr/local/x-ui/bin"

// Config holds the agent configuration
type Config struct {
	MasterURL    string        `json:"master_url"`
	Token        string        `json:"token"`
	NodeName     string        `json:"node_name"`
	PollInterval time.Duration `json:"poll_interval"`
	// ApiPort overrides the xray stats API inbound port from the master's
	// 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
type Agent struct {
	config         *Config
	httpClient     *http.Client
	stopCh         chan struct{}
	lastConfigHash string
	xrayProc       *xray.Process
	xrayAPI        xray.XrayAPI
	apiPort        int
	baselineLoaded bool
}

// New creates a new agent instance
func New(cfg *Config) *Agent {
	if cfg.PollInterval == 0 {
		cfg.PollInterval = defaultPollInterval
	}
	return &Agent{
		config: cfg,
		httpClient: &http.Client{
			Timeout: 30 * time.Second,
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: false},
			},
		},
		stopCh: make(chan struct{}),
	}
}

// Register registers the agent with the master
func (a *Agent) Register() error {
	url := fmt.Sprintf("%s/panel/api/agent/register", a.config.MasterURL)
	payload := map[string]string{
		"token": a.config.Token,
		"name":  a.config.NodeName,
	}
	body, _ := json.Marshal(payload)
	req, err := http.NewRequest("POST", url, bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+a.config.Token)

	resp, err := a.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("register failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		data, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("register failed with status %d: %s", resp.StatusCode, string(data))
	}

	logger.Info("Agent registered successfully with master")
	return nil
}

// Run starts the agent main loop
func (a *Agent) Run() error {
	logger.Info("Starting 3x-ui agent...")
	logger.Infof("Master: %s, Node: %s", a.config.MasterURL, a.config.NodeName)

	// Initial registration
	if err := a.Register(); err != nil {
		return fmt.Errorf("initial registration failed: %w", err)
	}

	// Start heartbeat ticker
	heartbeatTicker := time.NewTicker(defaultHeartbeatInterval)
	defer heartbeatTicker.Stop()

	// Start config poll ticker
	pollTicker := time.NewTicker(a.config.PollInterval)
	defer pollTicker.Stop()

	// Start traffic report ticker
	trafficTicker := time.NewTicker(defaultTrafficInterval)
	defer trafficTicker.Stop()

	// Initial config fetch
	if err := a.fetchAndApplyConfig(); err != nil {
		logger.Warningf("Initial config fetch failed: %v", err)
	}

	for {
		select {
		case <-a.stopCh:
			logger.Info("Agent stopping...")
			a.stopXray()
			return nil
		case <-heartbeatTicker.C:
			if err := a.sendHeartbeat(); err != nil {
				logger.Warningf("Heartbeat failed: %v", err)
			}
		case <-pollTicker.C:
			if err := a.fetchAndApplyConfig(); err != nil {
				logger.Warningf("Config fetch failed: %v", err)
			}
		case <-trafficTicker.C:
			if err := a.collectAndReportTraffic(); err != nil {
				logger.Warningf("Traffic report failed: %v", err)
			}
		}
	}
}

// Stop stops the agent
func (a *Agent) Stop() {
	close(a.stopCh)
}

// sendHeartbeat sends a heartbeat to the master
func (a *Agent) sendHeartbeat() error {
	url := fmt.Sprintf("%s/panel/api/agent/heartbeat", a.config.MasterURL)
	payload := map[string]interface{}{
		"token":     a.config.Token,
		"timestamp": time.Now().Unix(),
	}
	body, _ := json.Marshal(payload)
	req, err := http.NewRequest("POST", url, bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+a.config.Token)

	resp, err := a.httpClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("heartbeat failed with status %d", resp.StatusCode)
	}
	return nil
}

// fetchAndApplyConfig fetches the xray config from master and applies it if changed
func (a *Agent) fetchAndApplyConfig() error {
	url := fmt.Sprintf("%s/panel/api/agent/config", a.config.MasterURL)
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "Bearer "+a.config.Token)

	resp, err := a.httpClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		data, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("config fetch failed with status %d: %s", resp.StatusCode, string(data))
	}

	var result struct {
		ConfigHash string          `json:"config_hash"`
		XrayConfig json.RawMessage `json:"xray_config"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return err
	}

	// Skip if config hasn't changed
	if result.ConfigHash == a.lastConfigHash && result.ConfigHash != "" {
		return nil
	}

	logger.Info("New xray config received, applying...")
	if err := a.applyXrayConfig(result.XrayConfig); err != nil {
		return fmt.Errorf("failed to apply xray config: %w", err)
	}

	a.lastConfigHash = result.ConfigHash
	logger.Info("Xray config applied successfully")
	return nil
}

// ensureXrayBinary makes sure the xray binary exists at the panel-layout path.
// The install doc puts it at /usr/local/x-ui/bin/xray; the panel's process
// manager looks for the arch-specific name, so link one to the other.
func ensureXrayBinary() error {
	binPath := xray.GetBinaryPath()
	if _, err := os.Stat(binPath); err == nil {
		return nil
	}
	plain := filepath.Join(config.GetBinFolderPath(), "xray")
	if _, err := os.Stat(plain); err != nil {
		return fmt.Errorf("xray binary not found at %s or %s", plain, binPath)
	}
	// Best effort: link the arch-specific name the process manager expects.
	if err := os.Symlink(plain, binPath); err != nil && !os.IsExist(err) {
		return fmt.Errorf("failed to link xray binary: %w", err)
	}
	return nil
}

// applyXrayConfig writes the xray config and (re)starts xray through the
// panel's process manager.
func (a *Agent) applyXrayConfig(configJSON json.RawMessage) error {
	if err := ensureXrayBinary(); err != nil {
		return err
	}
	var cfg xray.Config
	if err := json.Unmarshal(configJSON, &cfg); err != nil {
		return fmt.Errorf("invalid xray config: %w", err)
	}
	// Same-host override: rewrite the stats API inbound port so the agent's
	// xray doesn't fight the master's xray over the default port.
	if a.config.ApiPort > 0 {
		for i := range cfg.InboundConfigs {
			if cfg.InboundConfigs[i].Tag == "api" {
				logger.Infof("Overriding xray API port %d -> %d (agent config)",
					cfg.InboundConfigs[i].Port, a.config.ApiPort)
				cfg.InboundConfigs[i].Port = a.config.ApiPort
			}
		}
		// Also move the prometheus metrics listener (default 127.0.0.1:11111)
		// which would collide with the master's xray on the same host.
		if len(cfg.Metrics) > 0 {
			var m map[string]any
			if err := json.Unmarshal(cfg.Metrics, &m); err == nil {
				if ls, ok := m["listen"].(string); ok && ls != "" {
					newListen := fmt.Sprintf("127.0.0.1:%d", a.config.ApiPort+1)
					logger.Infof("Overriding xray metrics listen %s -> %s (agent config)", ls, newListen)
					m["listen"] = newListen
					if raw, err := json.Marshal(m); err == nil {
						cfg.Metrics = raw
					}
				}
			}
		}
	}
	a.apiPort = apiPortFromConfig(&cfg)

	a.stopXray()

	// Use a dedicated config file so the agent's xray doesn't fight the
	// master's xray over the shared config.json when both run on one host.
	agentConfigPath := config.GetBinFolderPath() + "/config-agent.json"
	proc := xray.NewTestProcess(&cfg, agentConfigPath)
	if err := proc.Start(); err != nil {
		return fmt.Errorf("failed to start xray: %w", err)
	}
	a.xrayProc = proc
	return nil
}

func (a *Agent) stopXray() {
	if a.xrayProc != nil {
		_ = a.xrayProc.Stop()
		a.xrayProc = nil
	}
}

// apiPortFromConfig extracts the stats API port (the inbound tagged "api").
func apiPortFromConfig(cfg *xray.Config) int {
	for _, inbound := range cfg.InboundConfigs {
		if inbound.Tag == "api" {
			return inbound.Port
		}
	}
	return 0
}

type trafficReport struct {
	Inbounds []trafficEntry `json:"inbounds"`
	Users    []userEntry    `json:"users"`
	Onlines  []string       `json:"onlines"`
}

type trafficEntry struct {
	Tag  string `json:"tag"`
	Up   int64  `json:"up"`
	Down int64  `json:"down"`
}

type userEntry struct {
	Email string `json:"email"`
	Up    int64  `json:"up"`
	Down  int64  `json:"down"`
}

// 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)
}

// 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
	}
	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 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 snap.ClientTraffics {
		if ct == nil || ct.Email == "" || ct.Up+ct.Down <= 0 {
			continue
		}
		report.Users = append(report.Users, userEntry{Email: ct.Email, Up: ct.Up, Down: ct.Down})
	}
	for _, u := range onlineUsers {
		report.Onlines = append(report.Onlines, u.Email)
	}
	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))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+a.config.Token)

	resp, err := a.httpClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("traffic report failed with status %d", resp.StatusCode)
	}
	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)
	if err != nil {
		return nil, err
	}
	var cfg Config
	if err := json.Unmarshal(data, &cfg); err != nil {
		return nil, err
	}
	return &cfg, nil
}