go-gost/x
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add rate limiter

Browse Source
This commit is contained in:
ginuerzh
2022-09-05 22:47:51 +08:00
parent e23da0f319
commit 4c2131ca6d
40 changed files with 1622 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
config/config.go
View File

@ -185,19 +185,20 @@ type RecorderObject struct {
}
type LimiterConfig struct {
Name string `json:"name"`
RateLimit *RateLimitConfig `yaml:"rate" json:"rate"`
Name string `json:"name"`
Rate *RateLimiterConfig `yaml:"rate" json:"rate"`
}
type RateLimitConfig struct {
Input string `yaml:",omitempty" json:"input,omitempty"`
Output string `yaml:",omitempty" json:"output,omitempty"`
Conn *LimitConfig `yaml:",omitempty" json:"conn,omitempty"`
type RateLimiterConfig struct {
Limits []string `yaml:",omitempty" json:"limits,omitempty"`
Reload time.Duration `yaml:",omitempty" json:"reload,omitempty"`
File *FileLoader `yaml:",omitempty" json:"file,omitempty"`
Redis *RedisLoader `yaml:",omitempty" json:"redis,omitempty"`
}
type LimitConfig struct {
Input string `yaml:",omitempty" json:"input,omitempty"`
Output string `yaml:",omitempty" json:"output,omitempty"`
In string `yaml:",omitempty" json:"in,omitempty"`
Out string `yaml:",omitempty" json:"out,omitempty"`
}
type ListenerConfig struct {

55
config/parsing/parse.go
View File

@ -4,7 +4,6 @@ import (
"net"
"net/url"
"github.com/alecthomas/units"
"github.com/go-gost/core/admission"
"github.com/go-gost/core/auth"
"github.com/go-gost/core/bypass"
@ -323,33 +322,41 @@ func defaultChainSelector() selector.Selector[chain.Chainer] {
}
func ParseRateLimiter(cfg *config.LimiterConfig) (lim limiter.RateLimiter) {
if cfg == nil || cfg.RateLimit == nil {
if cfg == nil || cfg.Rate == nil {
return nil
}
var rlimiters []limiter.Limiter
var wlimiters []limiter.Limiter
if cfg.RateLimit.Conn != nil {
if v, _ := units.ParseBase2Bytes(cfg.RateLimit.Conn.Input); v > 0 {
rlimiters = append(rlimiters, xlimiter.Limiter(int(v)))
}
if v, _ := units.ParseBase2Bytes(cfg.RateLimit.Conn.Output); v > 0 {
wlimiters = append(wlimiters, xlimiter.Limiter(int(v)))
}
}
if v, _ := units.ParseBase2Bytes(cfg.RateLimit.Input); v > 0 {
rlimiters = append(rlimiters, xlimiter.Limiter(int(v)))
}
if v, _ := units.ParseBase2Bytes(cfg.RateLimit.Output); v > 0 {
wlimiters = append(wlimiters, xlimiter.Limiter(int(v)))
}
var opts []xlimiter.Option
var input, output limiter.Limiter
if len(rlimiters) > 0 {
input = xlimiter.MultiLimiter(rlimiters...)
if cfg.Rate.File != nil && cfg.Rate.File.Path != "" {
opts = append(opts, xlimiter.FileLoaderOption(loader.FileLoader(cfg.Rate.File.Path)))
}
if len(wlimiters) > 0 {
output = xlimiter.MultiLimiter(wlimiters...)
if cfg.Rate.Redis != nil && cfg.Rate.Redis.Addr != "" {
switch cfg.Rate.Redis.Type {
case "list": // redis list
opts = append(opts, xlimiter.RedisLoaderOption(loader.RedisListLoader(
cfg.Rate.Redis.Addr,
loader.DBRedisLoaderOption(cfg.Rate.Redis.DB),
loader.PasswordRedisLoaderOption(cfg.Rate.Redis.Password),
loader.KeyRedisLoaderOption(cfg.Rate.Redis.Key),
)))
default: // redis set
opts = append(opts, xlimiter.RedisLoaderOption(loader.RedisSetLoader(
cfg.Rate.Redis.Addr,
loader.DBRedisLoaderOption(cfg.Rate.Redis.DB),
loader.PasswordRedisLoaderOption(cfg.Rate.Redis.Password),
loader.KeyRedisLoaderOption(cfg.Rate.Redis.Key),
)))
}
}
return xlimiter.RateLimiter(input, output)
opts = append(opts,
xlimiter.LimitsOption(cfg.Rate.Limits...),
xlimiter.ReloadPeriodOption(cfg.Rate.Reload),
xlimiter.LoggerOption(logger.Default().WithFields(map[string]any{
"kind": "limiter",
"hosts": cfg.Name,
})),
)
return xlimiter.NewRateLimiter(opts...)
}