流量统计

.gitignore

@@ -188,9 +188,6 @@ go.work.sum
 # 编译器输出
 vendor/
 
-# 流量统计或其他运行时生成的文件
-traffic/
-
 # ==========================================
 # Docker
 # ==========================================

go-gost/x/limiter/traffic/cache/limiter.go

@@ -0,0 +1,161 @@
+package limiter
+
+import (
+	"context"
+	"time"
+
+	"github.com/go-gost/core/limiter"
+	"github.com/go-gost/core/limiter/traffic"
+	"github.com/go-gost/x/internal/util/cache"
+)
+
+const (
+	defaultRefreshInterval = 30 * time.Second
+	defaultCleanupInterval = 60 * time.Second
+)
+
+type options struct {
+	refreshInterval time.Duration
+	cleanupInterval time.Duration
+	scope           string
+}
+
+type Option func(*options)
+
+func RefreshIntervalOption(interval time.Duration) Option {
+	return func(o *options) {
+		o.refreshInterval = interval
+	}
+}
+
+func CleanupIntervalOption(interval time.Duration) Option {
+	return func(o *options) {
+		o.cleanupInterval = interval
+	}
+}
+
+func ScopeOption(scope string) Option {
+	return func(o *options) {
+		o.scope = scope
+	}
+}
+
+type cachedTrafficLimiter struct {
+	inLimits  *cache.Cache
+	outLimits *cache.Cache
+	limiter   traffic.TrafficLimiter
+	options   options
+}
+
+func NewCachedTrafficLimiter(limiter traffic.TrafficLimiter, opts ...Option) traffic.TrafficLimiter {
+	if limiter == nil {
+		return nil
+	}
+
+	var options options
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	if options.refreshInterval == 0 {
+		options.refreshInterval = defaultRefreshInterval
+	}
+	if options.refreshInterval < time.Second {
+		options.refreshInterval = time.Second
+	}
+
+	if options.cleanupInterval == 0 {
+		options.cleanupInterval = defaultCleanupInterval
+	}
+	if options.cleanupInterval < time.Second {
+		options.cleanupInterval = time.Second
+	}
+
+	lim := &cachedTrafficLimiter{
+		inLimits:  cache.NewCache(options.cleanupInterval),
+		outLimits: cache.NewCache(options.cleanupInterval),
+		limiter:   limiter,
+		options:   options,
+	}
+	return lim
+}
+
+func (p *cachedTrafficLimiter) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
+	if p.limiter == nil {
+		return nil
+	}
+
+	var options limiter.Options
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	if p.options.scope != "" && p.options.scope != options.Scope {
+		return nil
+	}
+
+	item := p.inLimits.Get(key)
+	lim, _ := item.Value().(traffic.Limiter)
+	if !item.Expired() {
+		return lim
+	}
+
+	limNew := p.limiter.In(ctx, key, opts...)
+	if limNew == nil {
+		limNew = lim
+	}
+	if item == nil || !p.equal(lim, limNew) {
+		p.inLimits.Set(key, cache.NewItem(limNew, p.options.refreshInterval))
+		return limNew
+	}
+
+	p.inLimits.Set(key, cache.NewItem(lim, p.options.refreshInterval))
+
+	return lim
+}
+
+func (p *cachedTrafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
+	if p.limiter == nil {
+		return nil
+	}
+
+	var options limiter.Options
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	if p.options.scope != "" && p.options.scope != options.Scope {
+		return nil
+	}
+
+	item := p.outLimits.Get(key)
+	lim, _ := item.Value().(traffic.Limiter)
+	if !item.Expired() {
+		return lim
+	}
+
+	limNew := p.limiter.Out(ctx, key, opts...)
+	if limNew == nil {
+		limNew = lim
+	}
+	if item == nil || !p.equal(lim, limNew) {
+		p.outLimits.Set(key, cache.NewItem(limNew, p.options.refreshInterval))
+		return limNew
+	}
+
+	p.outLimits.Set(key, cache.NewItem(lim, p.options.refreshInterval))
+
+	return lim
+}
+
+func (p *cachedTrafficLimiter) equal(lim1, lim2 traffic.Limiter) bool {
+	if lim1 == lim2 {
+		return true
+	}
+
+	if lim1 == nil || lim2 == nil {
+		return false
+	}
+
+	return lim1.Limit() == lim2.Limit()
+}

go-gost/x/limiter/traffic/generator.go

@@ -0,0 +1,31 @@
+package traffic
+
+import (
+	limiter "github.com/go-gost/core/limiter/traffic"
+)
+
+type limitGenerator struct {
+	in  int
+	out int
+}
+
+func newLimitGenerator(in, out int) *limitGenerator {
+	return &limitGenerator{
+		in:  in,
+		out: out,
+	}
+}
+
+func (p *limitGenerator) In() limiter.Limiter {
+	if p == nil || p.in <= 0 {
+		return nil
+	}
+	return NewLimiter(p.in)
+}
+
+func (p *limitGenerator) Out() limiter.Limiter {
+	if p == nil || p.out <= 0 {
+		return nil
+	}
+	return NewLimiter(p.out)
+}

go-gost/x/limiter/traffic/limiter.go

