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)
+}