add traffic limiter option for handler

This commit is contained in:
ginuerzh 2023-11-18 18:25:40 +08:00
parent a916f04016
commit 486f2cee61
2 changed files with 44 additions and 3 deletions

View File

@ -8,6 +8,7 @@ import (
"github.com/go-gost/core/bypass" "github.com/go-gost/core/bypass"
"github.com/go-gost/core/chain" "github.com/go-gost/core/chain"
"github.com/go-gost/core/limiter/rate" "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/logger"
"github.com/go-gost/core/metadata" "github.com/go-gost/core/metadata"
) )
@ -18,6 +19,7 @@ type Options struct {
Auth *url.Userinfo Auth *url.Userinfo
Auther auth.Authenticator Auther auth.Authenticator
RateLimiter rate.RateLimiter RateLimiter rate.RateLimiter
Limiter traffic.TrafficLimiter
TLSConfig *tls.Config TLSConfig *tls.Config
Logger logger.Logger Logger logger.Logger
Service string 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 { func TLSConfigOption(tlsConfig *tls.Config) Option {
return func(opts *Options) { return func(opts *Options) {
opts.TLSConfig = tlsConfig opts.TLSConfig = tlsConfig

View File

@ -10,7 +10,40 @@ type Limiter interface {
Set(n int) Set(n int)
} }
type TrafficLimiter interface { type Options struct {
In(key string) Limiter Network string
Out(key string) Limiter 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
} }