feat: PRT - add ban duration for websocket rate limit exceeded

protocol/chainlib/consumer_websocket_manager.go

@@ -17,7 +17,10 @@ import (
 	"github.com/tidwall/gjson"
 )
 
-var WebSocketRateLimit = -1 // rate limit requests per second on websocket connection
+var (
+	WebSocketRateLimit   = -1               // rate limit requests per second on websocket connection
+	WebSocketBanDuration = time.Duration(0) // once rate limit is reached, will not allow new incoming message for a duration
+)
 
 type ConsumerWebsocketManager struct {
 	websocketConn                 *websocket.Conn
@@ -148,6 +151,15 @@ func (cwm *ConsumerWebsocketManager) ListenToMessages() {
 			case <-webSocketCtx.Done():
 				return
 			case <-ticker.C:
+				// check if rate limit reached, and ban is required
+				if WebSocketBanDuration > 0 && requestsPerSecond.Load() > uint64(WebSocketRateLimit) {
+					// wait the ban duration before resetting the store.
+					select {
+					case <-webSocketCtx.Done():
+						return
+					case <-time.After(WebSocketBanDuration): // just continue
+					}
+				}
 				requestsPerSecond.Store(0)
 			}
 		}

protocol/common/cobra_common.go

@@ -40,7 +40,8 @@ const (
 	SetProviderOptimizerNumberOfTiersToCreate = "set-provider-optimizer-number-of-tiers-to-create"
 
 	// websocket flags
-	RateLimitWebSocketFlag = "rate-limit-websocket-requests-per-connection"
+	RateLimitWebSocketFlag                       = "rate-limit-websocket-requests-per-connection"
+	BanDurationForWebsocketRateLimitExceededFlag = "ban-duration-for-websocket-rate-limit-exceeded"
 )
 
 const (

protocol/rpcconsumer/rpcconsumer.go

@@ -619,6 +619,7 @@ rpcconsumer consumer_examples/full_consumer_example.yml --cache-be "127.0.0.1:77
 	cmdRPCConsumer.Flags().Float64Var(&provideroptimizer.LastTierChance, common.SetProviderOptimizerWorstTierPickChance, 0.0, "set the chances for picking a provider from the worse group, default is 0% -> 0.0")
 	cmdRPCConsumer.Flags().IntVar(&provideroptimizer.OptimizerNumTiers, common.SetProviderOptimizerNumberOfTiersToCreate, 4, "set the number of groups to create, default is 4")
 	cmdRPCConsumer.Flags().IntVar(&chainlib.WebSocketRateLimit, common.RateLimitWebSocketFlag, chainlib.WebSocketRateLimit, "rate limit (per second) websocket requests per user connection, default is unlimited")
+	cmdRPCConsumer.Flags().DurationVar(&chainlib.WebSocketBanDuration, common.BanDurationForWebsocketRateLimitExceededFlag, chainlib.WebSocketBanDuration, "once websocket rate limit is reached, user will be banned Xfor a duration, default no ban")
 	common.AddRollingLogConfig(cmdRPCConsumer)
 	return cmdRPCConsumer
 }