network proto: revert full buffer size read (#4497)

* network proto: revert full buffer size read

* fix read-all in network protocol
This commit is contained in:
Tarun Koyalwar 2023-12-17 19:32:19 +07:00 committed by GitHub
parent 1f04f253fa
commit e912c7c58d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -6,6 +6,7 @@ import (
"fmt" "fmt"
"net" "net"
"net/url" "net/url"
"os"
"strings" "strings"
"time" "time"
@ -265,7 +266,7 @@ func (request *Request) executeRequestWithPayloads(variables map[string]interfac
} }
if input.Read > 0 { if input.Read > 0 {
buffer, err := reader.ConnReadNWithTimeout(conn, int64(input.Read), DefaultReadTimeout) buffer, err := ConnReadNWithTimeout(conn, int64(input.Read), DefaultReadTimeout)
if err != nil { if err != nil {
return errorutil.NewWithErr(err).Msgf("could not read response from connection") 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 bufferSize = -1
} }
final, err := reader.ConnReadNWithTimeout(conn, int64(bufferSize), DefaultReadTimeout) final, err := ConnReadNWithTimeout(conn, int64(bufferSize), DefaultReadTimeout)
if err != nil { if err != nil {
request.options.Output.Request(request.options.TemplatePath, address, request.Type().String(), err) request.options.Output.Request(request.options.TemplatePath, address, request.Type().String(), err)
return errors.Wrap(err, "could not read from server") return errors.Wrap(err, "could not read from server")
@ -412,3 +413,27 @@ func getAddress(toTest string) (string, error) {
} }
return toTest, nil 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
}