3x-ui-agent/internal/agent/singbox_stats.go

package agent

import (
	"context"
	"errors"
	"fmt"
	"regexp"
	"time"

	"github.com/mhsanaei/3x-ui/v3/internal/logger"
	"github.com/mhsanaei/3x-ui/v3/internal/xray"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/encoding"
	"google.golang.org/protobuf/encoding/protowire"
)

// ITEM6 B3: per-user traffic accounting for the sing-box runtime.
//
// §4.4 of the scope offered two shapes: embed sing-box's v2rayapi Go
// package, or ship NexusGate's stats-helper binary. This client takes
// the first shape without the dependency: sing-box's v2ray_api is a
// different protobuf service from xray's, but its messages are just
// as small and stable ({patterns} in, {name, value} pairs out), so —
// exactly as this build already does for xray in internal/xray/api.go
// after dropping the core dependency — the client speaks it at the
// protobuf wire level with protowire. Verified against the real
// v1.13.21 binary in the B3 end-to-end test; the helper binary remains
// the documented fallback if the wire shape ever moves.
//
// Counter names are identical to xray's
// (inbound>>>tag>>>traffic>>>uplink, user>>>name>>>traffic>>>downlink),
// and the master-side write path keys users by the same name (=email),
// so only this agent-side query implementation changed.
//
// Epoch discipline (NexusGate): counters belong to one core process
// instance. A snapshot is only meaningful when the core's epoch is
// unchanged across the query — a sample straddling a restart mixes
// two instances' counters and is discarded whole. When the epoch has
// flipped since the committed baseline, the new instance's counters
// started at zero, so they count IN FULL as fresh deltas (the panel
// side's ApplyDelta rule, and the same outcome as the xray path's
// "counter below baseline counts from zero"). Together with
// flush-before-change in the activation path — which reports the
// old instance's tail before it dies — a core restart neither
// double-reports nor loses the new instance's traffic. The one
// conservative exception: the very first snapshot an agent ever
// takes is baseline-only (it may be reading counters that predate
// the agent), exactly like the xray path's first poll.

const singboxStatsServiceQuery = "/v2ray.core.app.stats.command.StatsService/QueryStats"

// errSingboxEpochChanged aborts a snapshot whose core instance moved
// mid-query (or is unknown); the next poll retries against the new
// instance. It is a routine event around config changes, not a fault.
var errSingboxEpochChanged = errors.New("sing-box restarted during stats query")

var (
	singboxTrafficRegex       = regexp.MustCompile(`(inbound|outbound)>>>([^>]+)>>>traffic>>>(downlink|uplink)`)
	singboxClientTrafficRegex = regexp.MustCompile(`user>>>([^>]+)>>>traffic>>>(downlink|uplink)`)
)

// sbRawBytes / sbRawCodec mirror the xray package's raw protobuf
// pass-through codec (the standard "proto" content subtype).
type sbRawBytes []byte

type sbRawCodec struct{}

func (sbRawCodec) Marshal(v any) ([]byte, error) {
	b, ok := v.(sbRawBytes)
	if !ok {
		return nil, fmt.Errorf("sing-box stats: raw codec cannot marshal %T", v)
	}
	return b, nil
}

func (sbRawCodec) Unmarshal(data []byte, v any) error {
	p, ok := v.(*sbRawBytes)
	if !ok {
		return fmt.Errorf("sing-box stats: raw codec cannot unmarshal into %T", v)
	}
	*p = append((*p)[:0], data...)
	return nil
}

func (sbRawCodec) Name() string { return "proto" }

var _ encoding.Codec = sbRawCodec{}

// singboxStatEntry is one counter from a QueryStats response
// (v2ray.core.app.stats.command.Stat: name = 1, value = 2).
type singboxStatEntry struct {
	Name  string
	Value int64
}

// encodeSingboxQuery builds the QueryStatsRequest wire form:
// patterns is field 3 (repeated string). Asking for the three counter
// families keeps the response free of anything else the core counts.
func encodeSingboxQuery() []byte {
	var b []byte
	for _, pattern := range []string{"inbound>>>", "outbound>>>", "user>>>"} {
		b = protowire.AppendTag(b, 3, protowire.BytesType)
		b = protowire.AppendString(b, pattern)
	}
	return b
}

