Merge pull request #56 from Sagit-chu/opencode/calm-orchid

fix(limiter): fix traffic limiter ScopeClient behavior to allow per-u…

go-gost/x/config/parsing/service/parse.go

@@ -3,6 +3,7 @@ package service
 import (
 	"fmt"
 	"runtime"
+	"strconv"
 	"strings"
 	"time"
 
@@ -30,6 +31,7 @@ import (
 	logger_parser "github.com/go-gost/x/config/parsing/logger"
 	selector_parser "github.com/go-gost/x/config/parsing/selector"
 	tls_util "github.com/go-gost/x/internal/util/tls"
+	xtraffic "github.com/go-gost/x/limiter/traffic"
 	cache_limiter "github.com/go-gost/x/limiter/traffic/cache"
 	"github.com/go-gost/x/metadata"
 	mdutil "github.com/go-gost/x/metadata/util"
@@ -181,6 +183,27 @@ func ParseService(cfg *config.ServiceConfig) (service.Service, error) {
 		)
 	}
 
+	var trafficLimiter listener.Option
+	if cfg.Limiter != "" {
+		lim := registry.TrafficLimiterRegistry().Get(cfg.Limiter)
+		if lim == nil {
+			// Try to parse as simple number (bandwidth in bytes/sec)
+			if val, err := strconv.Atoi(cfg.Limiter); err == nil && val > 0 {
+				lim = xtraffic.NewTrafficLimiter(
+					xtraffic.LimitsOption(fmt.Sprintf("%s %d %d", xtraffic.ServiceLimitKey, val, val)),
+				)
+			}
+		}
+		trafficLimiter = listener.TrafficLimiterOption(
+			cache_limiter.NewCachedTrafficLimiter(
+				lim,
+				cache_limiter.RefreshIntervalOption(limiterRefreshInterval),
+				cache_limiter.CleanupIntervalOption(limiterCleanupInterval),
+				cache_limiter.ScopeOption(limiterScope),
+			),
+		)
+	}
+
 	listenOpts := []listener.Option{
 		listener.AddrOption(cfg.Addr),
 		listener.RouterOption(xchain.NewRouter(routerOpts...)),
@@ -188,14 +211,7 @@ func ParseService(cfg *config.ServiceConfig) (service.Service, error) {
 		listener.AuthOption(auth_parser.Info(cfg.Listener.Auth)),
 		listener.TLSConfigOption(tlsConfig),
 		listener.AdmissionOption(xadmission.AdmissionGroup(admissions...)),
-		listener.TrafficLimiterOption(
-			cache_limiter.NewCachedTrafficLimiter(
-				registry.TrafficLimiterRegistry().Get(cfg.Limiter),
-				cache_limiter.RefreshIntervalOption(limiterRefreshInterval),
-				cache_limiter.CleanupIntervalOption(limiterCleanupInterval),
-				cache_limiter.ScopeOption(limiterScope),
-			),
-		),
+		trafficLimiter,
 		listener.ConnLimiterOption(registry.ConnLimiterRegistry().Get(cfg.CLimiter)),
 		listener.ServiceOption(cfg.Name),
 		listener.ProxyProtocolOption(ppv),

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

@@ -136,6 +136,9 @@ func (l *trafficLimiter) In(ctx context.Context, key string, opts ...limiter.Opt
 		return nil
 
 	case limiter.ScopeClient:
+		if lim, ok := l.inLimits.Get(key); ok && lim != nil {
+			return lim.(traffic.Limiter)
+		}
 		return nil
 
 	case limiter.ScopeConn:
@@ -215,6 +218,9 @@ func (l *trafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Op
 		return nil
 
 	case limiter.ScopeClient:
+		if lim, ok := l.outLimits.Get(key); ok && lim != nil {
+			return lim.(traffic.Limiter)
+		}
 		return nil
 
 	case limiter.ScopeConn: