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

package controller

import (
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/mhsanaei/3x-ui/v3/internal/database"
	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
	"github.com/mhsanaei/3x-ui/v3/internal/web/service/probe"

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

// Item 7 B4: probe task management — the operator-maintained latency
// target set (ITEM7-SCOPE §5.7, user ruling Q4: no preset list; the
// operator keeps targets, grouped by the free-form region label, from
// the panel). Endpoints ride the panel-authenticated agent admin
// group like B3's probe-meta, and validation mirrors hub api.rs's
// save_ping_task chain refusal for refusal — including refusing (not
// clamping) an out-of-range interval.
//
// The panel-facing id space is the panel NODE id, as on the
// probe-meta surface; assignments are stored against agent ids and
// translated at this layer, where the binding lives.

type ProbeTasksController struct{}

func NewProbeTasksController(g *gin.RouterGroup) *ProbeTasksController {
	pc := &ProbeTasksController{}
	g.GET("/probe-tasks", pc.listTasks)
	g.PUT("/probe-tasks", pc.saveTask)
	g.POST("/probe-tasks/batch", pc.saveTaskBatch)
	g.DELETE("/probe-tasks/:id", pc.deleteTask)
	g.POST("/probe-tasks/reorder", pc.reorderTasks)
	g.PUT("/probe-tasks/by-node/:nodeId", pc.setNodeTasks)
	return pc
}

// probeTaskForm is the wire shape of one task as the panel edits it.
type probeTaskForm struct {
	ID          int64    `json:"id"`
	Name        string   `json:"name"`
	Target      string   `json:"target"`
	Interval    int64    `json:"interval"`
	AutoJoin    bool     `json:"autoJoin"`
	Sort        int64    `json:"sort"`
	Region      string   `json:"region"`
	NodeIDs     []int64  `json:"nodeIds"`
	BaseNodeIDs *[]int64 `json:"baseNodeIds"`
}

// agentsByNode maps panel node id → agent row for every bound agent.
func agentsByNode() (map[int64]model.AgentNode, error) {
	var agents []model.AgentNode
	if err := database.GetDB().Where("node_id IS NOT NULL").Find(&agents).Error; err != nil {
		return nil, err
	}
	out := make(map[int64]model.AgentNode, len(agents))
	for _, ag := range agents {
		out[int64(*ag.NodeID)] = ag
	}
	return out, nil
}

// agentDisplayName renders an agent for operator-facing messages: the
// bound node's name where there is one (the panel identifies nodes by
// name and never shows an id, as the hub does).
func agentDisplayName(agentID int64) string {
	var ag model.AgentNode
	if err := database.GetDB().Where("id = ?", agentID).First(&ag).Error; err == nil {
		if ag.NodeID != nil {
			var node model.Node
			if err := database.GetDB().Where("id = ?", *ag.NodeID).First(&node).Error; err == nil && node.Name != "" {
				return node.Name
			}
		}
		if ag.Name != "" {
			return ag.Name
		}
	}
	return fmt.Sprintf("#%d", agentID)
}

// nodeName renders a panel node for messages, "" when it is gone.
func nodeName(nodeID int64) string {
	var node model.Node
	if err := database.GetDB().Where("id = ?", nodeID).First(&node).Error; err != nil {
		return ""
	}
	return node.Name
}

// resolveAgentIDs translates the form's node ids into agent ids,
// refusing the whole save when one names a node that is gone or has
// no agent bound (the hub's foreign key does this by construction;
// the naming keeps the refusal readable).
func resolveAgentIDs(nodeIDs []int64, bound map[int64]model.AgentNode) ([]int64, string) {
	out := make([]int64, 0, len(nodeIDs))
	for _, nodeID := range nodeIDs {
		ag, ok := bound[nodeID]
		if !ok {
			if name := nodeName(nodeID); name != "" {
				return nil, fmt.Sprintf("节点「%s」未绑定 agent，无法指派探测任务", name)
			}
			return nil, fmt.Sprintf("节点 %d 不存在，可能已被删除", nodeID)
		}
		out = append(out, int64(ag.Id))
	}
	return out, ""
}

func (pc *ProbeTasksController) listTasks(c *gin.Context) {
	tasks, err := probe.PingTasks()
	if err != nil {
		jsonMsg(c, "load probe tasks failed", err)
		return
	}
	bound, err := agentsByNode()
	if err != nil {
		jsonMsg(c, "load probe tasks failed", err)
		return
	}
	nodeByAgent := make(map[int64]int64, len(bound))
	for nodeID, ag := range bound {
		nodeByAgent[int64(ag.Id)] = nodeID
	}
	forms := make([]probeTaskForm, 0, len(tasks))
	for _, t := range tasks {
		f := probeTaskForm{
			ID: t.ID, Name: t.Name, Target: t.Target, Interval: t.Interval,
			AutoJoin: t.AutoJoin, Sort: t.Sort, Region: t.Region, NodeIDs: []int64{},
		}
		for _, agentID := range t.AgentIDs {
			if nodeID, ok := nodeByAgent[agentID]; ok {
				f.NodeIDs = append(f.NodeIDs, nodeID)
			}
		}
		forms = append(forms, f)
	}
	jsonObj(c, forms, nil)
}

// validProbeTarget mirrors hub valid_target: a target the agent can
// resolve — host:port, an IPv6 literal bracketed as a URL writes one.
// A bare contains(':') would admit a bare IPv6 address, ":443" with no
// host and "host:" with no port, each of which probes at 100% loss
// forever with nothing identifying the target as the cause.
func validProbeTarget(target string) bool {
	var host, port string
	if rest, ok := strings.CutPrefix(target, "["); ok {
		h, p, found := strings.Cut(rest, "]:")
		if !found {
			return false
		}
		host, port = h, p
	} else {
		// rsplit_once(':'): host:port splits on the last colon.
		i := strings.LastIndex(target, ":")
		if i < 0 {
			return false
		}
		host, port = target[:i], target[i+1:]
	}
	if host == "" || strings.Contains(host, ":") && !strings.HasPrefix(target, "[") {
		return false
	}
	n, err := strconv.Atoi(port)
	return err == nil && n > 0 && n <= 65535
}

// validateProbeTaskForm normalizes the form in place (hub order:
// trim, presence, name length, target shape, interval range) and
// returns the first refusal's message, "" when it may be stored. The
// region rule is ours (the hub has no region): the meta group rule.
func validateProbeTaskForm(f *probeTaskForm) string {
	// Trimmed into the stored value rather than a discarded copy:
	// what reaches the agent is the stored target, and a trailing
	// space from a paste would leave the probe reporting loss
	// indefinitely (hub discipline, same for the name).
	f.Name = strings.TrimSpace(f.Name)
	f.Target = strings.TrimSpace(f.Target)
	f.Region = strings.TrimSpace(f.Region)
	if f.Name == "" || f.Target == "" {
		return "请填写名称和目标"
	}
	if len([]rune(f.Name)) > 64 {
		return "名称最多 64 个字"
	}
	if !validProbeTarget(f.Target) {
		return "目标要写成「主机:端口」，例如 1.1.1.1:443 或 [2606:4700:4700::1111]:443"
	}
	// Refused rather than clamped: the agent clamps again on
	// arrival, so an out-of-range value would silently become a
	// different number while the panel still displays what was
	// entered (hub).
	if f.Interval < probe.MinProbeInterval || f.Interval > 3600 {
		return fmt.Sprintf("间隔要在 %d 到 3600 秒之间", probe.MinProbeInterval)
	}
	if len([]rune(f.Region)) > 13 || hasControl(f.Region) {
		return "地区最多 13 个字，不能含控制字符"
	}
	return ""
}

// probeTaskError renders a store refusal in the hub's words.
func probeTaskError(c *gin.Context, err error) {
	var crowded *probe.PingCrowdedError
	var autoJoin *probe.PingAutoJoinError
	switch {
	case errors.Is(err, probe.ErrPingTaskGone):
		jsonMsg(c, err.Error(), errRequest)
	case errors.As(err, &crowded):
		jsonMsg(c, fmt.Sprintf("节点「%s」会被分配超过 %d 个探测任务，agent 最多只跑这么多，多出来的会被静默丢掉",
			agentDisplayName(crowded.AgentID), probe.MaxProbesPerNode), errRequest)
	case errors.As(err, &autoJoin):
		jsonMsg(c, fmt.Sprintf("新节点会自动加入 %d 个探测任务，agent 最多只跑 %d 个",
			autoJoin.Count, probe.MaxProbesPerNode), errRequest)
	default:
		jsonMsg(c, err.Error(), errRequest)
	}
}

func (pc *ProbeTasksController) saveTask(c *gin.Context) {
	var form probeTaskForm
	if err := c.ShouldBindJSON(&form); err != nil {
		jsonMsg(c, "Invalid request", err)
		return
	}
	if form.NodeIDs == nil {
		form.NodeIDs = []int64{}
	}
	if msg := validateProbeTaskForm(&form); msg != "" {
		jsonMsg(c, msg, errRequest)
		return
	}
	bound, err := agentsByNode()
	if err != nil {
		jsonMsg(c, "save probe task failed", err)
		return
	}
	agentIDs, refusal := resolveAgentIDs(form.NodeIDs, bound)
	if refusal != "" {
		jsonMsg(c, refusal, errRequest)
		return
	}
	var base *[]int64
	if form.BaseNodeIDs != nil {
		baseIDs, refusal := resolveAgentIDs(*form.BaseNodeIDs, bound)
		if refusal != "" {
			jsonMsg(c, refusal, errRequest)
			return
		}
		base = &baseIDs
	}
	id, err := probe.SavePingTask(probe.PingTask{
		ID:       form.ID,
		Name:     form.Name,
		Target:   form.Target,
		Interval: form.Interval,
		AutoJoin: form.AutoJoin,
		Region:   form.Region,
		AgentIDs: agentIDs,
	}, base)
	if err != nil {
		probeTaskError(c, err)
		return
	}
	jsonObj(c, gin.H{"id": id}, nil)
}

// probeTaskBatchForm is the batch-create shape (V2.0.1 ⑤): many
// targets at once — already split client-side from the pasted
// lines/commas — sharing every other field. Name is a shared label:
// with several targets each task is auto-named after its target
// unless a name is given, in which case it prefixes the target
// ("<name> · <target>"); with exactly one target the name is used
// verbatim (and may be empty, falling back to the target).
type probeTaskBatchForm struct {
	Targets  []string `json:"targets"`
	Name     string   `json:"name"`
	Interval int64    `json:"interval"`
	AutoJoin bool     `json:"autoJoin"`
	Region   string   `json:"region"`
	NodeIDs  []int64  `json:"nodeIds"`
}

// saveTaskBatch creates one task per target, all-or-nothing on
// validation: every form is validated before the first insert, so a
// bad line refuses the whole paste with the hub-style message
// instead of leaving a half-entered list behind.
func (pc *ProbeTasksController) saveTaskBatch(c *gin.Context) {
	var form probeTaskBatchForm
	if err := c.ShouldBindJSON(&form); err != nil {
		jsonMsg(c, "Invalid request", err)
		return
	}
	if form.NodeIDs == nil {
		form.NodeIDs = []int64{}
	}
	// The batch form's interval is a shared convenience field: an
	// omitted one means the standard 60 s (the single-task form
	// keeps requiring an explicit value).
	if form.Interval == 0 {
		form.Interval = 60
	}
	// Split remnants and duplicates in the paste collapse here;
	// order is preserved (it becomes the tasks' sort order).
	seen := map[string]bool{}
	targets := make([]string, 0, len(form.Targets))
	for _, raw := range form.Targets {
		tgt := strings.TrimSpace(raw)
		if tgt == "" || seen[tgt] {
			continue
		}
		seen[tgt] = true
		targets = append(targets, tgt)
	}
	if len(targets) == 0 {
		jsonMsg(c, "请填写目标", errRequest)
		return
	}
	sharedName := strings.TrimSpace(form.Name)
	forms := make([]probeTaskForm, 0, len(targets))
	for _, tgt := range targets {
		name := sharedName
		switch {
		case len(targets) == 1:
			if name == "" {
				name = tgt
			}
		case name == "":
			name = tgt
		default:
			name = name + " · " + tgt
		}
		f := probeTaskForm{
			Name:     name,
			Target:   tgt,
			Interval: form.Interval,
			AutoJoin: form.AutoJoin,
			Region:   form.Region,
			NodeIDs:  form.NodeIDs,
		}
		if msg := validateProbeTaskForm(&f); msg != "" {
			jsonMsg(c, msg, errRequest)
			return
		}
		forms = append(forms, f)
	}
	bound, err := agentsByNode()
	if err != nil {
		jsonMsg(c, "save probe tasks failed", err)
		return
	}
	agentIDs, refusal := resolveAgentIDs(form.NodeIDs, bound)
	if refusal != "" {
		jsonMsg(c, refusal, errRequest)
		return
	}
	ids := make([]int64, 0, len(forms))
	for _, f := range forms {
		id, err := probe.SavePingTask(probe.PingTask{
			Name:     f.Name,
			Target:   f.Target,
			Interval: f.Interval,
			AutoJoin: f.AutoJoin,
			Region:   f.Region,
			AgentIDs: agentIDs,
		}, nil)
		if err != nil {
			probeTaskError(c, err)
			return
		}
		ids = append(ids, id)
	}
	jsonObj(c, gin.H{"ids": ids}, nil)
}

func (pc *ProbeTasksController) deleteTask(c *gin.Context) {
	id, err := strconv.ParseInt(c.Param("id"), 10, 64)
	if err != nil {
		jsonMsg(c, "Invalid request", err)
		return
	}
	if err := probe.DeletePingTask(id); err != nil {
		jsonMsg(c, "delete probe task failed", err)
		return
	}
	jsonObj(c, gin.H{"message": "ok"}, nil)
}

func (pc *ProbeTasksController) reorderTasks(c *gin.Context) {
	var form struct {
		IDs []int64 `json:"ids"`
	}
	if err := c.ShouldBindJSON(&form); err != nil {
		jsonMsg(c, "Invalid request", err)
		return
	}
	if err := probe.ReorderPingTasks(form.IDs); err != nil {
		jsonMsg(c, err.Error(), errRequest)
		return
	}
	jsonObj(c, gin.H{"message": "ok"}, nil)
}

// setNodeTasks edits one node's probe assignments from the node's
// side (hub set_node_ping_tasks): tasks/base are task ids, only the
// difference from base is applied.
func (pc *ProbeTasksController) setNodeTasks(c *gin.Context) {
	nodeID, err := strconv.Atoi(c.Param("nodeId"))
	if err != nil {
		c.JSON(404, gin.H{"message": "agent not found"})
		return
	}
	var agent model.AgentNode
	if err := database.GetDB().Where("node_id = ?", nodeID).First(&agent).Error; err != nil {
		c.JSON(404, gin.H{"message": "agent not found"})
		return
	}
	var form struct {
		Tasks []int64 `json:"tasks"`
		Base  []int64 `json:"base"`
	}
	if err := c.ShouldBindJSON(&form); err != nil {
		jsonMsg(c, "监控列表格式不对", errRequest)
		return
	}
	if form.Tasks == nil {
		form.Tasks = []int64{}
	}
	if form.Base == nil {
		form.Base = []int64{}
	}
	if err := probe.SetNodePingTasks(int64(agent.Id), form.Tasks, form.Base); err != nil {
		probeTaskError(c, err)
		return
	}
	jsonObj(c, gin.H{"message": "ok"}, nil)
}