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"`
}

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

// 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"`
}

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

	traffics, clientTraffics, err := a.xrayAPI.GetTraffic()
	if err != nil {
		return err
	}
	onlineUsers, _ := a.xrayAPI.GetOnlineUsers()

	report := trafficReport{}
	for _, t := range 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 {
		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)
	}
	if len(report.Inbounds) == 0 && len(report.Users) == 0 && len(report.Onlines) == 0 {
		return nil
	}

	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
}

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