Prism/internal/metrics/breakdown.go

package metrics

import (
	"github.com/prism/proxy/internal/mihomoapi"
	"github.com/prism/proxy/internal/store"
)

type TrafficItem struct {
	Key         string `json:"key"`
	Label       string `json:"label"`
	Node        string `json:"node,omitempty"`
	NodeKey     string `json:"node_key,omitempty"`
	Upload      int64  `json:"upload"`
	Download    int64  `json:"download"`
	Total       int64  `json:"total"`
	Connections int    `json:"connections"`
}

func AggregateConnections(conns []map[string]any, policyGroups map[string]bool) (domains, nodes []TrafficItem) {
	domainMap := map[string]*TrafficItem{}
	nodeMap := map[string]*TrafficItem{}

	for _, conn := range conns {
		host := connHost(conn)
		nodeKey := connLeafNode(conn, policyGroups)
		up := toInt64(conn["upload"])
		down := toInt64(conn["download"])

		if host != "" {
			key := host + "\x00" + nodeKey
			acc(domainMap, key, up, down)
			item := domainMap[key]
			item.Label = host
			item.NodeKey = nodeKey
		}
		if nodeKey != "" {
			acc(nodeMap, nodeKey, up, down)
		}
	}
	return sortTrafficItems(domainMap), sortTrafficItems(nodeMap)
}

func acc(m map[string]*TrafficItem, key string, up, down int64) {
	item, ok := m[key]
	if !ok {
		item = &TrafficItem{Key: key, Label: key}
		m[key] = item
	}
	item.Upload += up
	item.Download += down
	item.Total += up + down
	item.Connections++
}

func sortTrafficItems(m map[string]*TrafficItem) []TrafficItem {
	out := make([]TrafficItem, 0, len(m))
	for _, v := range m {
		out = append(out, *v)
	}
	for i := 0; i < len(out); i++ {
		for j := i + 1; j < len(out); j++ {
			if out[j].Total > out[i].Total {
				out[i], out[j] = out[j], out[i]
			}
		}
	}
	return out
}

func connHost(m map[string]any) string {
	meta, _ := m["metadata"].(map[string]any)
	if meta == nil {
		return ""
	}
	if h, ok := meta["host"].(string); ok && h != "" {
		return h
	}
	if h, ok := meta["sniffHost"].(string); ok && h != "" {
		return h
	}
	if h, ok := meta["remoteDestination"].(string); ok && h != "" {
		return h
	}
	if ip, ok := meta["destinationIP"].(string); ok && ip != "" {
		port, _ := meta["destinationPort"].(string)
		if port != "" {
			return ip + ":" + port
		}
		return ip
	}
	return ""
}

func connLeafNode(m map[string]any, policyGroups map[string]bool) string {
	return mihomoapi.ResolveLeafOutbound(mihomoapi.ChainStringsFromAny(m["chains"]), policyGroups)
}

func toInt64(v any) int64 {
	switch n := v.(type) {
	case float64:
		return int64(n)
	case int64:
		return n
	case int:
		return int64(n)
	default:
		return 0
	}
}

func ApplyTrafficLabels(items []TrafficItem, names map[string]string) []TrafficItem {
	out := make([]TrafficItem, len(items))
	for i, item := range items {
		out[i] = item
		if label, ok := names[item.Key]; ok && label != "" {
			out[i].Label = label
		}
	}
	return out
}

func ApplyDomainLabels(items []TrafficItem, names map[string]string) []TrafficItem {
	out := make([]TrafficItem, len(items))
	for i, item := range items {
		out[i] = item
		if label, ok := names[item.NodeKey]; ok && label != "" {
			out[i].Node = label
		} else if item.NodeKey != "" {
			out[i].Node = item.NodeKey
		}
	}
	return out
}

func TrafficAggToItems(aggs []store.TrafficAgg, names map[string]string) []TrafficItem {
	out := make([]TrafficItem, 0, len(aggs))
	for _, a := range aggs {
		label := a.Key
		if n, ok := names[a.Key]; ok && n != "" {
			label = n
		}
		out = append(out, TrafficItem{
			Key:         a.Key,
			Label:       label,
			Upload:      a.Upload,
			Download:    a.Download,
			Total:       a.Upload + a.Download,
			Connections: a.Connections,
		})
	}
	return out
}

func TrafficHostNodeAggToItems(aggs []store.TrafficHostNodeAgg, names map[string]string) []TrafficItem {
	out := make([]TrafficItem, 0, len(aggs))
	for _, a := range aggs {
		nodeLabel := a.NodeKey
		if n, ok := names[a.NodeKey]; ok && n != "" {
			nodeLabel = n
		}
		out = append(out, TrafficItem{
			Key:         a.Host + "\x00" + a.NodeKey,
			Label:       a.Host,
			Node:        nodeLabel,
			NodeKey:     a.NodeKey,
			Upload:      a.Upload,
			Download:    a.Download,
			Total:       a.Upload + a.Download,
			Connections: a.Connections,
		})
	}
	return out
}