// parseSingboxStatsResponse decodes a QueryStatsResponse
// (field 1: repeated Stat).
func parseSingboxStatsResponse(b []byte) ([]singboxStatEntry, error) {
	var out []singboxStatEntry
	for len(b) > 0 {
		num, typ, n := protowire.ConsumeTag(b)
		if n < 0 {
			return nil, protowire.ParseError(n)
		}
		b = b[n:]
		if num == 1 && typ == protowire.BytesType {
			v, n := protowire.ConsumeBytes(b)
			if n < 0 {
				return nil, protowire.ParseError(n)
			}
			st, err := parseSingboxStat(v)
			if err != nil {
				return nil, err
			}
			out = append(out, st)
			b = b[n:]
			continue
		}
		n = protowire.ConsumeFieldValue(num, typ, b)
		if n < 0 {
			return nil, protowire.ParseError(n)
		}
		b = b[n:]
	}
	return out, nil
}

func parseSingboxStat(b []byte) (singboxStatEntry, error) {
	var st singboxStatEntry
	for len(b) > 0 {
		num, typ, n := protowire.ConsumeTag(b)
		if n < 0 {
			return st, protowire.ParseError(n)
		}
		b = b[n:]
		switch {
		case num == 1 && typ == protowire.BytesType:
			v, n := protowire.ConsumeBytes(b)
			if n < 0 {
				return st, protowire.ParseError(n)
			}
			st.Name = string(v)
			b = b[n:]
		case num == 2 && typ == protowire.VarintType:
			v, n := protowire.ConsumeVarint(b)
			if n < 0 {
				return st, protowire.ParseError(n)
			}
			st.Value = int64(v)
			b = b[n:]
		default:
			n := protowire.ConsumeFieldValue(num, typ, b)
			if n < 0 {
				return st, protowire.ParseError(n)
			}
			b = b[n:]
		}
	}
	return st, nil
}

// singboxStatsClient is the sing-box counterpart of xray.XrayAPI for
// the agent reporting path: same snapshot/commit contract (snapshots
// never advance the baseline; only an acknowledged report commits),
// plus the epoch guard described above. Not safe for concurrent use;
// the agent serializes it under sbMu.
type singboxStatsClient struct {
	addr      string
	epochFunc func() string

	conn *grpc.ClientConn

	committed      map[string]int64
	committedReady bool
	committedEpoch string
	pendingEpoch   string

	// queryFn replaces the gRPC query in tests.
	queryFn func(ctx context.Context) ([]singboxStatEntry, error)
}

func newSingboxStatsClient(epochFunc func() string) *singboxStatsClient {
	return &singboxStatsClient{epochFunc: epochFunc}
}

// SetAddr points the client at the v2ray_api listen address of the
// currently applied config. A changed address drops the connection.
func (c *singboxStatsClient) SetAddr(addr string) {
	if addr == c.addr {
		return
	}
	c.addr = addr
	c.closeConn()
}

func (c *singboxStatsClient) closeConn() {
	if c.conn != nil {
		_ = c.conn.Close()
		c.conn = nil
	}
}

