1 files changed, 6 insertions, 1 deletions

diff --git a/cmd/relay/main.go b/cmd/relay/main.go
index e4afec8..86a29cb 100644
--- a/cmd/relay/main.go
+++ b/cmd/relay/main.go
@@ -79,8 +79,9 @@ func main() {
                 streamInterceptors = append(streamInterceptors, auth.NostrStreamInterceptor(authOpts))
         }
 
+        var limiter *ratelimit.Limiter
         if cfg.RateLimit.Enabled {
-                limiter := ratelimit.New(cfg.RateLimit.ToRateLimiter())
+                limiter = ratelimit.New(cfg.RateLimit.ToRateLimiter())
                 unaryInterceptors = append(unaryInterceptors, ratelimit.UnaryInterceptor(limiter))
                 streamInterceptors = append(streamInterceptors, ratelimit.StreamInterceptor(limiter))
         }
@@ -125,6 +126,10 @@ func main() {
                 })
         }
 
+        if limiter != nil {
+                wsHandler.SetRateLimiter(limiter)
+        }
+
         var grpcDisplay, httpDisplay, wsDisplay string
         if cfg.Server.PublicURL != "" {
                 grpcDisplay = cfg.Server.PublicURL + ":443"