Add client_connection_rate_limit option (#601)

Author: FZambia

go.mod

@@ -30,6 +30,7 @@ require (
 	go.uber.org/automaxprocs v1.5.1
 	golang.org/x/crypto v0.4.0
 	golang.org/x/sync v0.1.0
+	golang.org/x/time v0.3.0
 	google.golang.org/grpc v1.51.0
 	google.golang.org/protobuf v1.28.1
 )

go.sum

@@ -326,7 +326,8 @@ golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.5.0 h1:OLmvp0KP+FVG99Ct/qFiL/Fhk4zp4QQnZ7b2U+5piUM=
 golang.org/x/text v0.5.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
 golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
-golang.org/x/time v0.0.0-20211116232009-f0f3c7e86c11 h1:GZokNIeuVkl3aZHJchRrr13WCsols02MLUcz1U9is6M=
+golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4=
+golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180221164845-07fd8470d635/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=

internal/middleware/connlimit.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/centrifugal/centrifuge"
 	"github.com/prometheus/client_golang/prometheus"
+	"golang.org/x/time/rate"
 )
 
 var (
@@ -29,7 +30,12 @@ func init() {
 }
 
 func ConnLimit(node *centrifuge.Node, ruleContainer *rule.Container, h http.Handler) http.Handler {
+	rl := connectionRateLimiter(ruleContainer.Config().ClientConnectionRateLimit)
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if rl != nil && !rl.Allow() {
+			w.WriteHeader(http.StatusServiceUnavailable)
+			return
+		}
 		connLimit := ruleContainer.Config().ClientConnectionLimit
 		if connLimit > 0 && node.Hub().NumClients() >= connLimit {
 			connLimitReached.Inc()
@@ -45,3 +51,10 @@ func ConnLimit(node *centrifuge.Node, ruleContainer *rule.Container, h http.Hand
 		h.ServeHTTP(w, r)
 	})
 }
+
+func connectionRateLimiter(connRateLimit int) *rate.Limiter {
+	if connRateLimit > 0 {
+		return rate.NewLimiter(rate.Every(time.Second), connRateLimit)
+	}
+	return nil
+}

internal/middleware/connlimit_test.go

@@ -0,0 +1,36 @@
+package middleware
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/centrifugal/centrifugo/v4/internal/rule"
+	"github.com/centrifugal/centrifugo/v4/internal/tools"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestConnLimit_ConnectionRate(t *testing.T) {
+	node := tools.NodeWithMemoryEngineNoHandlers()
+	defer func() { _ = node.Shutdown(context.Background()) }()
+
+	ruleConfig, err := rule.NewContainer(rule.Config{
+		ClientConnectionRateLimit: 10,
+	})
+	require.NoError(t, err)
+
+	ts := httptest.NewServer(ConnLimit(node, ruleConfig, testHandler()))
+	defer ts.Close()
+
+	for i := 0; i < 20; i++ {
+		res, err := http.Post(ts.URL, "application/json", nil)
+		require.NoError(t, err)
+		if res.StatusCode == http.StatusServiceUnavailable {
+			require.True(t, i >= 10)
+			return
+		}
+	}
+	require.Fail(t, "no rate limit hit upon sending 10 requests to a server")
+}

internal/rule/rule.go

@@ -72,6 +72,9 @@ type Config struct {
 	// ClientConnectionLimit sets the maximum number of concurrent clients a single Centrifugo
 	// node will accept.
 	ClientConnectionLimit int
+	// ClientConnectionRateLimit sets the maximum number of new connections a single Centrifugo
+	// node will accept per second.
+	ClientConnectionRateLimit int
 }
 
 // DefaultConfig has default config options.

main.go

@@ -1413,6 +1413,7 @@ func ruleConfig() rule.Config {
 	cfg.RpcProxyName = v.GetString("rpc_proxy_name")
 	cfg.RpcNamespaces = rpcNamespacesFromConfig(v)
 	cfg.ClientConnectionLimit = v.GetInt("client_connection_limit")
+	cfg.ClientConnectionRateLimit = v.GetInt("client_connection_rate_limit")
 	return cfg
 }