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-04-21 21:37:30 +08:00
parent a04c8b45f3
commit e23da0f319
8 changed files with 218 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
config/parsing/parse.go
View File

@ -4,11 +4,13 @@ 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"
"github.com/go-gost/core/chain"
"github.com/go-gost/core/hosts"
"github.com/go-gost/core/limiter"
"github.com/go-gost/core/logger"
"github.com/go-gost/core/recorder"
"github.com/go-gost/core/resolver"
@ -17,9 +19,10 @@ import (
auth_impl "github.com/go-gost/x/auth"
bypass_impl "github.com/go-gost/x/bypass"
"github.com/go-gost/x/config"
hosts_impl "github.com/go-gost/x/hosts"
xhosts "github.com/go-gost/x/hosts"
"github.com/go-gost/x/internal/loader"
recorder_impl "github.com/go-gost/x/recorder"
xlimiter "github.com/go-gost/x/limiter"
xrecorder "github.com/go-gost/x/recorder"
"github.com/go-gost/x/registry"
resolver_impl "github.com/go-gost/x/resolver"
xs "github.com/go-gost/x/selector"
@ -224,7 +227,7 @@ func ParseHosts(cfg *config.HostsConfig) hosts.HostMapper {
return nil
}
var mappings []hosts_impl.Mapping
var mappings []xhosts.Mapping
for _, mapping := range cfg.Mappings {
if mapping.IP == "" || mapping.Hostname == "" {
continue
@ -234,33 +237,33 @@ func ParseHosts(cfg *config.HostsConfig) hosts.HostMapper {
if ip == nil {
continue
}
mappings = append(mappings, hosts_impl.Mapping{
mappings = append(mappings, xhosts.Mapping{
Hostname: mapping.Hostname,
IP: ip,
})
}
opts := []hosts_impl.Option{
hosts_impl.MappingsOption(mappings),
hosts_impl.ReloadPeriodOption(cfg.Reload),
hosts_impl.LoggerOption(logger.Default().WithFields(map[string]any{
opts := []xhosts.Option{
xhosts.MappingsOption(mappings),
xhosts.ReloadPeriodOption(cfg.Reload),
xhosts.LoggerOption(logger.Default().WithFields(map[string]any{
"kind": "hosts",
"hosts": cfg.Name,
})),
}
if cfg.File != nil && cfg.File.Path != "" {
opts = append(opts, hosts_impl.FileLoaderOption(loader.FileLoader(cfg.File.Path)))
opts = append(opts, xhosts.FileLoaderOption(loader.FileLoader(cfg.File.Path)))
}
if cfg.Redis != nil && cfg.Redis.Addr != "" {
switch cfg.Redis.Type {
case "list": // redis list
opts = append(opts, hosts_impl.RedisLoaderOption(loader.RedisListLoader(
opts = append(opts, xhosts.RedisLoaderOption(loader.RedisListLoader(
cfg.Redis.Addr,
loader.DBRedisLoaderOption(cfg.Redis.DB),
loader.PasswordRedisLoaderOption(cfg.Redis.Password),
loader.KeyRedisLoaderOption(cfg.Redis.Key),
)))
default: // redis set
opts = append(opts, hosts_impl.RedisLoaderOption(loader.RedisSetLoader(
opts = append(opts, xhosts.RedisLoaderOption(loader.RedisSetLoader(
cfg.Redis.Addr,
loader.DBRedisLoaderOption(cfg.Redis.DB),
loader.PasswordRedisLoaderOption(cfg.Redis.Password),
@ -268,7 +271,7 @@ func ParseHosts(cfg *config.HostsConfig) hosts.HostMapper {
)))
}
}
return hosts_impl.NewHostMapper(opts...)
return xhosts.NewHostMapper(opts...)
}
func ParseRecorder(cfg *config.RecorderConfig) (r recorder.Recorder) {
@ -277,8 +280,8 @@ func ParseRecorder(cfg *config.RecorderConfig) (r recorder.Recorder) {
}
if cfg.File != nil && cfg.File.Path != "" {
return recorder_impl.FileRecorder(cfg.File.Path,
recorder_impl.SepRecorderOption(cfg.File.Sep))
return xrecorder.FileRecorder(cfg.File.Path,
xrecorder.SepRecorderOption(cfg.File.Sep))
}
if cfg.Redis != nil &&
@ -286,16 +289,16 @@ func ParseRecorder(cfg *config.RecorderConfig) (r recorder.Recorder) {
cfg.Redis.Key != "" {
switch cfg.Redis.Type {
case "list": // redis list
return recorder_impl.RedisListRecorder(cfg.Redis.Addr,
recorder_impl.DBRedisRecorderOption(cfg.Redis.DB),
recorder_impl.KeyRedisRecorderOption(cfg.Redis.Key),
recorder_impl.PasswordRedisRecorderOption(cfg.Redis.Password),
return xrecorder.RedisListRecorder(cfg.Redis.Addr,
xrecorder.DBRedisRecorderOption(cfg.Redis.DB),
xrecorder.KeyRedisRecorderOption(cfg.Redis.Key),
xrecorder.PasswordRedisRecorderOption(cfg.Redis.Password),
)
default: // redis set
return recorder_impl.RedisSetRecorder(cfg.Redis.Addr,
recorder_impl.DBRedisRecorderOption(cfg.Redis.DB),
recorder_impl.KeyRedisRecorderOption(cfg.Redis.Key),
recorder_impl.PasswordRedisRecorderOption(cfg.Redis.Password),
return xrecorder.RedisSetRecorder(cfg.Redis.Addr,
xrecorder.DBRedisRecorderOption(cfg.Redis.DB),
xrecorder.KeyRedisRecorderOption(cfg.Redis.Key),
xrecorder.PasswordRedisRecorderOption(cfg.Redis.Password),
)
}
}
@ -318,3 +321,35 @@ func defaultChainSelector() selector.Selector[chain.Chainer] {
xs.BackupFilter[chain.Chainer](),
)
}
func ParseRateLimiter(cfg *config.LimiterConfig) (lim limiter.RateLimiter) {
if cfg == nil || cfg.RateLimit == 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 input, output limiter.Limiter
if len(rlimiters) > 0 {
input = xlimiter.MultiLimiter(rlimiters...)
}
if len(wlimiters) > 0 {
output = xlimiter.MultiLimiter(wlimiters...)
}
return xlimiter.RateLimiter(input, output)
}

1
config/parsing/service.go
View File

@ -75,6 +75,7 @@ func ParseService(cfg *config.ServiceConfig) (service.Service, error) {
listener.TLSConfigOption(tlsConfig),
listener.AdmissionOption(admission.AdmissionGroup(admissions...)),
listener.ChainOption(chainGroup(cfg.Listener.Chain, cfg.Listener.ChainGroup)),
listener.RateLimiterOption(registry.RateLimiterRegistry().Get(cfg.Limiter)),
listener.LoggerOption(listenerLogger),
listener.ServiceOption(cfg.Name),
)