feat: Prism HTTP/SOCKS5 代理网关及管理台

Co-authored-by: Cursor <cursoragent@cursor.com>

internal/trafficlog/collector.go

@@ -0,0 +1,227 @@
+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)
+}