update limiter

registry/limiter.go

@@ -2,6 +2,7 @@ package registry
 
 import (
 	"github.com/go-gost/core/limiter/conn"
+	"github.com/go-gost/core/limiter/rate"
 	"github.com/go-gost/core/limiter/traffic"
 )
 
@@ -82,3 +83,38 @@ func (w *connLimiterWrapper) Limiter(key string) conn.Limiter {
 	}
 	return v.Limiter(key)
 }
+
+type rateLimiterRegistry struct {
+	registry
+}
+
+func (r *rateLimiterRegistry) Register(name string, v rate.RateLimiter) error {
+	return r.registry.Register(name, v)
+}
+
+func (r *rateLimiterRegistry) Get(name string) rate.RateLimiter {
+	if name != "" {
+		return &rateLimiterWrapper{name: name, r: r}
+	}
+	return nil
+}
+
+func (r *rateLimiterRegistry) get(name string) rate.RateLimiter {
+	if v := r.registry.Get(name); v != nil {
+		return v.(rate.RateLimiter)
+	}
+	return nil
+}
+
+type rateLimiterWrapper struct {
+	name string
+	r    *rateLimiterRegistry
+}
+
+func (w *rateLimiterWrapper) Limiter(key string) rate.Limiter {
+	v := w.r.get(w.name)
+	if v == nil {
+		return nil
+	}
+	return v.Limiter(key)
+}

registry/registry.go

@@ -11,6 +11,7 @@ import (
 	"github.com/go-gost/core/chain"
 	"github.com/go-gost/core/hosts"
 	"github.com/go-gost/core/limiter/conn"
+	"github.com/go-gost/core/limiter/rate"
 	"github.com/go-gost/core/limiter/traffic"
 	"github.com/go-gost/core/recorder"
 	"github.com/go-gost/core/resolver"
@@ -27,16 +28,18 @@ var (
 	dialerReg    Registry[NewDialer]    = &dialerRegistry{}
 	connectorReg Registry[NewConnector] = &connectorRegistry{}
 
-	serviceReg        Registry[service.Service]        = &serviceRegistry{}
-	chainReg          Registry[chain.Chainer]          = &chainRegistry{}
-	autherReg         Registry[auth.Authenticator]     = &autherRegistry{}
-	admissionReg      Registry[admission.Admission]    = &admissionRegistry{}
-	bypassReg         Registry[bypass.Bypass]          = &bypassRegistry{}
-	resolverReg       Registry[resolver.Resolver]      = &resolverRegistry{}
-	hostsReg          Registry[hosts.HostMapper]       = &hostsRegistry{}
-	recorderReg       Registry[recorder.Recorder]      = &recorderRegistry{}
+	serviceReg   Registry[service.Service]     = &serviceRegistry{}
+	chainReg     Registry[chain.Chainer]       = &chainRegistry{}
+	autherReg    Registry[auth.Authenticator]  = &autherRegistry{}
+	admissionReg Registry[admission.Admission] = &admissionRegistry{}
+	bypassReg    Registry[bypass.Bypass]       = &bypassRegistry{}
+	resolverReg  Registry[resolver.Resolver]   = &resolverRegistry{}
+	hostsReg     Registry[hosts.HostMapper]    = &hostsRegistry{}
+	recorderReg  Registry[recorder.Recorder]   = &recorderRegistry{}
+
 	trafficLimiterReg Registry[traffic.TrafficLimiter] = &trafficLimiterRegistry{}
 	connLimiterReg    Registry[conn.ConnLimiter]       = &connLimiterRegistry{}
+	rateLimiterReg    Registry[rate.RateLimiter]       = &rateLimiterRegistry{}
 )
 
 type Registry[T any] interface {
@@ -138,3 +141,7 @@ func TrafficLimiterRegistry() Registry[traffic.TrafficLimiter] {
 func ConnLimiterRegistry() Registry[conn.ConnLimiter] {
 	return connLimiterReg
 }
+
+func RateLimiterRegistry() Registry[rate.RateLimiter] {
+	return rateLimiterReg
+}