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

package controller

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

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

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

func (a *AgentController) heartbeat(c *gin.Context) {
	agent, ok := a.findAgent(bearerToken(c))
	if !ok {
		c.JSON(401, gin.H{"message": "unauthorized"})
		return
	}
	a.touchAgent(uint(agent.Id), "")
	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)
	c.JSON(200, gin.H{
		"config_hash": hex.EncodeToString(sum[:]),
		"xray_config": json.RawMessage(raw),
	})
}

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 {
		if len(req.Onlines) > 0 {
			a.inboundService.SetNodeOnlineTree(*agent.NodeID, map[string][]string{
				"agent": req.Onlines,
			})
		} else {
			a.inboundService.ClearNodeOnlineClients(*agent.NodeID)
		}
	}
	a.touchAgent(uint(agent.Id), "")
	jsonObj(c, gin.H{"message": "ok"}, nil)
}

// issueAgentToken generates a new agent token bound to nodeID and stores only
// its hash. The raw token is returned once and must be shown to the operator.
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],
		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 once; only its hash is stored.
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)
}

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
}