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

package controller

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"strconv"
	"time"

	"github.com/mhsanaei/3x-ui/v3/internal/config"
	"github.com/mhsanaei/3x-ui/v3/internal/database"
	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
	"github.com/mhsanaei/3x-ui/v3/internal/logger"
	"github.com/mhsanaei/3x-ui/v3/internal/selfupdate"
	"github.com/mhsanaei/3x-ui/v3/internal/web/service"
	"github.com/mhsanaei/3x-ui/v3/internal/xray"

	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

// AgentController serves the lightweight agent API.
//
// Agent endpoints (register/heartbeat/config/traffic) authenticate with the
// agent's own Bearer token, which is NOT a panel API token — they are mounted
// on a public route group and must never go through the panel auth middleware.
// The panel-authed administration endpoints (token issuance, egress
// management) live under the authenticated /panel/api group instead.
type AgentController struct {
	xrayService    service.XrayService
	inboundService service.InboundService
}

func NewAgentController(g *gin.RouterGroup) *AgentController {
	a := &AgentController{}
	a.initRouter(g)
	return a
}

func (a *AgentController) initRouter(g *gin.RouterGroup) {
	g.POST("/register", a.register)
	g.POST("/heartbeat", a.heartbeat)
	g.GET("/config", a.getConfig)
	g.POST("/traffic", a.reportTraffic)
}

// RegisterAdminRoutes mounts the panel-authenticated administration endpoints
// on an already-authed route group.
func (a *AgentController) RegisterAdminRoutes(g *gin.RouterGroup) {
	g.POST("/tokens", a.createToken)
	g.GET("/tokens", a.listTokens)
	g.DELETE("/tokens/:id", a.deleteToken)
	g.GET("/egress/:token", a.getEgress)
	g.PUT("/egress/:token", a.setEgress)
	g.GET("/install/:nodeId", a.getAgentInstall)
}

func hashAgentToken(token string) string {
	sum := sha256.Sum256([]byte(token))
	return hex.EncodeToString(sum[:])
}

var errRequest = errors.New("bad request")

func bearerToken(c *gin.Context) string {
	auth := c.GetHeader("Authorization")
	if len(auth) > 7 && auth[:7] == "Bearer " {
		return auth[7:]
	}
	return ""
}

func (a *AgentController) findAgent(token string) (*model.AgentNode, bool) {
	if token == "" {
		return nil, false
	}
	var agent model.AgentNode
	if err := database.GetDB().Where("token_hash = ?", hashAgentToken(token)).First(&agent).Error; err != nil {
		return nil, false
	}
	// Item 19 backfill: the token has just passed the hash check, so for
	// rows created before plaintext storage existed, adopt the presented
	// token as the recoverable copy. The guarded update only fires while
	// token_plain is empty — an authenticated request can never overwrite
	// an already-stored token (only issuance writes a fresh one).
	if agent.TokenPlain == "" {
		if res := database.GetDB().Model(&model.AgentNode{}).
			Where("id = ? AND token_plain = ''", agent.Id).
			Update("token_plain", token); res.Error == nil && res.RowsAffected > 0 {
			agent.TokenPlain = token
		}
	}
	return &agent, true
}

func (a *AgentController) touchAgent(id uint, name string) {
	updates := map[string]any{"last_seen": time.Now().UnixMilli()}
	if name != "" {
		updates["name"] = name
	}
	db := database.GetDB()
	_ = db.Model(&model.AgentNode{}).Where("id = ?", id).Updates(updates).Error
	// If bound to a node, refresh the node's heartbeat so the legacy node
	// list shows it online (the legacy probe can't reach agent-only nodes).
	var nodeID *int
	_ = db.Model(&model.AgentNode{}).Select("node_id").Where("id = ?", id).Scan(&nodeID).Error
	if nodeID != nil {
		now := time.Now().UnixMilli()
		_ = db.Exec("UPDATE nodes SET last_heartbeat = ?, status = 'online' WHERE id = ?", now, *nodeID).Error
	}
}

