From 57b85f4f8dae611d346f6ca6059d5d08de49d8a4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= <zoltan.pmail@gmail.com>
Date: Tue, 9 Jul 2024 16:15:25 +0200
Subject: [PATCH] Optimisation for sonar

---
 relay/client/client.go | 101 ++++++++++++++++++++++++-----------------
 1 file changed, 60 insertions(+), 41 deletions(-)

diff --git a/relay/client/client.go b/relay/client/client.go
index 9c8bd0365..35c338e34 100644
--- a/relay/client/client.go
+++ b/relay/client/client.go
@@ -297,7 +297,7 @@ func (c *Client) readLoop(relayConn net.Conn) {
 				c.log.Debugf("failed to read message from relay server: %s", errExit)
 			}
 			c.mu.Unlock()
-			goto Exit
+			break
 		}
 
 		msgType, err := messages.DetermineServerMsgType(buf[:n])
@@ -306,55 +306,74 @@ func (c *Client) readLoop(relayConn net.Conn) {
 			continue
 		}
 
-		switch msgType {
-		case messages.MsgTypeHealthCheck:
-			msg := messages.MarshalHealthcheck()
-			_, wErr := c.relayConn.Write(msg)
-			if wErr != nil {
-				if c.serviceIsRunning && !internallyStoppedFlag.isSet() {
-					c.log.Errorf("failed to send heartbeat: %s", wErr)
-				}
-			}
-			hc.Heartbeat()
-		case messages.MsgTypeTransport:
-			peerID, payload, err := messages.UnmarshalTransportMsg(buf[:n])
-			if err != nil {
-				if c.serviceIsRunning && !internallyStoppedFlag.isSet() {
-					c.log.Errorf("failed to parse transport message: %v", err)
-				}
-				continue
-			}
-			stringID := messages.HashIDToString(peerID)
-
-			c.mu.Lock()
-			if !c.serviceIsRunning {
-				c.mu.Unlock()
-				goto Exit
-			}
-			container, ok := c.conns[stringID]
-			c.mu.Unlock()
-			if !ok {
-				c.log.Errorf("peer not found: %s", stringID)
-				continue
-			}
-
-			container.writeMsg(Msg{
-				bufPool: c.bufPool,
-				bufPtr:  bufPtr,
-				Payload: payload})
-		case messages.MsgTypeClose:
-			log.Debugf("relay connection close by server")
-			goto Exit
+		if !c.handleMsg(msgType, buf[:n], bufPtr, hc, internallyStoppedFlag) {
+			break
 		}
 	}
 
-Exit:
 	hc.Stop()
 	c.notifyDisconnected()
 	c.wgReadLoop.Done()
 	_ = c.close(false)
 }
 
+func (c *Client) handleMsg(msgType messages.MsgType, buf []byte, bufPtr *[]byte, hc *healthcheck.Receiver, internallyStoppedFlag *internalStopFlag) (continueLoop bool) {
+	switch msgType {
+	case messages.MsgTypeHealthCheck:
+		msg := messages.MarshalHealthcheck()
+		_, wErr := c.relayConn.Write(msg)
+		if wErr != nil {
+			if c.serviceIsRunning && !internallyStoppedFlag.isSet() {
+				c.log.Errorf("failed to send heartbeat: %s", wErr)
+			}
+		}
+		hc.Heartbeat()
+		c.bufPool.Put(bufPtr)
+	case messages.MsgTypeTransport:
+		return c.handleTrasnportMsg(buf, bufPtr, internallyStoppedFlag)
+	case messages.MsgTypeClose:
+		log.Debugf("relay connection close by server")
+		c.bufPool.Put(bufPtr)
+		return false
+	}
+
+	return true
+}
+
+func (c *Client) handleTrasnportMsg(buf []byte, bufPtr *[]byte, internallyStoppedFlag *internalStopFlag) bool {
+	peerID, payload, err := messages.UnmarshalTransportMsg(buf)
+	if err != nil {
+		if c.serviceIsRunning && !internallyStoppedFlag.isSet() {
+			c.log.Errorf("failed to parse transport message: %v", err)
+		}
+
+		c.bufPool.Put(bufPtr)
+		return true
+	}
+	stringID := messages.HashIDToString(peerID)
+
+	c.mu.Lock()
+	if !c.serviceIsRunning {
+		c.mu.Unlock()
+		c.bufPool.Put(bufPtr)
+		return false
+	}
+	container, ok := c.conns[stringID]
+	c.mu.Unlock()
+	if !ok {
+		c.log.Errorf("peer not found: %s", stringID)
+		c.bufPool.Put(bufPtr)
+		return true
+	}
+	msg := Msg{
+		bufPool: c.bufPool,
+		bufPtr:  bufPtr,
+		Payload: payload,
+	}
+	container.writeMsg(msg)
+	return true
+}
+
 // todo check by reference too, the id is not enought because the id come from the outer conn
 func (c *Client) writeTo(id string, dstID []byte, payload []byte) (int, error) {
 	c.mu.Lock()