From f8d5d025702da2e9a477b0f04b2602cd2053abc2 Mon Sep 17 00:00:00 2001 From: mzack Date: Tue, 12 Oct 2021 23:28:24 +0200 Subject: [PATCH 1/3] Adding support to read tcp data stream till the end --- v2/pkg/protocols/network/network.go | 7 +++++ v2/pkg/protocols/network/request.go | 40 ++++++++++++++++++++++++----- 2 files changed, 41 insertions(+), 6 deletions(-) diff --git a/v2/pkg/protocols/network/network.go b/v2/pkg/protocols/network/network.go index b7de7b920..ef47dc36e 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"` // Operators for the current request go here. operators.Operators `yaml:",inline,omitempty"` diff --git a/v2/pkg/protocols/network/request.go b/v2/pkg/protocols/network/request.go index aad3c60ef..ab5176685 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" @@ -183,13 +184,40 @@ func (r *Request) executeRequestWithPayloads(actualAddress, address, input strin if r.ReadSize != 0 { bufferSize = r.ReadSize } - final := make([]byte, bufferSize) - n, err := conn.Read(final) - if err != nil && err != io.EOF { - r.options.Output.Request(r.options.TemplateID, address, "network", err) - return errors.Wrap(err, "could not read from server") + + var ( + final []byte + n int + ) + + if r.ReadAll { + readInterval := time.After(time.Second * 1) + read_socket: + for { + select { + case <-readInterval: + break read_socket + default: + buf := make([]byte, bufferSize) + nBuf, err := conn.Read(buf) + if err != nil && !os.IsTimeout(err) { + r.options.Output.Request(r.options.TemplateID, address, "network", err) + 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 { + r.options.Output.Request(r.options.TemplateID, address, "network", err) + return errors.Wrap(err, "could not read from server") + } + responseBuilder.Write(final[:n]) } - responseBuilder.Write(final[:n]) if r.options.Options.Debug || r.options.Options.DebugResponse { responseOutput := responseBuilder.String() From 0ba3b19f1f7653a41327315644b7fcccb57de1c2 Mon Sep 17 00:00:00 2001 From: mzack Date: Tue, 19 Oct 2021 11:31:36 +0200 Subject: [PATCH 2/3] fixing merge issues --- v2/pkg/protocols/network/request.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/v2/pkg/protocols/network/request.go b/v2/pkg/protocols/network/request.go index 399699a0f..a1425ebae 100644 --- a/v2/pkg/protocols/network/request.go +++ b/v2/pkg/protocols/network/request.go @@ -191,13 +191,13 @@ func (request *Request) executeRequestWithPayloads(actualAddress, address, input if request.ReadSize != 0 { bufferSize = request.ReadSize } - - var ( + + var ( final []byte n int ) - if r.ReadAll { + if request.ReadAll { readInterval := time.After(time.Second * 1) read_socket: for { @@ -208,7 +208,7 @@ func (request *Request) executeRequestWithPayloads(actualAddress, address, input buf := make([]byte, bufferSize) nBuf, err := conn.Read(buf) if err != nil && !os.IsTimeout(err) { - r.options.Output.Request(r.options.TemplateID, address, "network", err) + request.options.Output.Request(request.options.TemplateID, address, "network", err) return errors.Wrap(err, "could not read from server") } responseBuilder.Write(buf[:nBuf]) From ae99a88f52b4ddb260818a8ba80a28f650233894 Mon Sep 17 00:00:00 2001 From: mzack Date: Tue, 19 Oct 2021 19:39:33 +0200 Subject: [PATCH 3/3] timer.after => timer --- v2/pkg/protocols/network/request.go | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/v2/pkg/protocols/network/request.go b/v2/pkg/protocols/network/request.go index a1425ebae..6beecc6e1 100644 --- a/v2/pkg/protocols/network/request.go +++ b/v2/pkg/protocols/network/request.go @@ -198,17 +198,25 @@ func (request *Request) executeRequestWithPayloads(actualAddress, address, input ) if request.ReadAll { - readInterval := time.After(time.Second * 1) + 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: + 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])