Adding prototype of request flow override annotations (#2161)

* Adding prototype of request flow override annotations

* reworking iteration engine

* adding directory to .gitignore
dev
Mzack9999 2022-07-18 10:46:03 +02:00 committed by GitHub
parent 4ae458df98
commit cf1039f49c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 86 additions and 19 deletions

1
.gitignore vendored
View File

@ -11,3 +11,4 @@ v2/cmd/docgen/docgen
v2/pkg/protocols/common/helpers/deserialization/testdata/Deserialize.class v2/pkg/protocols/common/helpers/deserialization/testdata/Deserialize.class
v2/pkg/protocols/common/helpers/deserialization/testdata/ValueObject.class v2/pkg/protocols/common/helpers/deserialization/testdata/ValueObject.class
v2/pkg/protocols/common/helpers/deserialization/testdata/ValueObject2.ser v2/pkg/protocols/common/helpers/deserialization/testdata/ValueObject2.ser
.vscode

View File

@ -24,8 +24,30 @@ var (
reSniAnnotation = regexp.MustCompile(`(?m)^@tls-sni:\s*(.+)\s*$`) reSniAnnotation = regexp.MustCompile(`(?m)^@tls-sni:\s*(.+)\s*$`)
// @timeout:duration overrides the input timout with a custom duration // @timeout:duration overrides the input timout with a custom duration
reTimeoutAnnotation = regexp.MustCompile(`(?m)^@timeout:\s*(.+)\s*$`) reTimeoutAnnotation = regexp.MustCompile(`(?m)^@timeout:\s*(.+)\s*$`)
// @once sets the request to be executed only once for a specific URL
reOnceAnnotation = regexp.MustCompile(`(?m)^@once\s*$`)
) )
type flowMark int
const (
Once flowMark = iota
)
// parseFlowAnnotations and override requests flow
func parseFlowAnnotations(rawRequest string) (flowMark, bool) {
var fm flowMark
// parse request for known ovverride annotations
var hasFlowOveride bool
// @once
if reOnceAnnotation.MatchString(rawRequest) {
fm = Once
hasFlowOveride = true
}
return fm, hasFlowOveride
}
// parseAnnotations and override requests settings // parseAnnotations and override requests settings
func (r *Request) parseAnnotations(rawRequest string, request *http.Request) (*http.Request, bool) { func (r *Request) parseAnnotations(rawRequest string, request *http.Request) (*http.Request, bool) {
// parse request for known ovverride annotations // parse request for known ovverride annotations

View File

@ -19,6 +19,7 @@ type requestGenerator struct {
options *protocols.ExecuterOptions options *protocols.ExecuterOptions
payloadIterator *generators.Iterator payloadIterator *generators.Iterator
interactshURLs []string interactshURLs []string
onceFlow map[string]struct{}
} }
// LeaveDefaultPorts skips normalization of default standard ports // LeaveDefaultPorts skips normalization of default standard ports
@ -26,7 +27,11 @@ var LeaveDefaultPorts = false
// newGenerator creates a new request generator instance // newGenerator creates a new request generator instance
func (request *Request) newGenerator() *requestGenerator { func (request *Request) newGenerator() *requestGenerator {
generator := &requestGenerator{request: request, options: request.options} generator := &requestGenerator{
request: request,
options: request.options,
onceFlow: make(map[string]struct{}),
}
if len(request.Payloads) > 0 { if len(request.Payloads) > 0 {
generator.payloadIterator = request.generator.NewIterator() generator.payloadIterator = request.generator.NewIterator()
@ -53,29 +58,68 @@ func (r *requestGenerator) nextValue() (value string, payloads map[string]interf
} }
hasPayloadIterator := r.payloadIterator != nil hasPayloadIterator := r.payloadIterator != nil
hasInitializedPayloads := r.currentPayloads != nil
if r.currentIndex == 0 && hasPayloadIterator && !hasInitializedPayloads { if hasPayloadIterator && r.currentPayloads == nil {
r.currentPayloads, r.okCurrentPayload = r.payloadIterator.Value() r.currentPayloads, r.okCurrentPayload = r.payloadIterator.Value()
} }
if r.currentIndex < len(sequence) {
currentRequest := sequence[r.currentIndex] var request string
r.currentIndex++ var shouldContinue bool
return currentRequest, r.currentPayloads, true if nextRequest, nextIndex, found := r.findNextIteration(sequence, r.currentIndex); found {
r.currentIndex = nextIndex + 1
request = nextRequest
shouldContinue = true
} else if nextRequest, nextIndex, found := r.findNextIteration(sequence, 0); found && hasPayloadIterator {
r.currentIndex = nextIndex + 1
request = nextRequest
shouldContinue = true
} }
if r.currentIndex == len(sequence) {
if r.okCurrentPayload { if shouldContinue {
r.currentIndex = 0 if r.hasMarker(request, Once) {
currentRequest := sequence[r.currentIndex] r.applyMark(request, Once)
if hasPayloadIterator { }
r.currentPayloads, r.okCurrentPayload = r.payloadIterator.Value() if hasPayloadIterator {
if r.okCurrentPayload { return request, r.currentPayloads, r.okCurrentPayload
r.currentIndex++ }
return currentRequest, r.currentPayloads, true return request, r.currentPayloads, true
} } else {
} return "", nil, false
}
}
func (r *requestGenerator) findNextIteration(sequence []string, index int) (string, int, bool) {
for i, request := range sequence[index:] {
if !r.wasMarked(request, Once) {
return request, index + i, true
} }
} }
return "", nil, false if r.payloadIterator != nil {
r.currentPayloads, r.okCurrentPayload = r.payloadIterator.Value()
}
return "", 0, false
}
func (r *requestGenerator) applyMark(request string, mark flowMark) {
switch mark {
case Once:
r.onceFlow[request] = struct{}{}
}
}
func (r *requestGenerator) wasMarked(request string, mark flowMark) bool {
switch mark {
case Once:
_, ok := r.onceFlow[request]
return ok
}
return false
}
func (r *requestGenerator) hasMarker(request string, mark flowMark) bool {
fo, hasOverrides := parseFlowAnnotations(request)
return hasOverrides && fo == mark
} }