func (c *singboxStatsClient) queryStats(ctx context.Context) ([]singboxStatEntry, error) {
	if c.queryFn != nil {
		return c.queryFn(ctx)
	}
	if c.addr == "" {
		return nil, errors.New("sing-box stats address unknown (config carries no v2ray_api listen)")
	}
	if c.conn == nil {
		conn, err := grpc.NewClient(c.addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
		if err != nil {
			return nil, fmt.Errorf("connect sing-box stats API: %w", err)
		}
		c.conn = conn
	}
	var resp sbRawBytes
	if err := c.conn.Invoke(ctx, singboxStatsServiceQuery, sbRawBytes(encodeSingboxQuery()), &resp, grpc.ForceCodec(sbRawCodec{})); err != nil {
		c.closeConn()
		return nil, err
	}
	return parseSingboxStatsResponse(resp)
}

// GetTrafficSnapshot reads the core counters and computes deltas
// against the committed baseline WITHOUT advancing it, under the
// epoch rules: an epoch that moves mid-query fails the round; a
// flipped epoch counts the new instance's counters in full; the
// agent's first-ever snapshot is baseline-only.
func (c *singboxStatsClient) GetTrafficSnapshot() (*xray.TrafficSnapshot, error) {
	epochBefore := c.epochFunc()
	if epochBefore == "" {
		return nil, errSingboxEpochChanged
	}
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	stats, err := c.queryStats(ctx)
	cancel()
	if err != nil {
		logger.Debug("Failed to query sing-box stats:", err)
		return nil, err
	}
	if epochAfter := c.epochFunc(); epochAfter != epochBefore {
		return nil, errSingboxEpochChanged
	}
	c.pendingEpoch = epochBefore

	snap := &xray.TrafficSnapshot{Values: make(map[string]int64, len(stats))}
	sameEpoch := c.committedReady && c.committedEpoch == epochBefore
	tagTrafficMap := make(map[string]*xray.Traffic)
	emailTrafficMap := make(map[string]*xray.ClientTraffic)

	for _, stat := range stats {
		snap.Values[stat.Name] = stat.Value
		if !c.committedReady {
			continue
		}
		var value int64
		if sameEpoch {
			lastValue, ok := c.committed[stat.Name]
			if !ok || stat.Value < lastValue {
				lastValue = 0
			}
			value = stat.Value - lastValue
		} else {
			// New core instance: its counters started at zero, so
			// everything they show is traffic since the restart.
			value = stat.Value
		}
		if matches := singboxTrafficRegex.FindStringSubmatch(stat.Name); len(matches) == 4 {
			processSingboxTraffic(matches, value, tagTrafficMap)
		} else if matches := singboxClientTrafficRegex.FindStringSubmatch(stat.Name); len(matches) == 3 {
			processSingboxClientTraffic(matches, value, emailTrafficMap)
		}
	}

	snap.Traffics = mapToSliceXrayTraffic(tagTrafficMap)
	snap.ClientTraffics = mapToSliceXrayClientTraffic(emailTrafficMap)
	return snap, nil
}

// CommitTrafficSnapshot marks a snapshot as acknowledged by the
// master: its Values and epoch become the committed baseline.
func (c *singboxStatsClient) CommitTrafficSnapshot(s *xray.TrafficSnapshot) {
	if s == nil {
		return
	}
	c.committed = s.BaselineValues()
	c.committedReady = true
	c.committedEpoch = c.pendingEpoch
}

// RestoreTrafficBaseline installs a committed baseline recovered from
// disk together with the epoch it was taken under. A baseline from a
// different core instance is dropped and the first snapshot rebases
// baseline-only: unlike a live epoch flip the agent observes itself,
// an agent restart leaves a blind gap — the running core's counters
// may include traffic the previous agent incarnation already
// reported, so counting them in full could double-report.
func (c *singboxStatsClient) RestoreTrafficBaseline(values map[string]int64, epoch string) {
	if len(values) == 0 || epoch == "" || epoch != c.epochFunc() {
		return
	}
	c.committed = make(map[string]int64, len(values))
	for k, v := range values {
		c.committed[k] = v
	}
	c.committedReady = true
	c.committedEpoch = epoch
}

// GetOnlineUsers is unused for sing-box (the traffic cycle merges the
// log tracker's view with the snapshot's active users itself); it
// exists so the client satisfies the same shape as the xray one.
func (c *singboxStatsClient) GetOnlineUsers() ([]xray.OnlineUser, error) {
	return nil, nil
}

func processSingboxTraffic(matches []string, value int64, trafficMap map[string]*xray.Traffic) {
	isInbound := matches[1] == "inbound"
	tag := matches[2]
	isDown := matches[3] == "downlink"
	if tag == "api" {
		return
	}
	traffic, ok := trafficMap[tag]
	if !ok {
		traffic = &xray.Traffic{IsInbound: isInbound, IsOutbound: !isInbound, Tag: tag}
		trafficMap[tag] = traffic
	}
	if isDown {
		traffic.Down = value
	} else {
		traffic.Up = value
	}
}

func processSingboxClientTraffic(matches []string, value int64, clientTrafficMap map[string]*xray.ClientTraffic) {
	email := matches[1]
	isDown := matches[2] == "downlink"
	traffic, ok := clientTrafficMap[email]
	if !ok {
		traffic = &xray.ClientTraffic{Email: email}
		clientTrafficMap[email] = traffic
	}
	if isDown {
		traffic.Down = value
	} else {
		traffic.Up = value
	}
}

func mapToSliceXrayTraffic(m map[string]*xray.Traffic) []*xray.Traffic {
	result := make([]*xray.Traffic, 0, len(m))
	for _, v := range m {
		result = append(result, v)
	}
	return result
}

func mapToSliceXrayClientTraffic(m map[string]*xray.ClientTraffic) []*xray.ClientTraffic {
	result := make([]*xray.ClientTraffic, 0, len(m))
	for _, v := range m {
		result = append(result, v)
	}
	return result
}