From a326f3925c5a00c8d8c214cdf839bf8075fe0cbe Mon Sep 17 00:00:00 2001 From: Dwi Siswanto <25837540+dwisiswant0@users.noreply.github.com> Date: Tue, 17 Jun 2025 06:23:32 +0700 Subject: [PATCH] fix(tmplexec): memory blowup in multiproto (#6258) * bugfix: fix memory blowup using previousEvent for multi-proto execution * refactor(tmplexec): uses supported protocol types Signed-off-by: Dwi Siswanto * add co-author Co-authored-by: Nakul Bharti Signed-off-by: Dwi Siswanto * refactor(tmplexec): mv builder inside loop scope Signed-off-by: Dwi Siswanto * refactor(tmplexec): skip existing keys in `FillPreviousEvent` The `FillPreviousEvent` func was modified to prevent overwriting/duplicating entries in the previous map. It now checks if a key `k` from `event.InternalEvent` already exists in the previous map. If it does, the key is skipped. This ensures that if `k` was already set (potentially w/o a prefix), it's not re-added with an `ID_` prefix. Additionally, keys in `event.InternalEvent` that already start with the current `ID_` prefix are also skipped to avoid redundant prefixing. This change simplifies the logic by removing the `reqTypeWithIndexRegex` and directly addresses the potential for duplicate / incorrectly prefixed keys when `event.InternalEvent` grows during protocol request execution. Signed-off-by: Dwi Siswanto * chore(tmplexec): naming convention, `ID` => `protoID` Signed-off-by: Dwi Siswanto * chore(tmplexec): it's request ID lol sorry Signed-off-by: Dwi Siswanto --------- Signed-off-by: Dwi Siswanto Co-authored-by: Ice3man Co-authored-by: Nakul Bharti --- pkg/tmplexec/generic/exec.go | 16 ++++----------- pkg/tmplexec/multiproto/multi.go | 14 ++----------- pkg/tmplexec/utils/utils.go | 34 ++++++++++++++++++++++++++++++++ 3 files changed, 40 insertions(+), 24 deletions(-) create mode 100644 pkg/tmplexec/utils/utils.go diff --git a/pkg/tmplexec/generic/exec.go b/pkg/tmplexec/generic/exec.go index c017810e7..25f88262e 100644 --- a/pkg/tmplexec/generic/exec.go +++ b/pkg/tmplexec/generic/exec.go @@ -1,13 +1,13 @@ package generic import ( - "strings" "sync/atomic" "github.com/projectdiscovery/gologger" "github.com/projectdiscovery/nuclei/v3/pkg/output" "github.com/projectdiscovery/nuclei/v3/pkg/protocols" "github.com/projectdiscovery/nuclei/v3/pkg/scan" + "github.com/projectdiscovery/nuclei/v3/pkg/tmplexec/utils" mapsutil "github.com/projectdiscovery/utils/maps" ) @@ -64,17 +64,9 @@ func (g *Generic) ExecuteWithResults(ctx *scan.ScanContext) error { // ideally this should never happen since protocol exits on error and callback is not called return } - ID := req.GetID() - if ID != "" { - builder := &strings.Builder{} - for k, v := range event.InternalEvent { - builder.WriteString(ID) - builder.WriteString("_") - builder.WriteString(k) - _ = previous.Set(builder.String(), v) - builder.Reset() - } - } + + utils.FillPreviousEvent(req.GetID(), event, previous) + if event.HasOperatorResult() { g.results.CompareAndSwap(false, true) } diff --git a/pkg/tmplexec/multiproto/multi.go b/pkg/tmplexec/multiproto/multi.go index d50e168ac..ef029861a 100644 --- a/pkg/tmplexec/multiproto/multi.go +++ b/pkg/tmplexec/multiproto/multi.go @@ -2,7 +2,6 @@ package multiproto import ( "strconv" - "strings" "sync/atomic" "github.com/projectdiscovery/nuclei/v3/pkg/output" @@ -10,6 +9,7 @@ import ( "github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/generators" "github.com/projectdiscovery/nuclei/v3/pkg/scan" "github.com/projectdiscovery/nuclei/v3/pkg/templates/types" + "github.com/projectdiscovery/nuclei/v3/pkg/tmplexec/utils" mapsutil "github.com/projectdiscovery/utils/maps" stringsutil "github.com/projectdiscovery/utils/strings" ) @@ -90,17 +90,7 @@ func (m *MultiProtocol) ExecuteWithResults(ctx *scan.ScanContext) error { return } - ID := req.GetID() - if ID != "" { - builder := &strings.Builder{} - for k, v := range event.InternalEvent { - builder.WriteString(ID) - builder.WriteString("_") - builder.WriteString(k) - _ = previous.Set(builder.String(), v) - builder.Reset() - } - } + utils.FillPreviousEvent(req.GetID(), event, previous) // log event and generate result for the event ctx.LogEvent(event) diff --git a/pkg/tmplexec/utils/utils.go b/pkg/tmplexec/utils/utils.go new file mode 100644 index 000000000..41d717769 --- /dev/null +++ b/pkg/tmplexec/utils/utils.go @@ -0,0 +1,34 @@ +package utils + +import ( + "strings" + + "github.com/projectdiscovery/nuclei/v3/pkg/output" + mapsutil "github.com/projectdiscovery/utils/maps" +) + +// FillPreviousEvent is a helper function to get the previous event from the event +// without leading to duplicate prefixes +func FillPreviousEvent(reqID string, event *output.InternalWrappedEvent, previous *mapsutil.SyncLockMap[string, any]) { + if reqID == "" { + return + } + + for k, v := range event.InternalEvent { + if _, ok := previous.Get(k); ok { + continue + } + + if strings.HasPrefix(k, reqID+"_") { + continue + } + + var builder strings.Builder + + builder.WriteString(reqID) + builder.WriteString("_") + builder.WriteString(k) + + _ = previous.Set(builder.String(), v) + } +}