setting user agent and better disconnection limit

lavanet · Dec 4, 2024 · ead2816 · ead2816
1 parent 0b5efac
commit ead2816
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/protocol/chainlib/websocket_connection_limiter.go b/protocol/chainlib/websocket_connection_limiter.go
@@ -20,10 +20,9 @@ type WebsocketConnectionLimiter struct {
 }
 
 func (wcl *WebsocketConnectionLimiter) handleFiberRateLimitFlags(c *fiber.Ctx) {
-	userAgent := c.Get("User-Agent")
-	utils.LavaFormatDebug("User-Agent", utils.LogAttr("userAgent", userAgent))
+	userAgent := c.Get(fiber.HeaderUserAgent)
 	// Store the User-Agent in locals for later use
-	c.Locals("User-Agent", userAgent)
+	c.Locals(fiber.HeaderUserAgent, userAgent)
 
 	forwardedFor := c.Get(common.IP_FORWARDING_HEADER_NAME)
 	if forwardedFor == "" {
@@ -79,7 +78,7 @@ func (wcl *WebsocketConnectionLimiter) canOpenConnection(websocketConn *websocke
 		numberOfActiveConnections := wcl.addIpConnectionAndGetCurrentAmount(key)
 
 		if numberOfActiveConnections > connectionLimit {
-			websocketConn.WriteMessage(1, []byte(fmt.Sprintf("Too Many Open Connections, limited to %d", connectionLimit)))
+			websocketConn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, fmt.Sprintf("Too Many Open Connections, limited to %d", connectionLimit)))
 			return false, func() { wcl.decreaseIpConnection(key) }
 		}
 	}