From e912c7c58d1e9bca09633d11714e4264a32a1b26 Mon Sep 17 00:00:00 2001
From: Tarun Koyalwar <45962551+tarunKoyalwar@users.noreply.github.com>
Date: Sun, 17 Dec 2023 19:32:19 +0700
Subject: [PATCH] network proto: revert full buffer size read (#4497)

* network proto: revert full buffer size read

* fix read-all in network protocol
---
 pkg/protocols/network/request.go | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/pkg/protocols/network/request.go b/pkg/protocols/network/request.go
index e6652476..c896b655 100644
--- a/pkg/protocols/network/request.go
+++ b/pkg/protocols/network/request.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net"
 	"net/url"
+	"os"
 	"strings"
 	"time"
 
@@ -265,7 +266,7 @@ func (request *Request) executeRequestWithPayloads(variables map[string]interfac
 		}
 
 		if input.Read > 0 {
-			buffer, err := reader.ConnReadNWithTimeout(conn, int64(input.Read), DefaultReadTimeout)
+			buffer, err := ConnReadNWithTimeout(conn, int64(input.Read), DefaultReadTimeout)
 			if err != nil {
 				return errorutil.NewWithErr(err).Msgf("could not read response from connection")
 			}
@@ -315,7 +316,7 @@ func (request *Request) executeRequestWithPayloads(variables map[string]interfac
 		bufferSize = -1
 	}
 
-	final, err := reader.ConnReadNWithTimeout(conn, int64(bufferSize), DefaultReadTimeout)
+	final, err := ConnReadNWithTimeout(conn, int64(bufferSize), DefaultReadTimeout)
 	if err != nil {
 		request.options.Output.Request(request.options.TemplatePath, address, request.Type().String(), err)
 		return errors.Wrap(err, "could not read from server")
@@ -412,3 +413,27 @@ func getAddress(toTest string) (string, error) {
 	}
 	return toTest, nil
 }
+
+func ConnReadNWithTimeout(conn net.Conn, n int64, timeout time.Duration) ([]byte, error) {
+	if timeout == 0 {
+		timeout = DefaultReadTimeout
+	}
+	if n == -1 {
+		// if n is -1 then read all available data from connection
+		return reader.ConnReadNWithTimeout(conn, -1, timeout)
+	} else if n == 0 {
+		n = 4096 // default buffer size
+	}
+	b := make([]byte, n)
+	_ = conn.SetDeadline(time.Now().Add(timeout))
+	count, err := conn.Read(b)
+	_ = conn.SetDeadline(time.Time{})
+	if err != nil && os.IsTimeout(err) && count > 0 {
+		// in case of timeout with some value read, return the value
+		return b[:count], nil
+	}
+	if err != nil {
+		return nil, err
+	}
+	return b[:count], nil
+}