From 37fec6114bac10074810c8f9ee1de74dce6c9b22 Mon Sep 17 00:00:00 2001
From: Mostafa Moradian <mostafa@gatewayd.io>
Date: Sat, 12 Oct 2024 14:41:59 +0200
Subject: [PATCH] Fix counters when reading or writing data to the socket

---
 network/client.go | 11 +++++++----
 network/proxy.go  |  9 +++++----
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/network/client.go b/network/client.go
index bc6ff5a2..2086bccf 100644
--- a/network/client.go
+++ b/network/client.go
@@ -187,9 +187,10 @@ func (c *Client) Send(data []byte) (int, *gerr.GatewayDError) {
 	}
 
 	sent := 0
-	received := len(data)
+	dataSize := len(data)
 	for {
-		if sent >= received {
+		// If we've sent all the data, we must break the loop.
+		if sent >= dataSize {
 			break
 		}
 
@@ -240,13 +241,15 @@ func (c *Client) Receive() (int, []byte, *gerr.GatewayDError) {
 	for ctx.Err() == nil {
 		chunk := make([]byte, c.ReceiveChunkSize)
 		read, err := c.conn.Read(chunk)
+		if read > 0 {
+			total += read
+			buffer.Write(chunk[:read])
+		}
 		if err != nil {
 			c.logger.Error().Err(err).Msg("Couldn't receive data from the server")
 			span.RecordError(err)
 			return total, buffer.Bytes(), gerr.ErrClientReceiveFailed.Wrap(err)
 		}
-		total += read
-		buffer.Write(chunk[:read])
 
 		if read < c.ReceiveChunkSize {
 			break
diff --git a/network/proxy.go b/network/proxy.go
index 6d60b310..f67d3b45 100644
--- a/network/proxy.go
+++ b/network/proxy.go
@@ -727,6 +727,10 @@ func (pr *Proxy) receiveTrafficFromClient(conn net.Conn) ([]byte, *gerr.GatewayD
 	for {
 		chunk := make([]byte, pr.ClientConfig.ReceiveChunkSize)
 		read, err := conn.Read(chunk)
+		if read > 0 {
+			total += read
+			buffer.Write(chunk[:read])
+		}
 		if read == 0 || err != nil {
 			pr.Logger.Debug().Err(err).Msg("Error reading from client")
 			span.RecordError(err)
@@ -734,12 +738,9 @@ func (pr *Proxy) receiveTrafficFromClient(conn net.Conn) ([]byte, *gerr.GatewayD
 			metrics.BytesReceivedFromClient.WithLabelValues(pr.GetGroupName(), pr.GetBlockName()).Observe(float64(read))
 			metrics.TotalTrafficBytes.WithLabelValues(pr.GetGroupName(), pr.GetBlockName()).Observe(float64(read))
 
-			return chunk[:read], gerr.ErrReadFailed.Wrap(err)
+			return buffer.Bytes(), gerr.ErrReadFailed.Wrap(err)
 		}
 
-		total += read
-		buffer.Write(chunk[:read])
-
 		if read < pr.ClientConfig.ReceiveChunkSize {
 			break
 		}