@@ -0,0 +1,76 @@
+package traffic
+
+import (
+	"context"
+	"fmt"
+	"sort"
+	"strconv"
+
+	limiter "github.com/go-gost/core/limiter/traffic"
+	"golang.org/x/time/rate"
+)
+
+type llimiter struct {
+	limiter *rate.Limiter
+}
+
+func NewLimiter(r int) limiter.Limiter {
+	return &llimiter{
+		limiter: rate.NewLimiter(rate.Limit(r), r),
+	}
+}
+
+func (l *llimiter) Wait(ctx context.Context, n int) int {
+	if l.limiter.Burst() < n {
+		n = l.limiter.Burst()
+	}
+	l.limiter.WaitN(ctx, n)
+	return n
+}
+
+func (l *llimiter) Limit() int {
+	return int(l.limiter.Limit())
+}
+
+func (l *llimiter) Set(n int) {
+	l.limiter.SetLimit(rate.Limit(n))
+	l.limiter.SetBurst(n)
+}
+
+func (l *llimiter) String() string {
+	return strconv.Itoa(int(l.limiter.Limit()))
+}
+
+type limiterGroup struct {
+	limiters []limiter.Limiter
+}
+
+func newLimiterGroup(limiters ...limiter.Limiter) *limiterGroup {
+	sort.Slice(limiters, func(i, j int) bool {
+		return limiters[i].Limit() < limiters[j].Limit()
+	})
+	return &limiterGroup{limiters: limiters}
+}
+
+func (l *limiterGroup) Wait(ctx context.Context, n int) int {
+	for i := range l.limiters {
+		if v := l.limiters[i].Wait(ctx, n); v < n {
+			n = v
+		}
+	}
+	return n
+}
+
+func (l *limiterGroup) Limit() int {
+	if len(l.limiters) == 0 {
+		return 0
+	}
+
+	return l.limiters[0].Limit()
+}
+
+func (l *limiterGroup) Set(n int) {}
+
+func (l *limiterGroup) String() string {
+	return fmt.Sprintf("%v", l.limiters)
+}

go-gost/x/limiter/traffic/plugin/grpc.go

@@ -0,0 +1,111 @@
+package traffic
+
+import (
+	"context"
+	"io"
+
+	"github.com/go-gost/core/limiter"
+	"github.com/go-gost/core/limiter/traffic"
+	"github.com/go-gost/core/logger"
+	"github.com/go-gost/plugin/limiter/traffic/proto"
+	"github.com/go-gost/x/internal/plugin"
+	xtraffic "github.com/go-gost/x/limiter/traffic"
+	"google.golang.org/grpc"
+)
+
+type grpcPlugin struct {
+	conn   grpc.ClientConnInterface
+	client proto.LimiterClient
+	log    logger.Logger
+}
+
+// NewGRPCPlugin creates a traffic limiter plugin based on gRPC.
+func NewGRPCPlugin(name string, addr string, opts ...plugin.Option) traffic.TrafficLimiter {
+	var options plugin.Options
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	log := logger.Default().WithFields(map[string]any{
+		"kind":    "limiter",
+		"limiter": name,
+	})
+	conn, err := plugin.NewGRPCConn(addr, &options)
+	if err != nil {
+		log.Error(err)
+	}
+
+	p := &grpcPlugin{
+		conn: conn,
+		log:  log,
+	}
+	if conn != nil {
+		p.client = proto.NewLimiterClient(conn)
+	}
+	return p
+}
+
+func (p *grpcPlugin) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
+	if p.client == nil {
+		return nil
+	}
+
+	var options limiter.Options
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	r, err := p.client.Limit(ctx,
+		&proto.LimitRequest{
+			Service: options.Service,
+			Scope:   options.Scope,
+			Network: options.Network,
+			Addr:    options.Addr,
+			Client:  options.Client,
+			Src:     options.Src,
+		})
+	if err != nil {
+		p.log.Error(err)
+		return nil
+	}
+
+	return xtraffic.NewLimiter(int(r.In))
+}
+
+func (p *grpcPlugin) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
+	if p.client == nil {
+		return nil
+	}
+
+	var options limiter.Options
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	r, err := p.client.Limit(ctx,
+		&proto.LimitRequest{
+			Service: options.Service,
+			Scope:   options.Scope,
+			Network: options.Network,
+			Addr:    options.Addr,
+			Client:  options.Client,
+			Src:     options.Src,
+		})
+	if err != nil {
+		p.log.Error(err)
+		return nil
+	}
+
+	return xtraffic.NewLimiter(int(r.Out))
+}
+
+func (p *grpcPlugin) Close() error {
+	if p.conn == nil {
+		return nil
+	}
+
+	if closer, ok := p.conn.(io.Closer); ok {
+		return closer.Close()
+	}
+	return nil
+}

go-gost/x/limiter/traffic/plugin/http.go

@@ -0,0 +1,161 @@
+package traffic
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"net/http"
+
+	"github.com/go-gost/core/limiter"
+	"github.com/go-gost/core/limiter/traffic"
+	"github.com/go-gost/core/logger"
+	"github.com/go-gost/x/internal/plugin"
+	xtraffic "github.com/go-gost/x/limiter/traffic"
+)
+
+type httpPluginRequest struct {
+	Service string `json:"service"`
+	Scope   string `json:"scope"`
+	Network string `json:"network"`
+	Addr    string `json:"addr"`
+	Client  string `json:"client"`
+	Src     string `json:"src"`
+}
+
+type httpPluginResponse struct {
+	In  int64 `json:"in"`
+	Out int64 `json:"out"`
+}
+
+type httpPlugin struct {
+	url    string
+	client *http.Client
+	header http.Header
+	log    logger.Logger
+}
+
+// NewHTTPPlugin creates a traffic limiter plugin based on HTTP.
+func NewHTTPPlugin(name string, url string, opts ...plugin.Option) traffic.TrafficLimiter {
+	var options plugin.Options
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	return &httpPlugin{
+		url:    url,
+		client: plugin.NewHTTPClient(&options),
+		header: options.Header,
+		log: logger.Default().WithFields(map[string]any{
+			"kind":    "limiter",
+			"limiter": name,
+		}),
+	}
+}
+
+func (p *httpPlugin) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
+	if p.client == nil {
+		return nil
+	}
+
+	var options limiter.Options
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	rb := httpPluginRequest{
+		Service: options.Service,
+		Scope:   options.Scope,
+		Network: options.Network,
+		Addr:    options.Addr,
+		Client:  options.Client,
+		Src:     options.Src,
+	}
+	v, err := json.Marshal(&rb)
+	if err != nil {
+		p.log.Error(err)
+		return nil
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.url, bytes.NewReader(v))
+	if err != nil {
+		p.log.Error(err)
+		return nil
+	}
+
+	if p.header != nil {
+		req.Header = p.header.Clone()
+	}
+	req.Header.Set("Content-Type", "application/json")
+	resp, err := p.client.Do(req)
+	if err != nil {
+		p.log.Error(err)
+		return nil
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		p.log.Errorf("server return non-200 code: %s", resp.Status)
+		return nil
+	}
+
+	res := httpPluginResponse{}
+	if err := json.NewDecoder(resp.Body).Decode(&res); err != nil {
+		p.log.Error(err)
+		return nil
+	}
+	return xtraffic.NewLimiter(int(res.In))
+}
+
+func (p *httpPlugin) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
+	if p.client == nil {
+		return nil
+	}
+
+	var options limiter.Options
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	rb := httpPluginRequest{
+		Service: options.Service,
+		Scope:   options.Scope,
+		Network: options.Network,
+		Addr:    options.Addr,
+		Client:  options.Client,
+		Src:     options.Src,
+	}
+	v, err := json.Marshal(&rb)
+	if err != nil {
+		p.log.Error(err)
+		return nil
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.url, bytes.NewReader(v))
+	if err != nil {
+		p.log.Error(err)
+		return nil
+	}
+
+	if p.header != nil {
+		req.Header = p.header.Clone()
+	}
+	req.Header.Set("Content-Type", "application/json")
+	resp, err := p.client.Do(req)
+	if err != nil {
+		p.log.Error(err)
+		return nil
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		p.log.Errorf("server return non-200 code: %s", resp.Status)
+		return nil
+	}
+
+	res := httpPluginResponse{}
+	if err := json.NewDecoder(resp.Body).Decode(&res); err != nil {
+		p.log.Error(err)
+		return nil
+	}
+	return xtraffic.NewLimiter(int(res.Out))
+}

