add rate limiter

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

48
registry/limiter.go Normal file
View File

@ -0,0 +1,48 @@
package registry
import (
"github.com/go-gost/core/limiter"
)
type rlimiterRegistry struct {
registry
}
func (r *rlimiterRegistry) Register(name string, v limiter.RateLimiter) error {
return r.registry.Register(name, v)
}
func (r *rlimiterRegistry) Get(name string) limiter.RateLimiter {
if name != "" {
return &rlimiterWrapper{name: name, r: r}
}
return nil
}
func (r *rlimiterRegistry) get(name string) limiter.RateLimiter {
if v := r.registry.Get(name); v != nil {
return v.(limiter.RateLimiter)
}
return nil
}
type rlimiterWrapper struct {
name string
r *rlimiterRegistry
}
func (w *rlimiterWrapper) Input() limiter.Limiter {
v := w.r.get(w.name)
if v == nil {
return nil
}
return v.Input()
}
func (w *rlimiterWrapper) Output() limiter.Limiter {
v := w.r.get(w.name)
if v == nil {
return nil
}
return v.Output()
}

View File

@ -10,6 +10,7 @@ import (
"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/recorder"
"github.com/go-gost/core/resolver"
"github.com/go-gost/core/service"
@ -33,6 +34,7 @@ var (
resolverReg Registry[resolver.Resolver] = &resolverRegistry{}
hostsReg Registry[hosts.HostMapper] = &hostsRegistry{}
recorderReg Registry[recorder.Recorder] = &recorderRegistry{}
rlimiterReg Registry[limiter.RateLimiter] = &rlimiterRegistry{}
)
type Registry[T any] interface {
@ -126,3 +128,7 @@ func HostsRegistry() Registry[hosts.HostMapper] {
func RecorderRegistry() Registry[recorder.Recorder] {
return recorderReg
}
func RateLimiterRegistry() Registry[limiter.RateLimiter] {
return rlimiterReg
}