From 1c51a6bef6a5d9a5b9c4967e4984595ef4bb48cd Mon Sep 17 00:00:00 2001
From: Mzack9999 <mzack9999@protonmail.com>
Date: Sat, 22 Jun 2024 19:12:00 +0200
Subject: [PATCH] improving workflows (#5318)

* improving workflows

* fixing count
---
 pkg/templates/compile.go                  |  5 +++++
 pkg/templates/compile_test.go             |  9 +++++++++
 pkg/templates/templates.go                |  4 ++--
 pkg/templates/tests/workflow-invalid.yaml | 18 ++++++++++++++++++
 4 files changed, 34 insertions(+), 2 deletions(-)
 create mode 100644 pkg/templates/tests/workflow-invalid.yaml

diff --git a/pkg/templates/compile.go b/pkg/templates/compile.go
index 39d6a6ea0..54708e7ea 100644
--- a/pkg/templates/compile.go
+++ b/pkg/templates/compile.go
@@ -342,6 +342,11 @@ func parseTemplate(data []byte, options protocols.ExecutorOptions) (*Template, e
 		return nil, errors.New("no template author field provided")
 	}
 
+	numberOfWorkflows := len(template.Workflows)
+	if numberOfWorkflows > 0 && numberOfWorkflows != template.Requests() {
+		return nil, errors.New("workflows cannot have other protocols")
+	}
+
 	// use default unknown severity
 	if len(template.Workflows) == 0 {
 		if template.Info.SeverityHolder.Severity == severity.Undefined {
diff --git a/pkg/templates/compile_test.go b/pkg/templates/compile_test.go
index 3a230cc52..81c5f8fd5 100644
--- a/pkg/templates/compile_test.go
+++ b/pkg/templates/compile_test.go
@@ -197,3 +197,12 @@ func Test_WrongTemplate(t *testing.T) {
 	require.Nil(t, got, "could not parse template")
 	require.ErrorContains(t, err, "no requests defined ")
 }
+
+func TestWrongWorkflow(t *testing.T) {
+	setup()
+
+	filePath := "tests/workflow-invalid.yaml"
+	got, err := templates.Parse(filePath, nil, executerOpts)
+	require.Nil(t, got, "could not parse template")
+	require.ErrorContains(t, err, "workflows cannot have other protocols")
+}
diff --git a/pkg/templates/templates.go b/pkg/templates/templates.go
index 045f947c6..df0c47648 100644
--- a/pkg/templates/templates.go
+++ b/pkg/templates/templates.go
@@ -177,8 +177,6 @@ func (template *Template) Type() types.ProtocolType {
 		return types.HeadlessProtocol
 	case len(template.RequestsNetwork) > 0:
 		return types.NetworkProtocol
-	case len(template.Workflow.Workflows) > 0:
-		return types.WorkflowProtocol
 	case len(template.RequestsSSL) > 0:
 		return types.SSLProtocol
 	case len(template.RequestsWebsocket) > 0:
@@ -189,6 +187,8 @@ func (template *Template) Type() types.ProtocolType {
 		return types.CodeProtocol
 	case len(template.RequestsJavascript) > 0:
 		return types.JavascriptProtocol
+	case len(template.Workflow.Workflows) > 0:
+		return types.WorkflowProtocol
 	default:
 		return types.InvalidProtocol
 	}
diff --git a/pkg/templates/tests/workflow-invalid.yaml b/pkg/templates/tests/workflow-invalid.yaml
new file mode 100644
index 000000000..981823aa0
--- /dev/null
+++ b/pkg/templates/tests/workflow-invalid.yaml
@@ -0,0 +1,18 @@
+id: workflow-example
+
+info:
+  name: Test Invalid Workflow Template
+  author: pdteam
+  severity: info
+
+http:
+  - raw:
+      - |
+        POST /re HTTP/1.1
+        Host: {{Hostname}}
+
+        {{code_response}}
+
+workflows:
+  - template: tests/match-1.yaml
+  - template: tests/match-1.yaml