go-gost/x/limiter/traffic/traffic.go

@@ -0,0 +1,626 @@
+package traffic
+
+import (
+	"bufio"
+	"context"
+	"io"
+	"net"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/alecthomas/units"
+	"github.com/go-gost/core/limiter"
+	"github.com/go-gost/core/limiter/traffic"
+	"github.com/go-gost/core/logger"
+	"github.com/go-gost/x/internal/loader"
+	"github.com/patrickmn/go-cache"
+	"github.com/yl2chen/cidranger"
+)
+
+const (
+	ServiceLimitKey = "$"
+	ConnLimitKey    = "$$"
+)
+
+const (
+	defaultExpiration = 15 * time.Second
+	cleanupInterval   = 30 * time.Second
+)
+
+type options struct {
+	limits      []string
+	fileLoader  loader.Loader
+	redisLoader loader.Loader
+	httpLoader  loader.Loader
+	period      time.Duration
+	logger      logger.Logger
+}
+
+type Option func(opts *options)
+
+func LimitsOption(limits ...string) Option {
+	return func(opts *options) {
+		opts.limits = limits
+	}
+}
+
+func ReloadPeriodOption(period time.Duration) Option {
+	return func(opts *options) {
+		opts.period = period
+	}
+}
+
+func FileLoaderOption(fileLoader loader.Loader) Option {
+	return func(opts *options) {
+		opts.fileLoader = fileLoader
+	}
+}
+
+func RedisLoaderOption(redisLoader loader.Loader) Option {
+	return func(opts *options) {
+		opts.redisLoader = redisLoader
+	}
+}
+
+func HTTPLoaderOption(httpLoader loader.Loader) Option {
+	return func(opts *options) {
+		opts.httpLoader = httpLoader
+	}
+}
+
+func LoggerOption(logger logger.Logger) Option {
+	return func(opts *options) {
+		opts.logger = logger
+	}
+}
+
+type limitValue struct {
+	in  int
+	out int
+}
+
+type trafficLimiter struct {
+	generators     sync.Map
+	cidrGenerators cidranger.Ranger
+	// connection level in/out limits
+	connInLimits  *cache.Cache
+	connOutLimits *cache.Cache
+	// service level in/out limits
+	inLimits   *cache.Cache
+	outLimits  *cache.Cache
+	mu         sync.RWMutex
+	cancelFunc context.CancelFunc
+	options    options
+}
+
+func NewTrafficLimiter(opts ...Option) traffic.TrafficLimiter {
+	var options options
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	ctx, cancel := context.WithCancel(context.TODO())
+	lim := &trafficLimiter{
+		cidrGenerators: cidranger.NewPCTrieRanger(),
+		connInLimits:   cache.New(defaultExpiration, cleanupInterval),
+		connOutLimits:  cache.New(defaultExpiration, cleanupInterval),
+		inLimits:       cache.New(defaultExpiration, cleanupInterval),
+		outLimits:      cache.New(defaultExpiration, cleanupInterval),
+		options:        options,
+		cancelFunc:     cancel,
+	}
+
+	if err := lim.reload(ctx); err != nil {
+		options.logger.Warnf("reload: %v", err)
+	}
+	if lim.options.period > 0 {
+		go lim.periodReload(ctx)
+	}
+	return lim
+}
+
+// In obtains a traffic input limiter based on key.
+// For connection scope, the key should be client connection address.
+func (l *trafficLimiter) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
+	var options limiter.Options
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	switch options.Scope {
+	case limiter.ScopeService:
+		if lim, ok := l.inLimits.Get(ServiceLimitKey); ok && lim != nil {
+			return lim.(traffic.Limiter)
+		}
+		return nil
+
+	case limiter.ScopeClient:
+		return nil
+
+	case limiter.ScopeConn:
+		fallthrough
+	default:
+	}
+
+	var lims []traffic.Limiter
+
+	// connection level limiter
+	if lim, ok := l.connInLimits.Get(key); ok {
+		if lim != nil {
+			// cached connection level limiter
+			lims = append(lims, lim.(traffic.Limiter))
+			// reset expiration
+			l.connInLimits.Set(key, lim, defaultExpiration)
+		}
+	} else {
+		// generate a new connection level limiter and cache it
+		if v, ok := l.generators.Load(ConnLimitKey); ok && v != nil {
+			lim := v.(*limitGenerator).In()
+			if lim != nil {
+				lims = append(lims, lim)
+				l.connInLimits.Set(key, lim, defaultExpiration)
+			}
+		}
+	}
+
+	host, _, _ := net.SplitHostPort(key)
+	// IP level limiter
+	if lim, ok := l.inLimits.Get(host); ok {
+		// cached IP limiter
+		if lim != nil {
+			lims = append(lims, lim.(traffic.Limiter))
+		}
+	} else {
+		l.mu.RLock()
+		ranger := l.cidrGenerators
+		l.mu.RUnlock()
+
+		// CIDR level limiter
+		if p, _ := ranger.ContainingNetworks(net.ParseIP(host)); len(p) > 0 {
+			if v, _ := p[0].(*cidrLimitEntry); v != nil {
+				if lim := v.generator.In(); lim != nil {
+					lims = append(lims, lim)
+					l.inLimits.Set(host, lim, cache.NoExpiration)
+				}
+			}
+		}
+	}
+
+	var lim traffic.Limiter
+	if len(lims) > 0 {
+		lim = newLimiterGroup(lims...)
+	}
+
+	if lim != nil && l.options.logger != nil {
+		l.options.logger.Debugf("input limit for %s: %s", key, lim)
+	}
+
+	return lim
+}
+
+// Out obtains a traffic output limiter based on key.
+// For connection scope, the key should be client connection address.
+func (l *trafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
+	var options limiter.Options
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	switch options.Scope {
+	case limiter.ScopeService:
+		if lim, ok := l.outLimits.Get(ServiceLimitKey); ok && lim != nil {
+			return lim.(traffic.Limiter)
+		}
+		return nil
+
+	case limiter.ScopeClient:
+		return nil
+
+	case limiter.ScopeConn:
+		fallthrough
+	default:
+	}
+
+	var lims []traffic.Limiter
+
+	// connection level limiter
+	if lim, ok := l.connOutLimits.Get(key); ok {
+		if lim != nil {
+			// cached connection level limiter
+			lims = append(lims, lim.(traffic.Limiter))
+			// reset expiration
+			l.connOutLimits.Set(key, lim, defaultExpiration)
+		}
+	} else {
+		// generate a new connection level limiter
+		if v, ok := l.generators.Load(ConnLimitKey); ok && v != nil {
+			lim := v.(*limitGenerator).Out()
+			if lim != nil {
+				lims = append(lims, lim)
+				l.connOutLimits.Set(key, lim, defaultExpiration)
+			}
+		}
+	}
+
+	host, _, _ := net.SplitHostPort(key)
+	// IP level limiter
+	if lim, ok := l.outLimits.Get(host); ok {
+		if lim != nil {
+			// cached IP level limiter
+			lims = append(lims, lim.(traffic.Limiter))
+		}
+	} else {
+		l.mu.RLock()
+		ranger := l.cidrGenerators
+		l.mu.RUnlock()
+
+		// CIDR level limiter
+		if p, _ := ranger.ContainingNetworks(net.ParseIP(host)); len(p) > 0 {
+			if v, _ := p[0].(*cidrLimitEntry); v != nil {
+				if lim := v.generator.Out(); lim != nil {
+					lims = append(lims, lim)
+					l.outLimits.Set(host, lim, cache.NoExpiration)
+				}
+			}
+		}
+	}
+
+	var lim traffic.Limiter
+	if len(lims) > 0 {
+		lim = newLimiterGroup(lims...)
+	}
+
+	if lim != nil && l.options.logger != nil {
+		l.options.logger.Debugf("output limit for %s: %s", key, lim)
+	}
+
+	return lim
+}
+
+func (l *trafficLimiter) periodReload(ctx context.Context) error {
+	period := l.options.period
+	if period < time.Second {
+		period = time.Second
+	}
+	ticker := time.NewTicker(period)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			if err := l.reload(ctx); err != nil {
+				l.options.logger.Warnf("reload: %v", err)
+				// return err
+			}
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}
+}
+
+func (l *trafficLimiter) reload(ctx context.Context) error {
+	values, err := l.load(ctx)
+	if err != nil {
+		return err
+	}
+
+	// service level limiter, never expired
+	{
+		value := values[ServiceLimitKey]
+		if v, _ := l.inLimits.Get(ServiceLimitKey); v != nil {
+			lim := v.(traffic.Limiter)
+			if value.in <= 0 {
+				l.inLimits.Delete(ServiceLimitKey)
+			} else {
+				lim.Set(value.in)
+			}
+		} else {
+			if value.in > 0 {
+				l.inLimits.Set(ServiceLimitKey, NewLimiter(value.in), cache.NoExpiration)
+			}
+		}
+
+		if v, _ := l.outLimits.Get(ServiceLimitKey); v != nil {
+			lim := v.(traffic.Limiter)
+			if value.out <= 0 {
+				l.outLimits.Delete(ServiceLimitKey)
+			} else {
+				lim.Set(value.out)
+			}
+		} else {
+			if value.out > 0 {
+				l.outLimits.Set(ServiceLimitKey, NewLimiter(value.out), cache.NoExpiration)
+			}
+		}
+		delete(values, ServiceLimitKey)
+	}
+
+	// connection level limiters
+	{
+		value := values[ConnLimitKey]
+
+		var in, out int
+		if v, _ := l.generators.Load(ConnLimitKey); v != nil {
+			in, out = v.(*limitGenerator).in, v.(*limitGenerator).out
+		}
+		l.generators.Store(ConnLimitKey, newLimitGenerator(value.in, value.out))
+
+		if value.in <= 0 {
+			l.connInLimits.Flush()
+		} else {
+			if in != value.in {
+				for _, item := range l.connInLimits.Items() {
+					if v := item.Object; v != nil {
+						v.(traffic.Limiter).Set(in)
+					}
+				}
+			}
+		}
+
+		if value.out <= 0 {
+			l.connOutLimits.Flush()
+		} else {
+			if out != value.out {
+				for _, item := range l.connOutLimits.Items() {
+					if v := item.Object; v != nil {
+						v.(traffic.Limiter).Set(out)
+					}
+				}
+			}
+		}
+		delete(values, ConnLimitKey)
+	}
+
+	cidrGenerators := cidranger.NewPCTrieRanger()
+	// IP/CIDR level limiters
+	{
+		// snapshot of the current limiters
+		inLimits := l.inLimits.Items()
+		outLimits := l.outLimits.Items()
+
+		delete(inLimits, ServiceLimitKey)
+		delete(outLimits, ServiceLimitKey)
+
+		for key, value := range values {
+			if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil {
+				cidrGenerators.Insert(&cidrLimitEntry{
+					ipNet:     *ipNet,
+					generator: newLimitGenerator(value.in, value.out),
+				})
+				continue
+			}
+
+			if v, _ := l.inLimits.Get(key); v != nil {
+				lim := v.(traffic.Limiter)
+				if value.in <= 0 {
+					l.inLimits.Delete(key)
+				} else {
+					lim.Set(value.in)
+				}
+				delete(inLimits, key)
+			} else {
+				if value.in > 0 {
+					l.inLimits.Set(key, NewLimiter(value.in), cache.NoExpiration)
+				}
+			}
+
+			if v, _ := l.outLimits.Get(key); v != nil {
+				lim := v.(traffic.Limiter)
+				if value.out <= 0 {
+					l.outLimits.Delete(key)
+				} else {
+					lim.Set(value.out)
+				}
+				delete(outLimits, key)
+			} else {
+				if value.out > 0 {
+					l.outLimits.Set(key, NewLimiter(value.out), cache.NoExpiration)
+				}
+			}
+		}
+
+		// check the CIDR for remain limiters, clean the unmatched ones.
+		for k, v := range inLimits {
+			if p, _ := cidrGenerators.ContainingNetworks(net.ParseIP(k)); len(p) > 0 {
+				if le, _ := p[0].(*cidrLimitEntry); le != nil {
+					in := le.generator.in
+					if in <= 0 {
+						l.inLimits.Delete(k)
+						continue
+					}
+					lim := v.Object.(traffic.Limiter)
+					if lim.Limit() != in {
+						lim.Set(in)
+					}
+				}
+			} else {
+				l.inLimits.Delete(k)
+			}
+		}
+		for k, v := range outLimits {
+			if p, _ := cidrGenerators.ContainingNetworks(net.ParseIP(k)); len(p) > 0 {
+				if le, _ := p[0].(*cidrLimitEntry); le != nil {
+					out := le.generator.out
+					if out <= 0 {
+						l.outLimits.Delete(k)
+						continue
+					}
+					lim := v.Object.(traffic.Limiter)
+					if lim.Limit() != out {
+						lim.Set(out)
+					}
+					delete(outLimits, k)
+				}
+			} else {
+				l.outLimits.Delete(k)
+			}
+		}
+	}
+
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	l.cidrGenerators = cidrGenerators
+
+	return nil
+}
+
+func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue, err error) {
+	values = make(map[string]limitValue)
+
+	for _, v := range l.options.limits {
+		key, in, out := l.parseLimit(v)
+		if key == "" {
+			continue
+		}
+		values[key] = limitValue{in: in, out: out}
+	}
+
+	if l.options.fileLoader != nil {
+		if lister, ok := l.options.fileLoader.(loader.Lister); ok {
+			list, er := lister.List(ctx)
+			if er != nil {
+				l.options.logger.Warnf("file loader: %v", er)
+			}
+			for _, s := range list {
+				key, in, out := l.parseLimit(l.parseLine(s))
+				if key == "" {
+					continue
+				}
+				values[key] = limitValue{in: in, out: out}
+			}
+		} else {
+			r, er := l.options.fileLoader.Load(ctx)
+			if er != nil {
+				l.options.logger.Warnf("file loader: %v", er)
+			}
+			patterns, _ := l.parsePatterns(r)
+			for _, s := range patterns {
+				key, in, out := l.parseLimit(l.parseLine(s))
+				if key == "" {
+					continue
+				}
+				values[key] = limitValue{in: in, out: out}
+			}
+		}
+	}
+	if l.options.redisLoader != nil {
+		if lister, ok := l.options.redisLoader.(loader.Lister); ok {
+			list, er := lister.List(ctx)
+			if er != nil {
+				l.options.logger.Warnf("redis loader: %v", er)
+			}
+			for _, s := range list {
+				key, in, out := l.parseLimit(l.parseLine(s))
+				if key == "" {
+					continue
+				}
+				values[key] = limitValue{in: in, out: out}
+			}
+		} else {
+			r, er := l.options.redisLoader.Load(ctx)
+			if er != nil {
+				l.options.logger.Warnf("redis loader: %v", er)
+			}
+			patterns, _ := l.parsePatterns(r)
+			for _, s := range patterns {
+				key, in, out := l.parseLimit(l.parseLine(s))
+				if key == "" {
+					continue
+				}
+				values[key] = limitValue{in: in, out: out}
+			}
+		}
+	}
+	if l.options.httpLoader != nil {
+		r, er := l.options.httpLoader.Load(ctx)
+		if er != nil {
+			l.options.logger.Warnf("http loader: %v", er)
+		}
+		patterns, _ := l.parsePatterns(r)
+		for _, s := range patterns {
+			key, in, out := l.parseLimit(l.parseLine(s))
+			if key == "" {
+				continue
+			}
+			values[key] = limitValue{in: in, out: out}
+		}
+	}
+
+	l.options.logger.Debugf("load items %d", len(values))
+	return
+}
+
+func (l *trafficLimiter) parsePatterns(r io.Reader) (patterns []string, err error) {
+	if r == nil {
+		return
+	}
+
+	scanner := bufio.NewScanner(r)
+	for scanner.Scan() {
+		if line := l.parseLine(scanner.Text()); line != "" {
+			patterns = append(patterns, line)
+		}
+	}
+
+	err = scanner.Err()
+	return
+}
+
+func (l *trafficLimiter) parseLine(s string) string {
+	if n := strings.IndexByte(s, '#'); n >= 0 {
+		s = s[:n]
+	}
+	return strings.TrimSpace(s)
+}
+
+func (l *trafficLimiter) parseLimit(s string) (key string, in, out int) {
+	s = strings.Replace(s, "\t", " ", -1)
+	s = strings.TrimSpace(s)
+	if s == "" {
+		return
+	}
+
+	var ss []string
+	for _, v := range strings.Split(s, " ") {
+		if v != "" {
+			ss = append(ss, v)
+		}
+	}
+	if len(ss) < 2 {
+		return
+	}
+
+	key = ss[0]
+	if v, _ := units.ParseBase2Bytes(ss[1]); v > 0 {
+		in = int(v)
+	}
+	if len(ss) > 2 {
+		if v, _ := units.ParseBase2Bytes(ss[2]); v > 0 {
+			out = int(v)
+		}
+	}
+
+	return
+}
+
+func (l *trafficLimiter) Close() error {
+	l.cancelFunc()
+	if l.options.fileLoader != nil {
+		l.options.fileLoader.Close()
+	}
+	if l.options.redisLoader != nil {
+		l.options.redisLoader.Close()
+	}
+	return nil
+}
+
+type cidrLimitEntry struct {
+	ipNet     net.IPNet
+	generator *limitGenerator
+}
+
+func (p *cidrLimitEntry) Network() net.IPNet {
+	return p.ipNet
+}

