From 4ec06b231da9bde8993aa37404a7057a2719135f Mon Sep 17 00:00:00 2001
From: Tarun Koyalwar <tarun@projectdiscovery.io>
Date: Fri, 26 Jan 2024 18:08:37 +0530
Subject: [PATCH] remove use of iterate() in flow

---
 cmd/integration-test/flow.go                  | 11 ++++++
 .../flow/iterate-one-value-flow.yaml          | 34 +++++++++++++++++++
 pkg/tmplexec/flow/flow_executor.go            | 12 +++++++
 pkg/tmplexec/flow/flow_internal.go            |  9 ++---
 pkg/tmplexec/flow/util.go                     | 17 ++++++++++
 5 files changed, 79 insertions(+), 4 deletions(-)
 create mode 100644 integration_tests/flow/iterate-one-value-flow.yaml

diff --git a/cmd/integration-test/flow.go b/cmd/integration-test/flow.go
index 334f7756fe..cfcb504489 100644
--- a/cmd/integration-test/flow.go
+++ b/cmd/integration-test/flow.go
@@ -14,6 +14,7 @@ var flowTestcases = []TestCaseInfo{
 	{Path: "flow/conditional-flow.yaml", TestCase: &conditionalFlow{}},
 	{Path: "flow/conditional-flow-negative.yaml", TestCase: &conditionalFlowNegative{}},
 	{Path: "flow/iterate-values-flow.yaml", TestCase: &iterateValuesFlow{}},
+	{Path: "flow/iterate-one-value-flow.yaml", TestCase: &iterateOneValueFlow{}},
 	{Path: "flow/dns-ns-probe.yaml", TestCase: &dnsNsProbe{}},
 	{Path: "flow/flow-hide-matcher.yaml", TestCase: &flowHideMatcher{}},
 }
@@ -70,6 +71,16 @@ func (t *iterateValuesFlow) Execute(filePath string) error {
 	return expectResultsCount(results, 2)
 }
 
+type iterateOneValueFlow struct{}
+
+func (t *iterateOneValueFlow) Execute(filePath string) error {
+	results, err := testutils.RunNucleiTemplateAndGetResults(filePath, "https://scanme.sh", debug)
+	if err != nil {
+		return err
+	}
+	return expectResultsCount(results, 1)
+}
+
 type dnsNsProbe struct{}
 
 func (t *dnsNsProbe) Execute(filePath string) error {
diff --git a/integration_tests/flow/iterate-one-value-flow.yaml b/integration_tests/flow/iterate-one-value-flow.yaml
new file mode 100644
index 0000000000..d45eaa586b
--- /dev/null
+++ b/integration_tests/flow/iterate-one-value-flow.yaml
@@ -0,0 +1,34 @@
+id: flow-iterate-one-value-flow
+
+info:
+  name: Test Flow Iterate One Value Flow
+  author: pdteam
+  severity: info
+
+flow: |
+  http(1)
+  for(let value of template.extracted){
+    set("value", value)
+    http(2)
+  }
+
+http:
+  - method: GET
+    path:
+      - "{{BaseURL}}"
+
+    extractors:
+      - type: regex
+        name: extracted
+        internal: true
+        regex:
+          - "[ok]+"
+
+  - method: GET
+    path:
+      - "{{BaseURL}}/{{value}}"
+
+    matchers:
+      - type: word
+        words:
+          - "ok"
\ No newline at end of file
diff --git a/pkg/tmplexec/flow/flow_executor.go b/pkg/tmplexec/flow/flow_executor.go
index bfa35095cf..0cae628fbd 100644
--- a/pkg/tmplexec/flow/flow_executor.go
+++ b/pkg/tmplexec/flow/flow_executor.go
@@ -50,6 +50,9 @@ type FlowExecutor struct {
 	// logic related variables
 	results *atomic.Bool
 	allErrs mapsutil.SyncLockMap[string, error]
+	// these are keys whose values are meant to be flatten before executing
+	// a request ex: if dynamic extractor returns ["value"] it will be converted to "value"
+	flattenKeys []string
 }
 
 // NewFlowExecutor creates a new flow executor from a list of requests
@@ -157,6 +160,15 @@ func (f *FlowExecutor) Compile() error {
 					opts.reqIDS = append(opts.reqIDS, types.ToString(value))
 				}
 			}
+			// before executing any protocol function flatten tracked values
+			if len(f.flattenKeys) > 0 {
+				ctx := f.options.GetTemplateCtx(f.ctx.Input.MetaInput)
+				for _, key := range f.flattenKeys {
+					if value, ok := ctx.Get(key); ok {
+						ctx.Set(key, flatten(value))
+					}
+				}
+			}
 			return f.jsVM.ToValue(f.requestExecutor(reqMap, opts))
 		}
 	}
diff --git a/pkg/tmplexec/flow/flow_internal.go b/pkg/tmplexec/flow/flow_internal.go
index 24e2256942..2032a74443 100644
--- a/pkg/tmplexec/flow/flow_internal.go
+++ b/pkg/tmplexec/flow/flow_internal.go
@@ -100,11 +100,12 @@ func (f *FlowExecutor) protocolResultCallback(req protocols.Request, matcherStat
 					for k, v := range result.OperatorsResult.DynamicValues {
 						// if length of v is 1 then remove slice and convert it to single value
 						if len(v) == 1 {
-							f.options.GetTemplateCtx(f.ctx.Input.MetaInput).Set(k, v[0])
-						} else {
-							// if not let user handle it in flow ex: `for(let val of template.extracted)`
-							f.options.GetTemplateCtx(f.ctx.Input.MetaInput).Set(k, v)
+							// add it to flatten keys list so it will be flattened to a string later
+							f.flattenKeys = append(f.flattenKeys, k)
 						}
+						// always preserve extracted value type
+						f.options.GetTemplateCtx(f.ctx.Input.MetaInput).Set(k, v)
+
 					}
 				}
 			} else if !result.HasOperatorResult() && !hasOperators(req.GetCompiledOperators()) {
diff --git a/pkg/tmplexec/flow/util.go b/pkg/tmplexec/flow/util.go
index c7143c55d7..0e3e6fbd08 100644
--- a/pkg/tmplexec/flow/util.go
+++ b/pkg/tmplexec/flow/util.go
@@ -21,3 +21,20 @@ func hasOperators(all []*operators.Operators) bool {
 	}
 	return false
 }
+
+func flatten(v interface{}) interface{} {
+	switch v := v.(type) {
+	case []interface{}:
+		if len(v) == 1 {
+			return v[0]
+		}
+		return v
+	case []string:
+		if len(v) == 1 {
+			return v[0]
+		}
+		return v
+	default:
+		return v
+	}
+}