diff --git a/protocol/chainlib/consumer_websocket_manager_test.go b/protocol/chainlib/consumer_websocket_manager_test.go index 645a3a55db..c501a663c9 100644 --- a/protocol/chainlib/consumer_websocket_manager_test.go +++ b/protocol/chainlib/consumer_websocket_manager_test.go @@ -105,7 +105,7 @@ func TestWebsocketConnectionLimiter(t *testing.T) { // Test the connection for _, expectSuccess := range tt.expectSuccess { - canOpen, _ := wcl.canOpenConnection(mockWsConn) + canOpen, _ := wcl.CanOpenConnection(mockWsConn) if expectSuccess { assert.True(t, canOpen, "Expected connection to be allowed") } else { diff --git a/protocol/chainlib/jsonRPC.go b/protocol/chainlib/jsonRPC.go index 9abff54d02..de48263375 100644 --- a/protocol/chainlib/jsonRPC.go +++ b/protocol/chainlib/jsonRPC.go @@ -360,7 +360,7 @@ func (apil *JsonRPCChainListener) Serve(ctx context.Context, cmdFlags common.Con app := createAndSetupBaseAppListener(cmdFlags, apil.endpoint.HealthCheckPath, apil.healthReporter) app.Use("/ws", func(c *fiber.Ctx) error { - apil.websocketConnectionLimiter.handleFiberRateLimitFlags(c) + apil.websocketConnectionLimiter.HandleFiberRateLimitFlags(c) // IsWebSocketUpgrade returns true if the client // requested upgrade to the WebSocket protocol. @@ -375,7 +375,7 @@ func (apil *JsonRPCChainListener) Serve(ctx context.Context, cmdFlags common.Con apiInterface := apil.endpoint.ApiInterface webSocketCallback := websocket.New(func(websocketConn *websocket.Conn) { - canOpenConnection, decreaseIpConnection := apil.websocketConnectionLimiter.canOpenConnection(websocketConn) + canOpenConnection, decreaseIpConnection := apil.websocketConnectionLimiter.CanOpenConnection(websocketConn) defer decreaseIpConnection() if !canOpenConnection { return diff --git a/protocol/chainlib/tendermintRPC.go b/protocol/chainlib/tendermintRPC.go index 3cf5503e6d..3bb8867095 100644 --- a/protocol/chainlib/tendermintRPC.go +++ b/protocol/chainlib/tendermintRPC.go @@ -390,7 +390,7 @@ func (apil *TendermintRpcChainListener) Serve(ctx context.Context, cmdFlags comm apiInterface := apil.endpoint.ApiInterface app.Use("/ws", func(c *fiber.Ctx) error { - apil.websocketConnectionLimiter.handleFiberRateLimitFlags(c) + apil.websocketConnectionLimiter.HandleFiberRateLimitFlags(c) // IsWebSocketUpgrade returns true if the client // requested upgrade to the WebSocket protocol. if websocket.IsWebSocketUpgrade(c) { @@ -400,7 +400,7 @@ func (apil *TendermintRpcChainListener) Serve(ctx context.Context, cmdFlags comm return fiber.ErrUpgradeRequired }) webSocketCallback := websocket.New(func(websocketConn *websocket.Conn) { - canOpenConnection, decreaseIpConnection := apil.websocketConnectionLimiter.canOpenConnection(websocketConn) + canOpenConnection, decreaseIpConnection := apil.websocketConnectionLimiter.CanOpenConnection(websocketConn) defer decreaseIpConnection() if !canOpenConnection { return diff --git a/protocol/chainlib/websocket_connection_limiter.go b/protocol/chainlib/websocket_connection_limiter.go index aee2ba293c..9be6a27e73 100644 --- a/protocol/chainlib/websocket_connection_limiter.go +++ b/protocol/chainlib/websocket_connection_limiter.go @@ -27,7 +27,7 @@ type WebsocketConnectionLimiter struct { lock sync.RWMutex } -func (wcl *WebsocketConnectionLimiter) handleFiberRateLimitFlags(c *fiber.Ctx) { +func (wcl *WebsocketConnectionLimiter) HandleFiberRateLimitFlags(c *fiber.Ctx) { userAgent := c.Get(fiber.HeaderUserAgent) // Store the User-Agent in locals for later use c.Locals(fiber.HeaderUserAgent, userAgent) @@ -68,7 +68,7 @@ func (wcl *WebsocketConnectionLimiter) getConnectionLimit(websocketConn Websocke return utils.Max(MaximumNumberOfParallelWebsocketConnectionsPerIp, connectionLimitHeaderValue) } -func (wcl *WebsocketConnectionLimiter) canOpenConnection(websocketConn WebsocketConnection) (bool, func()) { +func (wcl *WebsocketConnectionLimiter) CanOpenConnection(websocketConn WebsocketConnection) (bool, func()) { // Check which connection limit is higher and use that. connectionLimit := wcl.getConnectionLimit(websocketConn) decreaseIpConnectionCallback := func() {} @@ -93,7 +93,7 @@ func (wcl *WebsocketConnectionLimiter) canOpenConnection(websocketConn Websocket return false, decreaseIpConnectionCallback } // If under limit, increment and return cleanup function - wcl.addIpConnectionAndGetCurrentAmount(key) + wcl.addIpConnection(key) decreaseIpConnectionCallback = func() { wcl.decreaseIpConnection(key) } } return true, decreaseIpConnectionCallback @@ -105,7 +105,7 @@ func (wcl *WebsocketConnectionLimiter) getCurrentAmountOfConnections(key string) return wcl.ipToNumberOfActiveConnections[key] } -func (wcl *WebsocketConnectionLimiter) addIpConnectionAndGetCurrentAmount(key string) { +func (wcl *WebsocketConnectionLimiter) addIpConnection(key string) { wcl.lock.Lock() defer wcl.lock.Unlock() // wether it exists or not we add 1.