diff --git a/v2/pkg/protocols/network/network.go b/v2/pkg/protocols/network/network.go index a5274eff6..12094ae6e 100644 --- a/v2/pkg/protocols/network/network.go +++ b/v2/pkg/protocols/network/network.go @@ -60,6 +60,13 @@ type Request struct { // examples: // - value: "2048" ReadSize int `yaml:"read-size,omitempty" jsonschema:"title=size of network response to read,description=Size of response to read at the end. Default is 1024 bytes"` + // description: | + // ReadAll determines if the data stream should be read till the end regardless of the size + // + // Default value for read-all is false. + // examples: + // - value: false + ReadAll bool `yaml:"read-all,omitempty" jsonschema:"title=read all response stream,description=Read all response stream till the server stops sending"` // description: | // SelfContained specifies if the request is self contained. diff --git a/v2/pkg/protocols/network/request.go b/v2/pkg/protocols/network/request.go index acbbac9af..66da31e9b 100644 --- a/v2/pkg/protocols/network/request.go +++ b/v2/pkg/protocols/network/request.go @@ -6,6 +6,7 @@ import ( "io" "net" "net/url" + "os" "strings" "time" @@ -200,13 +201,48 @@ func (request *Request) executeRequestWithPayloads(actualAddress, address, input if request.ReadSize != 0 { bufferSize = request.ReadSize } - final := make([]byte, bufferSize) - n, err := conn.Read(final) - if err != nil && err != io.EOF { - request.options.Output.Request(request.options.TemplateID, address, "network", err) - return errors.Wrap(err, "could not read from server") + + var ( + final []byte + n int + ) + + if request.ReadAll { + readInterval := time.NewTimer(time.Second * 1) + // stop the timer and drain the channel + closeTimer := func(t *time.Timer) { + if !t.Stop() { + <-t.C + } + } + read_socket: + for { + select { + case <-readInterval.C: + closeTimer(readInterval) + break read_socket + default: + buf := make([]byte, bufferSize) + nBuf, err := conn.Read(buf) + if err != nil && !os.IsTimeout(err) { + request.options.Output.Request(request.options.TemplateID, address, "network", err) + closeTimer(readInterval) + return errors.Wrap(err, "could not read from server") + } + responseBuilder.Write(buf[:nBuf]) + final = append(final, buf...) + n += nBuf + } + } + } else { + final = make([]byte, bufferSize) + n, err = conn.Read(final) + if err != nil && err != io.EOF { + request.options.Output.Request(request.options.TemplateID, address, "network", err) + return errors.Wrap(err, "could not read from server") + } + responseBuilder.Write(final[:n]) } - responseBuilder.Write(final[:n]) response := responseBuilder.String() outputEvent := request.responseToDSLMap(reqBuilder.String(), string(final[:n]), response, input, actualAddress)