diff --git a/go-gost/x/config/parsing/service/parse.go b/go-gost/x/config/parsing/service/parse.go index 3dbcb2b..0b7957c 100644 --- a/go-gost/x/config/parsing/service/parse.go +++ b/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), diff --git a/go-gost/x/limiter/traffic/traffic.go b/go-gost/x/limiter/traffic/traffic.go index 6f9c2b2..785fc86 100644 --- a/go-gost/x/limiter/traffic/traffic.go +++ b/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: