Fix counters when reading or writing data to the socket

Author: mostafa

Diff for: 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 for: network/proxy.go

@@ -727,19 +727,20 @@ 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)
 
 			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
 		}