close protocolstate shared resources in nuclei sdk/lib

projectdiscovery · Apr 29, 2024 · 20e0895 · 20e0895
1 parent c9fbbe8
commit 20e0895
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/lib/sdk.go b/lib/sdk.go
@@ -18,6 +18,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/hosterrorscache"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/interactsh"
+	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols/headless/engine"
 	"github.com/projectdiscovery/nuclei/v3/pkg/reporting"
 	"github.com/projectdiscovery/nuclei/v3/pkg/templates"
@@ -192,6 +193,8 @@ func (e *NucleiEngine) Close() {
 	if e.rateLimiter != nil {
 		e.rateLimiter.Stop()
 	}
+	// close global shared resources
+	protocolstate.Close()
 }
 
 // ExecuteWithCallback executes templates on targets and calls callback on each result(only if results are found)

diff --git a/pkg/protocols/common/protocolstate/memguardian.go b/pkg/protocols/common/protocolstate/memguardian.go
@@ -18,14 +18,13 @@ var (
 	cancelFunc                     context.CancelFunc
 )
 
-func StartActiveMemGuardian() {
+func StartActiveMemGuardian(ctx context.Context) {
 	if memguardian.DefaultMemGuardian == nil {
 		return
 	}
 
 	memTimer = time.NewTicker(memguardian.DefaultInterval)
-	var ctx context.Context
-	ctx, cancelFunc = context.WithCancel(context.Background())
+	ctx, cancelFunc = context.WithCancel(ctx)
 	go func() {
 		for {
 			select {

diff --git a/pkg/protocols/common/protocolstate/state.go b/pkg/protocols/common/protocolstate/state.go
@@ -148,7 +148,7 @@ func Init(options *types.Options) error {
 		return Dialer.Dial(ctx, "tcp", addr)
 	})
 
-	StartActiveMemGuardian()
+	StartActiveMemGuardian(context.Background())
 
 	return nil
 }