// register activates a pre-issued token. Tokens are issued bound to a node via
// POST /panel/api/agent/tokens; the agent only claims its token here.
func (a *AgentController) register(c *gin.Context) {
	var req struct {
		Token string `json:"token"`
		Name  string `json:"name"`
	}
	if err := c.ShouldBindJSON(&req); err != nil || req.Token == "" {
		jsonMsg(c, "Token required", errRequest)
		return
	}
	agent, ok := a.findAgent(req.Token)
	if !ok {
		c.JSON(401, gin.H{"message": "unknown token"})
		return
	}
	a.touchAgent(uint(agent.Id), req.Name)
	logger.Infof("Agent registered: %s (token %s...)", req.Name, agent.TokenPrefix)
	jsonObj(c, gin.H{"message": "registered"}, nil)
}

// heartbeatTelemetry is the optional telemetry object an agent piggybacks
// on its heartbeat (P13). Every field is a pointer: a missing field means
// "not reported this round" and must not overwrite the stored value. Older
// agents send no telemetry object at all and are handled exactly as before.
type heartbeatTelemetry struct {
	AgentVersion  *string  `json:"agentVersion"`
	XrayVersion   *string  `json:"xrayVersion"`
	UptimeSeconds *uint64  `json:"uptimeSeconds"`
	LatencyMs     *int     `json:"latencyMs"`
	CpuPercent    *float64 `json:"cpuPercent"`
	MemPercent    *float64 `json:"memPercent"`
	IPv4          *string  `json:"ipv4"`
	IPv6          *string  `json:"ipv6"`
}

// updateAgentTelemetry persists the reported telemetry fields onto the
// agent row. Only fields actually present in the report are written.
func (a *AgentController) updateAgentTelemetry(id uint, t *heartbeatTelemetry) {
	updates := map[string]any{"telemetry_at": time.Now().UnixMilli()}
	if t.AgentVersion != nil {
		updates["agent_version"] = *t.AgentVersion
	}
	if t.XrayVersion != nil {
		updates["xray_version"] = *t.XrayVersion
	}
	if t.UptimeSeconds != nil {
		updates["uptime_secs"] = *t.UptimeSeconds
	}
	if t.LatencyMs != nil {
		updates["latency_ms"] = *t.LatencyMs
	}
	if t.CpuPercent != nil {
		updates["cpu_pct"] = *t.CpuPercent
	}
	if t.MemPercent != nil {
		updates["mem_pct"] = *t.MemPercent
	}
	if t.IPv4 != nil {
		updates["ipv4"] = *t.IPv4
	}
	if t.IPv6 != nil {
		updates["ipv6"] = *t.IPv6
	}
	_ = database.GetDB().Model(&model.AgentNode{}).Where("id = ?", id).Updates(updates).Error
}

// heartbeatUpdateStatus is the agent's report about a one-click update
// intent (item 20), piggybacked on the heartbeat like telemetry.
type heartbeatUpdateStatus struct {
	RequestedAt int64  `json:"requestedAt"`
	State       string `json:"state"`
	Error       string `json:"error"`
}

// applyAgentUpdateStatus folds an agent's update report into the stored
// intent. Reports for a superseded intent (requested_at mismatch while an
// intent is pending) are ignored. Terminal reports (failed / current)
// clear the intent so the config poll stops re-offering it; the recorded
// state stays visible on the node row.
func (a *AgentController) applyAgentUpdateStatus(id uint, st *heartbeatUpdateStatus) {
	db := database.GetDB()
	var row model.AgentNode
	if err := db.First(&row, id).Error; err != nil {
		return
	}
	if row.UpdateRequestedAt == 0 || st.RequestedAt != row.UpdateRequestedAt {
		return
	}
	updates := map[string]any{"update_state_at": time.Now().UnixMilli()}
	switch st.State {
	case "downloading", "swapped":
		updates["update_state"] = st.State
		updates["update_error"] = ""
	case "failed":
		updates["update_state"] = "failed"
		updates["update_error"] = st.Error
		updates["update_requested_at"] = int64(0)
	case "current":
		updates["update_state"] = "current"
		updates["update_error"] = ""
		updates["update_requested_at"] = int64(0)
	default:
		return
	}
	_ = db.Model(&model.AgentNode{}).Where("id = ?", id).Updates(updates).Error
}

