From c4c281e241d580d56682da4736e8852268f88ec4 Mon Sep 17 00:00:00 2001
From: Tarun Koyalwar <tarun@projectdiscovery.io>
Date: Mon, 6 May 2024 21:16:43 +0530
Subject: [PATCH] add context support in sdk

---
 lib/multi.go | 10 ++++++++--
 lib/sdk.go   | 13 +++++++++++--
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/lib/multi.go b/lib/multi.go
index 726b59c5d9..5b2c7d6776 100644
--- a/lib/multi.go
+++ b/lib/multi.go
@@ -111,11 +111,11 @@ func (e *ThreadSafeNucleiEngine) GlobalResultCallback(callback func(event *outpu
 	e.eng.resultCallbacks = []func(*output.ResultEvent){callback}
 }
 
-// ExecuteNucleiWithOpts executes templates on targets and calls callback on each result(only if results are found)
+// ExecuteNucleiWithOptsCtx executes templates on targets and calls callback on each result(only if results are found)
 // This method can be called concurrently and it will use some global resources but can be runned parallelly
 // by invoking this method with different options and targets
 // Note: Not all options are thread-safe. this method will throw error if you try to use non-thread-safe options
-func (e *ThreadSafeNucleiEngine) ExecuteNucleiWithOpts(targets []string, opts ...NucleiSDKOptions) error {
+func (e *ThreadSafeNucleiEngine) ExecuteNucleiWithOptsCtx(ctx context.Context, targets []string, opts ...NucleiSDKOptions) error {
 	baseOpts := *e.eng.opts
 	tmpEngine := &NucleiEngine{opts: &baseOpts, mode: threadSafe}
 	for _, option := range opts {
@@ -163,6 +163,12 @@ func (e *ThreadSafeNucleiEngine) ExecuteNucleiWithOpts(targets []string, opts ..
 	return nil
 }
 
+// ExecuteNucleiWithOpts is same as ExecuteNucleiWithOptsCtx but with default context
+// This is a placeholder and will be deprecated in future major release
+func (e *ThreadSafeNucleiEngine) ExecuteNucleiWithOpts(targets []string, opts ...NucleiSDKOptions) error {
+	return e.ExecuteNucleiWithOptsCtx(context.Background(), targets, opts...)
+}
+
 // Close all resources used by nuclei engine
 func (e *ThreadSafeNucleiEngine) Close() {
 	e.eng.Close()
diff --git a/lib/sdk.go b/lib/sdk.go
index 7959d4c1d5..63925f47ca 100644
--- a/lib/sdk.go
+++ b/lib/sdk.go
@@ -219,8 +219,9 @@ func (e *NucleiEngine) Close() {
 	}
 }
 
-// ExecuteWithCallback executes templates on targets and calls callback on each result(only if results are found)
-func (e *NucleiEngine) ExecuteWithCallback(callback ...func(event *output.ResultEvent)) error {
+// ExecuteCallbackWithCtx executes templates on targets and calls callback on each result(only if results are found)
+// enable matcher-status option if you expect this callback to be called for all results regardless if it matched or not
+func (e *NucleiEngine) ExecuteCallbackWithCtx(ctx context.Context, callback ...func(event *output.ResultEvent)) error {
 	if !e.templatesLoaded {
 		_ = e.LoadAllTemplates()
 	}
@@ -244,10 +245,18 @@ func (e *NucleiEngine) ExecuteWithCallback(callback ...func(event *output.Result
 	return nil
 }
 
+// ExecuteWithCallback is same as ExecuteCallbackWithCtx but with default context
+// Note this is deprecated and will be removed in future major release
+func (e *NucleiEngine) ExecuteWithCallback(callback ...func(event *output.ResultEvent)) error {
+	return e.ExecuteCallbackWithCtx(context.Background(), callback...)
+}
+
+// Options return nuclei Type Options
 func (e *NucleiEngine) Options() *types.Options {
 	return e.opts
 }
 
+// Engine returns core Executer of nuclei
 func (e *NucleiEngine) Engine() *core.Engine {
 	return e.engine
 }