From 1874d5b90070632367321f0705f3834c0597a3de Mon Sep 17 00:00:00 2001
From: Nakul Bharti <knakul853@gmail.com>
Date: Fri, 7 Mar 2025 20:34:27 +0530
Subject: [PATCH] fix: prevent nil pointer dereference on context cancellation
 (#6085)

---
 pkg/core/execute_options.go | 4 ++--
 pkg/core/executors.go       | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/pkg/core/execute_options.go b/pkg/core/execute_options.go
index 4d27b5f66..aa47bc44f 100644
--- a/pkg/core/execute_options.go
+++ b/pkg/core/execute_options.go
@@ -107,6 +107,7 @@ func (e *Engine) executeTemplateSpray(ctx context.Context, templatesList []*temp
 	// wp is workpool that contains different waitgroups for
 	// headless and non-headless templates
 	wp := e.GetWorkPool()
+	defer wp.Wait()
 
 	for _, template := range templatesList {
 		select {
@@ -135,7 +136,6 @@ func (e *Engine) executeTemplateSpray(ctx context.Context, templatesList []*temp
 			e.executeTemplateWithTargets(ctx, tpl, target, results)
 		}(template)
 	}
-	wp.Wait()
 	return results
 }
 
@@ -143,6 +143,7 @@ func (e *Engine) executeTemplateSpray(ctx context.Context, templatesList []*temp
 func (e *Engine) executeHostSpray(ctx context.Context, templatesList []*templates.Template, target provider.InputProvider) *atomic.Bool {
 	results := &atomic.Bool{}
 	wp, _ := syncutil.New(syncutil.WithSize(e.options.BulkSize + e.options.HeadlessBulkSize))
+	defer wp.Wait()
 
 	target.Iterate(func(value *contextargs.MetaInput) bool {
 		select {
@@ -158,7 +159,6 @@ func (e *Engine) executeHostSpray(ctx context.Context, templatesList []*template
 		}(value)
 		return true
 	})
-	wp.Wait()
 	return results
 }
 
diff --git a/pkg/core/executors.go b/pkg/core/executors.go
index 1f9af56ac..2e8c4d18d 100644
--- a/pkg/core/executors.go
+++ b/pkg/core/executors.go
@@ -163,6 +163,7 @@ func (e *Engine) executeTemplatesOnTarget(ctx context.Context, alltemplates []*t
 	// headless and non-headless templates
 	// global waitgroup should not be used here
 	wp := e.GetWorkPool()
+	defer wp.Wait()
 
 	for _, tpl := range alltemplates {
 		select {
@@ -210,5 +211,4 @@ func (e *Engine) executeTemplatesOnTarget(ctx context.Context, alltemplates []*t
 			results.CompareAndSwap(false, match)
 		}(tpl, target, sg)
 	}
-	wp.Wait()
 }