// confirmAgentUpdate marks a pending intent successful once the agent
// reports an A-series version at or beyond the intent's target — the new
// version appearing is the end-to-end success signal of the update.
func (a *AgentController) confirmAgentUpdate(id uint, reportedVersion string) {
	if reportedVersion == "" {
		return
	}
	db := database.GetDB()
	var row model.AgentNode
	if err := db.First(&row, id).Error; err != nil {
		return
	}
	if row.UpdateRequestedAt == 0 || row.UpdateTarget == "" {
		return
	}
	if cmp, ok := config.CompareDisplayVersion(reportedVersion, row.UpdateTarget); !ok || cmp < 0 {
		return
	}
	_ = db.Model(&model.AgentNode{}).Where("id = ?", id).Updates(map[string]any{
		"update_state":        "success",
		"update_error":        "",
		"update_requested_at": int64(0),
		"update_state_at":     time.Now().UnixMilli(),
	}).Error
}

func (a *AgentController) heartbeat(c *gin.Context) {
	agent, ok := a.findAgent(bearerToken(c))
	if !ok {
		c.JSON(401, gin.H{"message": "unauthorized"})
		return
	}
	// Telemetry is optional and best-effort: a body that fails to parse
	// (or no body at all) must never fail the heartbeat itself.
	var req struct {
		Telemetry    *heartbeatTelemetry    `json:"telemetry"`
		UpdateStatus *heartbeatUpdateStatus `json:"updateStatus"`
	}
	_ = c.ShouldBindJSON(&req)
	a.touchAgent(uint(agent.Id), "")
	reportedVersion := agent.AgentVersion
	if req.Telemetry != nil {
		a.updateAgentTelemetry(uint(agent.Id), req.Telemetry)
		if req.Telemetry.AgentVersion != nil && *req.Telemetry.AgentVersion != "" {
			reportedVersion = *req.Telemetry.AgentVersion
		}
	}
	if req.UpdateStatus != nil {
		a.applyAgentUpdateStatus(uint(agent.Id), req.UpdateStatus)
	}
	a.confirmAgentUpdate(uint(agent.Id), reportedVersion)
	jsonObj(c, gin.H{"message": "ok"}, nil)
}

// getConfig returns the full xray config for the agent's bound node.
func (a *AgentController) getConfig(c *gin.Context) {
	agent, ok := a.findAgent(bearerToken(c))
	if !ok {
		c.JSON(401, gin.H{"message": "unauthorized"})
		return
	}
	if agent.NodeID == nil {
		c.JSON(409, gin.H{"message": "agent token is not bound to a node"})
		return
	}
	cfg, err := a.xrayService.GetXrayConfigForAgent(*agent.NodeID, agent.Outbounds, agent.RoutingRules)
	if err != nil {
		jsonMsg(c, "Failed to build xray config", err)
		return
	}
	raw, err := json.Marshal(cfg)
	if err != nil {
		jsonMsg(c, "Failed to encode xray config", err)
		return
	}
	sum := sha256.Sum256(raw)
	payload := gin.H{
		"config_hash": hex.EncodeToString(sum[:]),
		"xray_config": json.RawMessage(raw),
	}
	// Item 20: a pending one-click update intent rides along on the config
	// poll (no new channel — the agent already polls this endpoint).
	if agent.UpdateRequestedAt > 0 &&
		(agent.UpdateState == "pending" || agent.UpdateState == "downloading") {
		payload["agent_update"] = gin.H{
			"requested_at": agent.UpdateRequestedAt,
			"target":       agent.UpdateTarget,
			"tag":          agent.UpdateTag,
			"base_url":     selfupdate.BaseURL(),
		}
	}
	c.JSON(200, payload)
}

type agentTrafficReport struct {
	Inbounds []struct {
		Tag  string `json:"tag"`
		Up   int64  `json:"up"`
		Down int64  `json:"down"`
	} `json:"inbounds"`
	Users []struct {
		Email string `json:"email"`
		Up    int64  `json:"up"`
		Down  int64  `json:"down"`
	} `json:"users"`
	// Onlines lists client emails with at least one live connection.
	Onlines []string `json:"onlines"`
}

