Prism/internal/trafficlog/collector.go

package trafficlog

import (
	"context"
	"log"
	"strings"
	"sync"
	"time"

	"github.com/prism/proxy/internal/country"
	"github.com/prism/proxy/internal/mihomoapi"
	"github.com/prism/proxy/internal/store"
)

type Collector struct {
	store    *store.Store
	mu       sync.Mutex
	client   *mihomoapi.Client
	active   map[string]mihomoapi.Connection
	names    map[string]string
	groups   map[string]bool
	keep     int

	lastPollErr   string
	lastPollErrAt time.Time
}

func NewCollector(s *store.Store, keep int) *Collector {
	if keep <= 0 {
		keep = 5000
	}
	return &Collector{
		store:  s,
		active: make(map[string]mihomoapi.Connection),
		names:  make(map[string]string),
		keep:   keep,
	}
}

func (c *Collector) UpdateClient(client *mihomoapi.Client) {
	c.mu.Lock()
	c.client = client
	c.mu.Unlock()
}

func (c *Collector) Run(ctx context.Context) {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			c.poll(ctx)
		}
	}
}

func (c *Collector) poll(ctx context.Context) {
	c.mu.Lock()
	client := c.client
	c.mu.Unlock()
	if client == nil {
		return
	}
	snap, err := client.ConnectionSnapshot()
	if err != nil {
		c.logPollError(err)
		return
	}
	if snap == nil {
		return
	}
	current := make(map[string]mihomoapi.Connection, len(snap.Connections))
	for _, conn := range snap.Connections {
		if conn.ID == "" {
			continue
		}
		current[conn.ID] = conn
	}

	c.mu.Lock()
	groups := c.groups
	names := c.names
	for id, conn := range current {
		c.active[id] = conn
	}
	var ended []mihomoapi.Connection
	for id, conn := range c.active {
		if _, ok := current[id]; ok {
			continue
		}
		ended = append(ended, conn)
		delete(c.active, id)
	}
	c.mu.Unlock()

	for _, conn := range ended {
		c.save(ctx, conn, groups, names)
	}
}

func (c *Collector) save(ctx context.Context, conn mihomoapi.Connection, groups map[string]bool, names map[string]string) {
	nodeKey := conn.LeafOutbound(groups)
	node := c.resolveName(nodeKey, names)
	chain := conn.ChainString()
	resolvedChain := c.resolveChainJSON(chain)
	if err := c.store.InsertProxyRequestLog(ctx, &store.ProxyRequestLog{
		Host:          conn.RequestHost(),
		Network:       conn.Metadata.Network,
		RuleType:      conn.Rule,
		RulePayload:   conn.RulePayload,
		RuleLine:      conn.RuleLine(),
		OutboundNode:  node,
		OutboundChain: resolvedChain,
		Upload:        conn.Upload,
		Download:      conn.Download,
		Success:       conn.IsSuccess(),
		StartedAt:     conn.Start,
		EndedAt:       time.Now(),
	}); err != nil {
		log.Printf("trafficlog: save request log: %v", err)
		return
	}
	_ = c.store.PruneProxyRequestLogs(ctx, c.keep)
}

func (c *Collector) logPollError(err error) {
	msg := err.Error()
	now := time.Now()
	c.mu.Lock()
	defer c.mu.Unlock()
	if msg == c.lastPollErr && now.Sub(c.lastPollErrAt) < time.Minute {
		return
	}
	c.lastPollErr = msg
	c.lastPollErrAt = now
	log.Printf("trafficlog: poll connections: %v", err)
}

func (c *Collector) refreshNames(ctx context.Context) {
	nodes, err := c.store.ListProxyNodes(ctx, "", nil)
	if err != nil {
		return
	}
	groups, _ := c.store.ListProxyGroups(ctx)
	m := map[string]string{
		"DIRECT": "DIRECT",
		"REJECT": "REJECT",
	}
	g := map[string]bool{
		"GLOBAL": true,
		"PROXY":  true,
	}
	for _, n := range nodes {
		m[n.RuntimeName] = n.Name
	}
	for _, grp := range groups {
		if grp.Enabled {
			g[grp.RuntimeName] = true
		}
		m[grp.RuntimeName] = grp.Name
	}
	seen := map[string]bool{}
	for _, n := range nodes {
		if n.Country == "" || seen[n.Country] {
			continue
		}
		seen[n.Country] = true
		rn := country.GroupRuntime(n.Country)
		m[rn] = country.Name(n.Country) + "（国家池）"
	}
	c.mu.Lock()
	c.names = m
	c.groups = g
	c.mu.Unlock()
}

func (c *Collector) resolveName(runtime string, names map[string]string) string {
	if runtime == "" {
		return ""
	}
	if name, ok := names[runtime]; ok && name != "" {
		return name
	}
	if strings.HasPrefix(runtime, "country_") {
		code := strings.TrimPrefix(runtime, "country_")
		return code + "（国家池）"
	}
	return runtime
}

func (c *Collector) resolveChainJSON(chain string) string {
	// chain stored as JSON array string; map runtime names to display names
	if chain == "" || chain == "[]" {
		return chain
	}
	// lightweight replace for display
	out := chain
	c.mu.Lock()
	names := c.names
	c.mu.Unlock()
	for rt, name := range names {
		if rt != "" && name != "" {
			out = strings.ReplaceAll(out, rt, name)
		}
	}
	return out
}

func (c *Collector) Start(ctx context.Context, client *mihomoapi.Client) {
	c.UpdateClient(client)
	c.refreshNames(ctx)
	go func() {
		t := time.NewTicker(30 * time.Second)
		defer t.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-t.C:
				c.refreshNames(ctx)
			}
		}
	}()
	go c.Run(ctx)
}