3x-ui/internal/web/service/node_tree.go

package service

import (
	"context"
	"sync"

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

// LocalDescendants returns this panel's read-only summaries of the nodes it
// directly manages, so a parent panel can surface them as transitive sub-nodes
// (#4983). Only nodes with a known GUID are included — a stable identity is
// required to attribute them one hop up. Not recursive: each panel reports its
// own direct nodes, and a master walks one level via each direct node's
// endpoint, which covers the Node1 -> Node2 -> Node3 case.
func (s *NodeService) LocalDescendants() ([]model.NodeSummary, error) {
	selfGuid, _ := (&SettingService{}).GetPanelGuid()
	db := database.GetDB()
	var nodes []*model.Node
	if err := db.Model(model.Node{}).Order("id asc").Find(&nodes).Error; err != nil {
		return nil, err
	}
	out := make([]model.NodeSummary, 0, len(nodes))
	for _, n := range nodes {
		if n.Guid == "" {
			continue
		}
		out = append(out, model.NodeSummary{
			Guid:          n.Guid,
			ParentGuid:    selfGuid,
			Name:          n.Name,
			Address:       n.Address,
			Scheme:        n.Scheme,
			Port:          n.Port,
			Status:        n.Status,
			LastHeartbeat: n.LastHeartbeat,
			LatencyMs:     n.LatencyMs,
			PanelVersion:  n.PanelVersion,
			XrayVersion:   n.XrayVersion,
			XrayState:     n.XrayState,
			XrayError:     n.XrayError,
		})
	}
	return out, nil
}

var (
	nodeDescendantsMu    sync.RWMutex
	nodeDescendantsCache = map[int][]model.NodeSummary{}
)

// RefreshDescendants pulls a direct node's published sub-node summaries and
// caches them keyed by node id. Best-effort: a fetch error keeps the last good
// set (the node may be briefly unreachable). Called from the heartbeat job.
func (s *NodeService) RefreshDescendants(ctx context.Context, n *model.Node) {
	if n == nil {
		return
	}
	mgr := runtime.GetManager()
	if mgr == nil {
		return
	}
	rt, err := mgr.RemoteFor(n)
	if err != nil {
		return
	}
	summaries, err := rt.GetDescendants(ctx)
	if err != nil {
		return
	}
	nodeDescendantsMu.Lock()
	if len(summaries) == 0 {
		delete(nodeDescendantsCache, n.Id)
	} else {
		nodeDescendantsCache[n.Id] = summaries
	}
	nodeDescendantsMu.Unlock()
}

// ClearDescendants drops a node's cached sub-node summaries (its probe failed).
func (s *NodeService) ClearDescendants(nodeID int) {
	nodeDescendantsMu.Lock()
	delete(nodeDescendantsCache, nodeID)
	nodeDescendantsMu.Unlock()
}

// RetainEnabledNodeDescendants drops sub-nodes learned from nodes the heartbeat no
// longer probes: disabled ones it skips, deleted ones missing from nodes.
func (s *NodeService) RetainEnabledNodeDescendants(nodes []*model.Node) {
	enabled := make(map[int]bool, len(nodes))
	for _, n := range nodes {
		enabled[n.Id] = n.Enable
	}
	nodeDescendantsMu.Lock()
	for nodeID := range nodeDescendantsCache {
		if !enabled[nodeID] {
			delete(nodeDescendantsCache, nodeID)
		}
	}
	nodeDescendantsMu.Unlock()
}

func cachedDescendants() []model.NodeSummary {
	nodeDescendantsMu.RLock()
	defer nodeDescendantsMu.RUnlock()
	out := make([]model.NodeSummary, 0)
	for _, list := range nodeDescendantsCache {
		out = append(out, list...)
	}
	return out
}

// GetNodeTree returns the direct nodes plus any transitive sub-nodes learned
// from them, with per-GUID counts so each node shows only the inbounds/online
// it physically hosts (#4983). Direct nodes carry the master's own GUID as
// ParentGuid; a transitive node carries its parent node's GUID. Transitive
// nodes are read-only projections (Id == 0). Used by the Nodes page and the
// heartbeat broadcast — never for probing/syncing, which stay on GetAll.
func (s *NodeService) GetNodeTree() ([]*model.Node, error) {
	nodes, err := s.GetAll()
	if err != nil {
		return nodes, err
	}
	selfGuid, _ := (&SettingService{}).GetPanelGuid()
	directGuids := make(map[string]struct{}, len(nodes))
	for _, n := range nodes {
		n.ParentGuid = selfGuid
		if n.Guid != "" {
			directGuids[n.Guid] = struct{}{}
		}
	}

	seen := make(map[string]struct{})
	var transitive []*model.Node
	for _, sum := range cachedDescendants() {
		if sum.Guid == "" {
			continue
		}
		if _, ok := directGuids[sum.Guid]; ok {
			continue // already shown as a direct node
		}
		if _, ok := seen[sum.Guid]; ok {
			continue
		}
		seen[sum.Guid] = struct{}{}
		transitive = append(transitive, &model.Node{
			Guid:          sum.Guid,
			ParentGuid:    sum.ParentGuid,
			Name:          sum.Name,
			Address:       sum.Address,
			Scheme:        sum.Scheme,
			Port:          sum.Port,
			Status:        sum.Status,
			LastHeartbeat: sum.LastHeartbeat,
			LatencyMs:     sum.LatencyMs,
			PanelVersion:  sum.PanelVersion,
			XrayVersion:   sum.XrayVersion,
			XrayState:     sum.XrayState,
			XrayError:     sum.XrayError,
			Transitive:    true,
		})
	}
	if len(transitive) == 0 {
		return nodes, nil
	}

	all := make([]*model.Node, 0, len(nodes)+len(transitive))
	all = append(all, nodes...)
	all = append(all, transitive...)
	s.recountByGuid(all, selfGuid)
	return all, nil
}

func (s *NodeService) GetNodeTreeView() ([]*NodeView, error) {
	nodes, err := s.GetNodeTree()
	if err != nil {
		return nil, err
	}
	return toNodeViews(nodes), nil
}

// recountByGuid recomputes InboundCount/OnlineCount/DepletedCount for every node
// in the tree, keyed by the GUID that physically hosts each inbound, so a direct
// node shows only its own inbounds and each transitive node shows its own
// (#4983). In a flat topology the per-GUID and per-node-id counts coincide, so
// this only changes behaviour once a transitive node exists.
func (s *NodeService) recountByGuid(nodes []*model.Node, selfGuid string) {
	db := database.GetDB()
	type ibRow struct {
		Id             int
		NodeID         *int   `gorm:"column:node_id"`
		OriginNodeGuid string `gorm:"column:origin_node_guid"`
	}
	var ibRows []ibRow
	if err := db.Table("inbounds").Select("id, node_id, origin_node_guid").Scan(&ibRows).Error; err != nil {
		return
	}
	ambiguous := ambiguousNodeGuids(nodes, selfGuid)
	effByInbound := make(map[int]string, len(ibRows))
	inboundCountByGuid := make(map[string]int)
	for _, r := range ibRows {
		guid := r.OriginNodeGuid
		if guid == "" {
			if r.NodeID != nil {
				guid = synthNodeGuid(*r.NodeID)
			} else {
				guid = selfGuid
			}
		} else if r.NodeID != nil {
			// Origin still holds an ambiguous GUID (cloned server / master-shared,
			// not yet re-attributed): bucket under the hosting node's unique id so
			// the clones don't merge.
			if _, bad := ambiguous[guid]; bad {
				guid = synthNodeGuid(*r.NodeID)
			}
		}
		effByInbound[r.Id] = guid
		inboundCountByGuid[guid]++
	}

	// Classify by EMAIL (not the stale client_traffics.inbound_id) and bucket
	// each client under its inbound's effective attribution GUID, deduping a
	// client attached to several inbounds under the same GUID.
	depletedByGuid := make(map[string]int)
	disabledByGuid := make(map[string]int)
	activeByGuid := make(map[string]int)
	if statuses, err := s.nodeClientStatuses(); err == nil {
		seen := make(map[string]map[int]struct{})
		for _, st := range statuses {
			guid, ok := effByInbound[st.InboundID]
			if !ok {
				continue
			}
			clientsSeen := seen[guid]
			if clientsSeen == nil {
				clientsSeen = make(map[int]struct{})
				seen[guid] = clientsSeen
			}
			if _, dup := clientsSeen[st.ClientID]; dup {
				continue
			}
			clientsSeen[st.ClientID] = struct{}{}
			switch {
			case st.Depleted:
				depletedByGuid[guid]++
			case st.Disabled:
				disabledByGuid[guid]++
			default:
				activeByGuid[guid]++
			}
		}
	}

	onlineByGuid := s.onlineEmailsByGuid()
	for _, n := range nodes {
		guid := effectiveNodeGuid(n, ambiguous)
		n.InboundCount = inboundCountByGuid[guid]
		n.OnlineCount = len(onlineByGuid[guid])
		n.DepletedCount = depletedByGuid[guid]
		n.DisabledCount = disabledByGuid[guid]
		n.ActiveCount = activeByGuid[guid]
	}
}