// reportTraffic ingests per-poll traffic deltas from the agent. Inbound deltas
// land on the agent's node inbound rows; client deltas reuse the shared
// email-keyed write path (plus its lifecycle maintenance, which also drives
// config changes the agent picks up on its next poll).
func (a *AgentController) reportTraffic(c *gin.Context) {
	agent, ok := a.findAgent(bearerToken(c))
	if !ok {
		c.JSON(401, gin.H{"message": "unauthorized"})
		return
	}
	var req agentTrafficReport
	if err := c.ShouldBindJSON(&req); err != nil {
		jsonMsg(c, "Invalid request", err)
		return
	}
	db := database.GetDB()
	if agent.NodeID != nil {
		for _, ib := range req.Inbounds {
			if ib.Tag == "" || ib.Up+ib.Down <= 0 {
				continue
			}
			_ = db.Model(&model.Inbound{}).
				Where("tag = ? AND node_id = ?", ib.Tag, *agent.NodeID).
				Updates(map[string]any{
					"up":   gorm.Expr(database.ClampedAddExpr("up"), ib.Up),
					"down": gorm.Expr(database.ClampedAddExpr("down"), ib.Down),
				}).Error
		}
	}
	var clientTraffics []*xray.ClientTraffic
	for _, u := range req.Users {
		if u.Email == "" || u.Up+u.Down <= 0 {
			continue
		}
		clientTraffics = append(clientTraffics, &xray.ClientTraffic{
			Email: u.Email,
			Up:    u.Up,
			Down:  u.Down,
		})
	}
	if len(clientTraffics) > 0 {
		if _, _, err := a.inboundService.AddTraffic(nil, clientTraffics); err != nil {
			jsonMsg(c, "Failed to store traffic", err)
			return
		}
	}
	if agent.NodeID != nil {
		a.inboundService.SetAgentOnlineClients(*agent.NodeID, req.Onlines)
	}
	a.touchAgent(uint(agent.Id), "")
	jsonObj(c, gin.H{"message": "ok"}, nil)
}

// issueAgentToken generates a new agent token bound to nodeID. The hash is
// what authenticates; the raw token is also kept recoverably (TokenPlain,
// item 19) so the node edit page can re-display the install command.
func issueAgentToken(nodeID *int, name string) (string, int, error) {
	raw := make([]byte, 32)
	if _, err := rand.Read(raw); err != nil {
		return "", 0, err
	}
	token := hex.EncodeToString(raw)
	agent := model.AgentNode{
		TokenHash:   hashAgentToken(token),
		TokenPrefix: token[:12],
		TokenPlain:  token,
		Name:        name,
		NodeID:      nodeID,
	}
	if err := database.GetDB().Create(&agent).Error; err != nil {
		return "", 0, err
	}
	return token, agent.Id, nil
}

