fix: prevent nil pointer dereference on context cancellation

projectdiscovery · Mar 4, 2025 · ee90916 · ee90916
1 parent 3a3f5e2
commit ee90916
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git 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,14 +136,14 @@ func (e *Engine) executeTemplateSpray(ctx context.Context, templatesList []*temp
 			e.executeTemplateWithTargets(ctx, tpl, target, results)
 		}(template)
 	}
-	wp.Wait()
 	return results
 }
 
 // executeHostSpray executes scan using host spray strategy where templates are iterated over each target
 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
@@ -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()
 }