add rate limiter
This commit is contained in:
48
registry/limiter.go
Normal file
48
registry/limiter.go
Normal 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()
|
||||
}
|
@ -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
|
||||
}
|
||||
|
Reference in New Issue
Block a user