go-gost/x/limiter/traffic/wrapper/conn.go

@@ -0,0 +1,420 @@
+package wrapper
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"io"
+	"net"
+	"syscall"
+
+	"github.com/go-gost/core/limiter"
+	"github.com/go-gost/core/limiter/traffic"
+	"github.com/go-gost/core/metadata"
+	xnet "github.com/go-gost/x/internal/net"
+	"github.com/go-gost/x/internal/net/udp"
+)
+
+var (
+	errUnsupport = errors.New("unsupported operation")
+)
+
+// limitConn is a Conn with traffic limiter supported.
+type limitConn struct {
+	net.Conn
+	rbuf    bytes.Buffer
+	limiter traffic.TrafficLimiter
+	opts    []limiter.Option
+	key     string
+}
+
+func WrapConn(c net.Conn, tlimiter traffic.TrafficLimiter, key string, opts ...limiter.Option) net.Conn {
+	if tlimiter == nil {
+		return c
+	}
+
+	return &limitConn{
+		Conn:    c,
+		limiter: tlimiter,
+		opts:    opts,
+		key:     key,
+	}
+}
+
+func (c *limitConn) Read(b []byte) (n int, err error) {
+	limiter := c.limiter.In(context.Background(), c.key, c.opts...)
+	if limiter == nil || limiter.Limit() <= 0 {
+		return c.Conn.Read(b)
+	}
+
+	if c.rbuf.Len() > 0 {
+		burst := len(b)
+		if c.rbuf.Len() < burst {
+			burst = c.rbuf.Len()
+		}
+		lim := limiter.Wait(context.Background(), burst)
+		return c.rbuf.Read(b[:lim])
+	}
+
+	nn, err := c.Conn.Read(b)
+	if err != nil {
+		return nn, err
+	}
+
+	n = limiter.Wait(context.Background(), nn)
+	if n < nn {
+		if _, err = c.rbuf.Write(b[n:nn]); err != nil {
+			return 0, err
+		}
+	}
+
+	return
+}
+
+func (c *limitConn) Write(b []byte) (n int, err error) {
+	limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
+	if limiter == nil || limiter.Limit() <= 0 {
+		return c.Conn.Write(b)
+	}
+
+	nn := 0
+	for len(b) > 0 {
+		nn, err = c.Conn.Write(b[:limiter.Wait(context.Background(), len(b))])
+		n += nn
+		if err != nil {
+			return
+		}
+		b = b[nn:]
+	}
+
+	return
+}
+
+func (c *limitConn) SyscallConn() (rc syscall.RawConn, err error) {
+	if sc, ok := c.Conn.(syscall.Conn); ok {
+		rc, err = sc.SyscallConn()
+		return
+	}
+	err = errUnsupport
+	return
+}
+
+func (c *limitConn) Metadata() metadata.Metadata {
+	if md, ok := c.Conn.(metadata.Metadatable); ok {
+		return md.Metadata()
+	}
+	return nil
+}
+
+type packetConn struct {
+	net.PacketConn
+	limiter traffic.TrafficLimiter
+	opts    []limiter.Option
+	key     string
+}
+
+func WrapPacketConn(pc net.PacketConn, lim traffic.TrafficLimiter, key string, opts ...limiter.Option) net.PacketConn {
+	if lim == nil {
+		return pc
+	}
+	return &packetConn{
+		PacketConn: pc,
+		limiter:    lim,
+		opts:       opts,
+		key:        key,
+	}
+}
+
+func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
+	for {
+		n, addr, err = c.PacketConn.ReadFrom(p)
+		if err != nil {
+			return
+		}
+
+		limiter := c.limiter.In(context.Background(), c.key, c.opts...)
+		if limiter == nil || limiter.Limit() <= 0 {
+			return
+		}
+
+		// discard when exceed the limit size.
+		if limiter.Wait(context.Background(), n) < n {
+			continue
+		}
+
+		return
+	}
+}
+
+func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
+	// discard when exceed the limit size.
+	limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
+	if limiter != nil && limiter.Limit() > 0 &&
+		limiter.Wait(context.Background(), len(p)) < len(p) {
+		n = len(p)
+		return
+	}
+
+	return c.PacketConn.WriteTo(p, addr)
+}
+
+func (c *packetConn) Metadata() metadata.Metadata {
+	if md, ok := c.PacketConn.(metadata.Metadatable); ok {
+		return md.Metadata()
+	}
+	return nil
+}
+
+type udpConn struct {
+	net.PacketConn
+	limiter traffic.TrafficLimiter
+	opts    []limiter.Option
+	key     string
+}
+
+func WrapUDPConn(pc net.PacketConn, limiter traffic.TrafficLimiter, key string, opts ...limiter.Option) udp.Conn {
+	return &udpConn{
+		PacketConn: pc,
+		limiter:    limiter,
+		opts:       opts,
+		key:        key,
+	}
+}
+
+func (c *udpConn) RemoteAddr() net.Addr {
+	if nc, ok := c.PacketConn.(xnet.RemoteAddr); ok {
+		return nc.RemoteAddr()
+	}
+	return nil
+}
+
+func (c *udpConn) SetReadBuffer(n int) error {
+	if nc, ok := c.PacketConn.(xnet.SetBuffer); ok {
+		return nc.SetReadBuffer(n)
+	}
+	return errUnsupport
+}
+
+func (c *udpConn) SetWriteBuffer(n int) error {
+	if nc, ok := c.PacketConn.(xnet.SetBuffer); ok {
+		return nc.SetWriteBuffer(n)
+	}
+	return errUnsupport
+}
+
+func (c *udpConn) Read(b []byte) (n int, err error) {
+	nc, ok := c.PacketConn.(io.Reader)
+	if !ok {
+		err = errUnsupport
+		return
+	}
+
+	for {
+		n, err = nc.Read(b)
+		if err != nil {
+			return
+		}
+
+		if c.limiter == nil {
+			return
+		}
+
+		limiter := c.limiter.In(context.Background(), c.key, c.opts...)
+		if limiter == nil || limiter.Limit() <= 0 {
+			return
+		}
+
+		// discard when exceed the limit size.
+		if limiter.Wait(context.Background(), n) < n {
+			continue
+		}
+
+		return
+	}
+}
+
+func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
+	for {
+		n, addr, err = c.PacketConn.ReadFrom(p)
+		if err != nil {
+			return
+		}
+
+		if c.limiter == nil {
+			return
+		}
+
+		limiter := c.limiter.In(context.Background(), c.key, c.opts...)
+		if limiter == nil || limiter.Limit() <= 0 {
+			return
+		}
+
+		// discard when exceed the limit size.
+		if limiter.Wait(context.Background(), n) < n {
+			continue
+		}
+
+		return
+	}
+}
+
+func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) {
+	nc, ok := c.PacketConn.(udp.ReadUDP)
+	if !ok {
+		err = errUnsupport
+		return
+	}
+
+	for {
+		n, addr, err = nc.ReadFromUDP(b)
+		if err != nil {
+			return
+		}
+
+		if c.limiter == nil {
+			return
+		}
+
+		limiter := c.limiter.In(context.Background(), c.key, c.opts...)
+		if limiter == nil || limiter.Limit() <= 0 {
+			return
+		}
+
+		// discard when exceed the limit size.
+		if limiter.Wait(context.Background(), n) < n {
+			continue
+		}
+
+		return
+	}
+}
+
+func (c *udpConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error) {
+	nc, ok := c.PacketConn.(udp.ReadUDP)
+	if !ok {
+		err = errUnsupport
+		return
+	}
+
+	for {
+		n, oobn, flags, addr, err = nc.ReadMsgUDP(b, oob)
+		if err != nil {
+			return
+		}
+
+		if c.limiter == nil {
+			return
+		}
+
+		limiter := c.limiter.In(context.Background(), c.key, c.opts...)
+		if limiter == nil || limiter.Limit() <= 0 {
+			return
+		}
+
+		// discard when exceed the limit size.
+		if limiter.Wait(context.Background(), n) < n {
+			continue
+		}
+		return
+	}
+}
+
+func (c *udpConn) Write(p []byte) (n int, err error) {
+	nc, ok := c.PacketConn.(io.Writer)
+	if !ok {
+		err = errUnsupport
+		return
+	}
+
+	if c.limiter != nil {
+		// discard when exceed the limit size.
+		limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
+		if limiter != nil && limiter.Limit() > 0 &&
+			limiter.Wait(context.Background(), len(p)) < len(p) {
+			n = len(p)
+			return
+		}
+	}
+
+	n, err = nc.Write(p)
+	return
+}
+
+func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
+	if c.limiter != nil {
+		// discard when exceed the limit size.
+		limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
+		if limiter != nil && limiter.Limit() > 0 &&
+			limiter.Wait(context.Background(), len(p)) < len(p) {
+			n = len(p)
+			return
+		}
+	}
+
+	n, err = c.PacketConn.WriteTo(p, addr)
+	return
+}
+
+func (c *udpConn) WriteToUDP(p []byte, addr *net.UDPAddr) (n int, err error) {
+	nc, ok := c.PacketConn.(udp.WriteUDP)
+	if !ok {
+		err = errUnsupport
+		return
+	}
+
+	if c.limiter != nil {
+		// discard when exceed the limit size.
+		limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
+		if limiter != nil && limiter.Limit() > 0 &&
+			limiter.Wait(context.Background(), len(p)) < len(p) {
+			n = len(p)
+			return
+		}
+	}
+
+	n, err = nc.WriteToUDP(p, addr)
+	return
+}
+
+func (c *udpConn) WriteMsgUDP(p, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) {
+	nc, ok := c.PacketConn.(udp.WriteUDP)
+	if !ok {
+		err = errUnsupport
+		return
+	}
+
+	if c.limiter != nil {
+		// discard when exceed the limit size.
+		limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
+		if limiter != nil && limiter.Limit() > 0 &&
+			limiter.Wait(context.Background(), len(p)) < len(p) {
+			n = len(p)
+			return
+		}
+	}
+
+	n, oobn, err = nc.WriteMsgUDP(p, oob, addr)
+	return
+}
+
+func (c *udpConn) SyscallConn() (rc syscall.RawConn, err error) {
+	if nc, ok := c.PacketConn.(xnet.SyscallConn); ok {
+		return nc.SyscallConn()
+	}
+	err = errUnsupport
+	return
+}
+
+func (c *udpConn) SetDSCP(n int) error {
+	if nc, ok := c.PacketConn.(xnet.SetDSCP); ok {
+		return nc.SetDSCP(n)
+	}
+	return nil
+}
+
+func (c *udpConn) Metadata() metadata.Metadata {
+	if md, ok := c.PacketConn.(metadata.Metadatable); ok {
+		return md.Metadata()
+	}
+	return nil
+}