// createToken issues a new agent token bound to a node. The raw token is
// returned in this response and kept recoverably for the install command.
func (a *AgentController) createToken(c *gin.Context) {
	var req struct {
		NodeID *int   `json:"node_id"`
		Name   string `json:"name"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		jsonMsg(c, "Invalid request", err)
		return
	}
	if req.NodeID != nil {
		var node model.Node
		if err := database.GetDB().First(&node, *req.NodeID).Error; err != nil {
			jsonMsg(c, "Node not found", errRequest)
			return
		}
	}
	token, _, err := issueAgentToken(req.NodeID, req.Name)
	if err != nil {
		jsonMsg(c, "Failed to generate token", err)
		return
	}
	var agent model.AgentNode
	_ = database.GetDB().Where("token_hash = ?", hashAgentToken(token)).First(&agent).Error
	jsonObj(c, gin.H{
		"id":           agent.Id,
		"token":        token,
		"token_prefix": agent.TokenPrefix,
		"node_id":      agent.NodeID,
		"name":         agent.Name,
	}, nil)
}

// getAgentInstall returns what the node edit page needs to re-display the
// full install command for an agent node (item 19): the agent's name and,
// when the master holds a recoverable copy, its raw token. Legacy rows
// whose token has not been backfilled yet answer hasToken=false with an
// empty token — the UI then shows the command without the token and points
// at /etc/x-ui/agent.json on the node itself.
func (a *AgentController) getAgentInstall(c *gin.Context) {
	nodeID, err := strconv.Atoi(c.Param("nodeId"))
	if err != nil {
		jsonMsg(c, "Invalid node id", errRequest)
		return
	}
	var agent model.AgentNode
	if err := database.GetDB().Where("node_id = ?", nodeID).
		Order("last_seen DESC, id DESC").First(&agent).Error; err != nil {
		jsonMsg(c, "Agent not found for node", errRequest)
		return
	}
	jsonObj(c, gin.H{
		"nodeId":   nodeID,
		"name":     agent.Name,
		"token":    agent.TokenPlain,
		"hasToken": agent.TokenPlain != "",
	}, nil)
}

func (a *AgentController) listTokens(c *gin.Context) {
	var agents []model.AgentNode
	_ = database.GetDB().Order("id DESC").Find(&agents).Error
	out := make([]gin.H, 0, len(agents))
	for _, ag := range agents {
		out = append(out, gin.H{
			"id":           ag.Id,
			"token_prefix": ag.TokenPrefix,
			"name":         ag.Name,
			"node_id":      ag.NodeID,
			"last_seen":    ag.LastSeen,
			"has_egress":   ag.Outbounds != "" || ag.RoutingRules != "",
			"created_at":   ag.CreatedAt,
		})
	}
	jsonObj(c, out, nil)
}

func (a *AgentController) deleteToken(c *gin.Context) {
	_ = database.GetDB().Delete(&model.AgentNode{}, "id = ?", c.Param("id")).Error
	jsonObj(c, gin.H{"message": "deleted"}, nil)
}

// getEgress returns the panel-managed egress (outbounds + routing rules) for
// the agent identified by its token.
func (a *AgentController) getEgress(c *gin.Context) {
	agent, ok := a.findAgent(c.Param("token"))
	if !ok {
		c.JSON(404, gin.H{"message": "agent not found"})
		return
	}
	jsonObj(c, gin.H{
		"id":            agent.Id,
		"token_prefix":  agent.TokenPrefix,
		"name":          agent.Name,
		"node_id":       agent.NodeID,
		"outbounds":     json.RawMessage(orEmptyArray(agent.Outbounds)),
		"routing_rules": json.RawMessage(orEmptyArray(agent.RoutingRules)),
	}, nil)
}

// setEgress replaces the panel-managed egress for the agent. The panel is the
// only source of truth: agents never report local egress, they only apply
// what the generated config carries.
func (a *AgentController) setEgress(c *gin.Context) {
	agent, ok := a.findAgent(c.Param("token"))
	if !ok {
		c.JSON(404, gin.H{"message": "agent not found"})
		return
	}
	var req struct {
		Outbounds    json.RawMessage `json:"outbounds"`
		RoutingRules json.RawMessage `json:"routing_rules"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		jsonMsg(c, "Invalid request", err)
		return
	}
	outbounds, err := canonicalJSONArray(req.Outbounds)
	if err != nil {
		jsonMsg(c, "outbounds must be a JSON array", err)
		return
	}
	rules, err := canonicalJSONArray(req.RoutingRules)
	if err != nil {
		jsonMsg(c, "routing_rules must be a JSON array", err)
		return
	}
	if err := database.GetDB().Model(&model.AgentNode{}).Where("id = ?", agent.Id).
		Updates(map[string]any{"outbounds": outbounds, "routing_rules": rules}).Error; err != nil {
		jsonMsg(c, "Failed to store egress", err)
		return
	}
	jsonObj(c, gin.H{"message": "saved"}, nil)
}

func orEmptyArray(s string) string {
	if s == "" {
		return "[]"
	}
	return s
}

// canonicalJSONArray validates that raw is a JSON array and returns it in
// canonical form. Empty input yields an empty string (template defaults).
func canonicalJSONArray(raw json.RawMessage) (string, error) {
	if len(raw) == 0 {
		return "", nil
	}
	var arr []any
	if err := json.Unmarshal(raw, &arr); err != nil {
		return "", err
	}
	out, err := json.Marshal(arr)
	if err != nil {
		return "", err
	}
	return string(out), nil
}