diff --git a/handler/option.go b/handler/option.go index b1faf53..8176081 100644 --- a/handler/option.go +++ b/handler/option.go @@ -8,6 +8,7 @@ import ( "github.com/go-gost/core/bypass" "github.com/go-gost/core/chain" "github.com/go-gost/core/limiter/rate" + "github.com/go-gost/core/limiter/traffic" "github.com/go-gost/core/logger" "github.com/go-gost/core/metadata" ) @@ -18,6 +19,7 @@ type Options struct { Auth *url.Userinfo Auther auth.Authenticator RateLimiter rate.RateLimiter + Limiter traffic.TrafficLimiter TLSConfig *tls.Config Logger logger.Logger Service string @@ -55,6 +57,12 @@ func RateLimiterOption(limiter rate.RateLimiter) Option { } } +func TrafficLimiterOption(limiter traffic.TrafficLimiter) Option { + return func(opts *Options) { + opts.Limiter = limiter + } +} + func TLSConfigOption(tlsConfig *tls.Config) Option { return func(opts *Options) { opts.TLSConfig = tlsConfig diff --git a/limiter/traffic/limiter.go b/limiter/traffic/limiter.go index 523fab9..0f902fc 100644 --- a/limiter/traffic/limiter.go +++ b/limiter/traffic/limiter.go @@ -10,7 +10,40 @@ type Limiter interface { Set(n int) } -type TrafficLimiter interface { - In(key string) Limiter - Out(key string) Limiter +type Options struct { + Network string + Addr string + Client string + Src string +} + +type Option func(opts *Options) + +func NetworkOption(network string) Option { + return func(opts *Options) { + opts.Network = network + } +} + +func AddrOption(addr string) Option { + return func(opts *Options) { + opts.Addr = addr + } +} + +func ClientOption(client string) Option { + return func(opts *Options) { + opts.Client = client + } +} + +func SrcOption(src string) Option { + return func(opts *Options) { + opts.Src = src + } +} + +type TrafficLimiter interface { + In(ctx context.Context, key string, opts ...Option) Limiter + Out(ctx context.Context, key string, opts ...Option) Limiter }