go-gost/x/limiter/traffic/wrapper/io.go

@@ -0,0 +1,81 @@
+package wrapper
+
+import (
+	"bytes"
+	"context"
+	"io"
+
+	"github.com/go-gost/core/limiter"
+	"github.com/go-gost/core/limiter/traffic"
+)
+
+// readWriter is an io.ReadWriter with traffic limiter supported.
+type readWriter struct {
+	io.ReadWriter
+	rbuf    bytes.Buffer
+	limiter traffic.TrafficLimiter
+	opts    []limiter.Option
+	key     string
+}
+
+func WrapReadWriter(limiter traffic.TrafficLimiter, rw io.ReadWriter, key string, opts ...limiter.Option) io.ReadWriter {
+	if limiter == nil {
+		return rw
+	}
+
+	return &readWriter{
+		ReadWriter: rw,
+		limiter:    limiter,
+		opts:       opts,
+		key:        key,
+	}
+}
+
+func (p *readWriter) Read(b []byte) (n int, err error) {
+	limiter := p.limiter.In(context.Background(), p.key, p.opts...)
+	if limiter == nil || limiter.Limit() <= 0 {
+		return p.ReadWriter.Read(b)
+	}
+
+	if p.rbuf.Len() > 0 {
+		burst := len(b)
+		if p.rbuf.Len() < burst {
+			burst = p.rbuf.Len()
+		}
+		lim := limiter.Wait(context.Background(), burst)
+		return p.rbuf.Read(b[:lim])
+	}
+
+	nn, err := p.ReadWriter.Read(b)
+	if err != nil {
+		return nn, err
+	}
+
+	n = limiter.Wait(context.Background(), nn)
+	if n < nn {
+		if _, err = p.rbuf.Write(b[n:nn]); err != nil {
+			return 0, err
+		}
+	}
+
+	return
+}
+
+func (p *readWriter) Write(b []byte) (n int, err error) {
+	limiter := p.limiter.Out(context.Background(), p.key, p.opts...)
+	if limiter == nil || limiter.Limit() <= 0 {
+		return p.ReadWriter.Write(b)
+	}
+
+	nn := 0
+	for len(b) > 0 {
+		nn, err = p.ReadWriter.Write(b[:limiter.Wait(context.Background(), len(b))])
+		n += nn
+		if err != nil {
+			return
+		}
+		b = b[nn:]
+	}
+
+	return
+}

go-gost/x/limiter/traffic/wrapper/listener.go

@@ -0,0 +1,40 @@
+package wrapper
+
+import (
+	"net"
+
+	"github.com/go-gost/core/limiter"
+	"github.com/go-gost/core/limiter/traffic"
+	traffic_limiter "github.com/go-gost/x/limiter/traffic"
+)
+
+type listener struct {
+	net.Listener
+	limiter traffic.TrafficLimiter
+	service string
+}
+
+func WrapListener(service string, ln net.Listener, limiter traffic.TrafficLimiter) net.Listener {
+	if limiter == nil {
+		return ln
+	}
+
+	return &listener{
+		Listener: ln,
+		limiter:  limiter,
+		service:  service,
+	}
+}
+
+func (ln *listener) Accept() (net.Conn, error) {
+	c, err := ln.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+
+	return WrapConn(c, ln.limiter, traffic_limiter.ServiceLimitKey,
+		limiter.ScopeOption(limiter.ScopeService),
+		limiter.ServiceOption(ln.service),
+		limiter.NetworkOption(ln.Addr().Network()),
+	), nil
+}

go-gost/x/traffic/memory_manager.go

@@ -0,0 +1,91 @@
+package traffic
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+)
+
+// MemoryManager 内存流量管理器
+type MemoryManager struct {
+	mu    sync.RWMutex
+	stats map[string]*trafficStats
+}
+
+// trafficStats 流量统计
+type trafficStats struct {
+	upload   atomic.Int64
+	download atomic.Int64
+}
+
+// NewMemoryManager 创建内存流量管理器
+func NewMemoryManager() *MemoryManager {
+	return &MemoryManager{
+		stats: make(map[string]*trafficStats),
+	}
+}
+
+// RecordTraffic 记录流量
+func (m *MemoryManager) RecordTraffic(ctx context.Context, service string, upload, download int64) error {
+	m.mu.RLock()
+	stats, exists := m.stats[service]
+	m.mu.RUnlock()
+
+	if !exists {
+		m.mu.Lock()
+		stats, exists = m.stats[service]
+		if !exists {
+			stats = &trafficStats{}
+			m.stats[service] = stats
+		}
+		m.mu.Unlock()
+	}
+
+	if upload > 0 {
+		stats.upload.Add(upload)
+	}
+	if download > 0 {
+		stats.download.Add(download)
+	}
+
+	return nil
+}
+
+// GetAllServicesStats 获取所有服务的流量统计
+func (m *MemoryManager) GetAllServicesStats(ctx context.Context) (map[string]map[string]int64, error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
+	result := make(map[string]map[string]int64)
+	for service, stats := range m.stats {
+		result[service] = map[string]int64{
+			"upload":   stats.upload.Load(),
+			"download": stats.download.Load(),
+		}
+	}
+
+	return result, nil
+}
+
+// ClearAllTrafficStats 清零所有流量统计
+func (m *MemoryManager) ClearAllTrafficStats(ctx context.Context) error {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
+	for _, stats := range m.stats {
+		stats.upload.Store(0)
+		stats.download.Store(0)
+	}
+
+	return nil
+}
+
+// Close 关闭管理器（内存管理器无需特殊清理）
+func (m *MemoryManager) Close() error {
+	return nil
+}
+
+// TestConnection 测试连接（内存管理器总是返回成功）
+func (m *MemoryManager) TestConnection(ctx context.Context) error {
+	return nil
+}

go-gost/x/traffic/traffic.go

@@ -0,0 +1,33 @@
+package traffic
+
+import (
+	"context"
+	"sync"
+)
+
+// Manager 流量管理器接口
+type Manager interface {
+	RecordTraffic(ctx context.Context, service string, upload, download int64) error
+	GetAllServicesStats(ctx context.Context) (map[string]map[string]int64, error)
+	ClearAllTrafficStats(ctx context.Context) error
+	Close() error
+	TestConnection(ctx context.Context) error
+}
+
+var (
+	globalManager Manager
+	once          sync.Once
+)
+
+// GetGlobalManager 获取全局流量管理器实例
+func GetGlobalManager() Manager {
+	once.Do(func() {
+		globalManager = NewMemoryManager()
+	})
+	return globalManager
+}
+
+// SetGlobalManager 设置全局流量管理器（用于测试）
+func SetGlobalManager(m Manager) {
+	globalManager = m
+}