perf: replace bytes.Buffer & strings.Builder with bytebufferpool

cmd/docgen/docgen.go

@@ -1,14 +1,14 @@
 package main
 
 import (
-	"bytes"
 	"log"
 	"os"
 	"reflect"
 	"regexp"
 	"strings"
 
 	"github.com/invopop/jsonschema"
+	"github.com/valyala/bytebufferpool"
 
 	"github.com/projectdiscovery/nuclei/v3/pkg/templates"
 	"github.com/projectdiscovery/nuclei/v3/pkg/utils/json"
@@ -42,8 +42,9 @@ func main() {
 	}
 	jsonschemaData := r.Reflect(&templates.Template{})
 
-	var buf bytes.Buffer
-	encoder := json.NewEncoder(&buf)
+	buf := bytebufferpool.Get()
+	defer bytebufferpool.Put(buf)
+	encoder := json.NewEncoder(buf)
 	encoder.SetIndent("", "  ")
 	_ = encoder.Encode(jsonschemaData)
 

cmd/tmc/main.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"bytes"
 	"fmt"
 	"log"
 	"os"
@@ -24,6 +23,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/utils/json"
 	"github.com/projectdiscovery/retryablehttp-go"
 	errorutil "github.com/projectdiscovery/utils/errors"
+	"github.com/valyala/bytebufferpool"
 	"gopkg.in/yaml.v3"
 )
 
@@ -379,8 +379,9 @@ func parseAndAddMaxRequests(catalog catalog.Catalog, path, data string) (string,
 	}
 	infoBlock.Info.Metadata["max-request"] = template.TotalRequests
 
-	var newInfoBlock bytes.Buffer
-	yamlEncoder := yaml.NewEncoder(&newInfoBlock)
+	newInfoBlock := bytebufferpool.Get()
+	defer bytebufferpool.Put(newInfoBlock)
+	yamlEncoder := yaml.NewEncoder(newInfoBlock)
 	yamlEncoder.SetIndent(yamlIndentSpaces)
 	err = yamlEncoder.Encode(infoBlock)
 	if err != nil {

internal/pdcp/writer.go

@@ -23,6 +23,7 @@ import (
 	unitutils "github.com/projectdiscovery/utils/unit"
 	updateutils "github.com/projectdiscovery/utils/update"
 	urlutil "github.com/projectdiscovery/utils/url"
+	"github.com/valyala/bytebufferpool"
 )
 
 const (
@@ -151,7 +152,8 @@ func (u *UploadWriter) autoCommit(ctx context.Context, r *io.PipeReader) {
 		}
 	}()
 	// temporary buffer to store the results
-	buff := &bytes.Buffer{}
+	buff := bytebufferpool.Get()
+	defer bytebufferpool.Put(buff)
 	ticker := time.NewTicker(flushTimer)
 	defer ticker.Stop()
 	for {
@@ -193,7 +195,7 @@ func (u *UploadWriter) autoCommit(ctx context.Context, r *io.PipeReader) {
 }
 
 // uploadChunk uploads a chunk of data to the server
-func (u *UploadWriter) uploadChunk(buff *bytes.Buffer) error {
+func (u *UploadWriter) uploadChunk(buff *bytebufferpool.ByteBuffer) error {
 	if err := u.upload(buff.Bytes()); err != nil {
 		return errorutil.NewWithErr(err).Msgf("could not upload chunk")
 	}

internal/runner/healthcheck.go

@@ -4,17 +4,18 @@ import (
 	"fmt"
 	"net"
 	"runtime"
-	"strings"
 
 	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/config"
 	"github.com/projectdiscovery/nuclei/v3/pkg/types"
 	fileutil "github.com/projectdiscovery/utils/file"
+	"github.com/valyala/bytebufferpool"
 )
 
 // DoHealthCheck performs self-diagnostic checks
 func DoHealthCheck(options *types.Options) string {
 	// RW permissions on config file
-	var test strings.Builder
+	test := bytebufferpool.Get()
+	defer bytebufferpool.Put(test)
 	test.WriteString(fmt.Sprintf("Version: %s\n", config.Version))
 	test.WriteString(fmt.Sprintf("Operating System: %s\n", runtime.GOOS))
 	test.WriteString(fmt.Sprintf("Architecture: %s\n", runtime.GOARCH))

internal/runner/templates.go

@@ -1,7 +1,6 @@
 package runner
 
 import (
-	"bytes"
 	"path/filepath"
 	"sort"
 	"strings"
@@ -11,6 +10,7 @@ import (
 	"github.com/logrusorgru/aurora"
 	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/config"
 	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/loader"
+	"github.com/valyala/bytebufferpool"
 
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/nuclei/v3/pkg/templates"
@@ -108,9 +108,10 @@ func (r *Runner) listAvailableStoreTags(store *loader.Store) {
 }
 
 func (r *Runner) highlightTemplate(body *[]byte) ([]byte, error) {
-	var buf bytes.Buffer
+	buf := bytebufferpool.Get()
+	defer bytebufferpool.Put(buf)
 	// YAML lexer, true color terminal formatter and monokai style
-	err := quick.Highlight(&buf, string(*body), "yaml", "terminal16m", "monokai")
+	err := quick.Highlight(buf, string(*body), "yaml", "terminal16m", "monokai")
 	if err != nil {
 		return nil, err
 	}

internal/server/dedupe.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/projectdiscovery/nuclei/v3/pkg/input/types"
 	mapsutil "github.com/projectdiscovery/utils/maps"
+	"github.com/valyala/bytebufferpool"
 )
 
 var dynamicHeaders = map[string]bool{
@@ -68,7 +69,9 @@ func hashRequest(req *types.RequestResponse) (string, error) {
 		return "", err
 	}
 
-	var hashContent strings.Builder
+	hashContent := bytebufferpool.Get()
+	defer bytebufferpool.Put(hashContent)
+
 	hashContent.WriteString(req.Request.Method)
 	hashContent.WriteString(normalizedURL)
 

internal/server/server.go

@@ -6,7 +6,6 @@ import (
 	"html/template"
 	"net/http"
 	"net/url"
-	"strings"
 	"sync/atomic"
 	"time"
 
@@ -20,6 +19,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/output"
 	"github.com/projectdiscovery/nuclei/v3/pkg/protocols"
 	"github.com/projectdiscovery/utils/env"
+	"github.com/valyala/bytebufferpool"
 )
 
 // DASTServer is a server that performs execution of fuzzing templates
@@ -98,7 +98,8 @@ func New(options *Options) (*DASTServer, error) {
 	}
 	server.scopeManager = scopeManager
 
-	var builder strings.Builder
+	builder := bytebufferpool.Get()
+	defer bytebufferpool.Put(builder)
 	gologger.Debug().Msgf("Using %d parallel tasks with %d buffer", maxWorkers, bufferSize)
 	if options.Token != "" {
 		builder.WriteString(" (with token)")

lib/sdk.go

@@ -27,6 +27,7 @@ import (
 	"github.com/projectdiscovery/ratelimit"
 	"github.com/projectdiscovery/retryablehttp-go"
 	errorutil "github.com/projectdiscovery/utils/errors"
+	"github.com/valyala/bytebufferpool"
 )
 
 // NucleiSDKOptions contains options for nuclei SDK
@@ -188,9 +189,11 @@ func (e *NucleiEngine) SignTemplate(tmplSigner *signer.TemplateSigner, data []by
 		return data, err
 	}
 	_, content := signer.ExtractSignatureAndContent(data)
-	buff := bytes.NewBuffer(content)
+	buff := bytebufferpool.Get()
+	defer bytebufferpool.Put(buff)
+	buff.Write(content)
 	buff.WriteString("\n" + signatureData)
-	return buff.Bytes(), err
+	return buff.Bytes(), nil
 }
 
 func (e *NucleiEngine) closeInternal() {

lib/sdk_private.go

@@ -3,11 +3,11 @@ package nuclei
 import (
 	"context"
 	"fmt"
-	"strings"
 	"sync"
 	"time"
 
 	"github.com/projectdiscovery/nuclei/v3/pkg/input"
+	"github.com/valyala/bytebufferpool"
 
 	"github.com/logrusorgru/aurora"
 	"github.com/pkg/errors"
@@ -51,7 +51,8 @@ func (e *NucleiEngine) applyRequiredDefaults(ctx context.Context) {
 			}
 			return
 		}
-		sb := strings.Builder{}
+		sb := bytebufferpool.Get()
+		defer bytebufferpool.Put(sb)
 		sb.WriteString(fmt.Sprintf("[%v] ", event.TemplateID))
 		if event.Matched != "" {
 			sb.WriteString(event.Matched)

pkg/catalog/config/nucleiconfig.go

@@ -1,7 +1,6 @@
 package config
 
 import (
-	"bytes"
 	"crypto/md5"
 	"fmt"
 	"log"
@@ -16,6 +15,7 @@ import (
 	errorutil "github.com/projectdiscovery/utils/errors"
 	fileutil "github.com/projectdiscovery/utils/file"
 	folderutil "github.com/projectdiscovery/utils/folder"
+	"github.com/valyala/bytebufferpool"
 )
 
 // DefaultConfig is the default nuclei configuration
@@ -292,7 +292,8 @@ func (c *Config) WriteTemplatesConfig() error {
 // WriteTemplatesIndex writes the nuclei templates index file
 func (c *Config) WriteTemplatesIndex(index map[string]string) error {
 	indexFile := c.GetTemplateIndexFilePath()
-	var buff bytes.Buffer
+	buff := bytebufferpool.Get()
+	defer bytebufferpool.Put(buff)
 	for k, v := range index {
 		_, _ = buff.WriteString(k + "," + v + "\n")
 	}

pkg/catalog/loader/ai_loader.go

@@ -16,6 +16,7 @@ import (
 	"github.com/projectdiscovery/retryablehttp-go"
 	pdcpauth "github.com/projectdiscovery/utils/auth/pdcp"
 	errorutil "github.com/projectdiscovery/utils/errors"
+	"github.com/valyala/bytebufferpool"
 )
 
 const (
@@ -70,8 +71,9 @@ func getAIGeneratedTemplates(prompt string, options *types.Options) ([]string, e
 	if hasNoTargets && hasNoStdin {
 		// Display the template content with syntax highlighting
 		if !options.NoColor {
-			var buf bytes.Buffer
-			err = quick.Highlight(&buf, template, "yaml", "terminal16m", "monokai")
+			buf := bytebufferpool.Get()
+			defer bytebufferpool.Put(buf)
+			err = quick.Highlight(buf, template, "yaml", "terminal16m", "monokai")
 			if err == nil {
 				template = buf.String()
 			}

pkg/external/customtemplates/azure_blob.go

@@ -1,7 +1,6 @@
 package customtemplates
 
 import (
-	"bytes"
 	"context"
 	"os"
 	"path/filepath"
@@ -13,6 +12,7 @@ import (
 	"github.com/projectdiscovery/nuclei/v3/pkg/catalog/config"
 	"github.com/projectdiscovery/nuclei/v3/pkg/types"
 	errorutil "github.com/projectdiscovery/utils/errors"
+	"github.com/valyala/bytebufferpool"
 )
 
 var _ Provider = &customTemplateAzureBlob{}
@@ -121,7 +121,8 @@ func downloadTemplate(client *azblob.Client, containerName string, path string,
 		return err
 	}
 
-	downloadedData := bytes.Buffer{}
+	downloadedData := bytebufferpool.Get()
+	defer bytebufferpool.Put(downloadedData)
 	retryReader := get.NewRetryReader(ctx, &azblob.RetryReaderOptions{})
 	_, err = downloadedData.ReadFrom(retryReader)
 	if err != nil {

pkg/fuzz/analyzers/time/time_delay.go

@@ -28,7 +28,8 @@ import (
 	"errors"
 	"fmt"
 	"math"
-	"strings"
+
+	"github.com/valyala/bytebufferpool"
 )
 
 type timeDelayRequestSender func(delay int) (float64, error)
@@ -101,7 +102,8 @@ func checkTimingDependency(
 
 	result := regression.IsWithinConfidence(correlationErrorRange, 1.0, slopeErrorRange)
 	if result {
-		var resultReason strings.Builder
+		resultReason := bytebufferpool.Get()
+		defer bytebufferpool.Put(resultReason)
 		resultReason.WriteString(fmt.Sprintf(
 			"[time_delay] made %d requests (baseline: %.2fs) successfully, with a regression slope of %.2f and correlation %.2f",
 			requestsLimit,

pkg/fuzz/component/body_test.go

@@ -1,7 +1,6 @@
 package component
 
 import (
-	"bytes"
 	"io"
 	"mime/multipart"
 	"strings"
@@ -10,6 +9,7 @@ import (
 	"github.com/projectdiscovery/retryablehttp-go"
 	urlutil "github.com/projectdiscovery/utils/url"
 	"github.com/stretchr/testify/require"
+	"github.com/valyala/bytebufferpool"
 )
 
 func TestBodyComponent(t *testing.T) {
@@ -122,7 +122,8 @@ func TestBodyFormComponent(t *testing.T) {
 }
 
 func TestMultiPartFormComponent(t *testing.T) {
-	formData := &bytes.Buffer{}
+	formData := bytebufferpool.Get()
+	defer bytebufferpool.Put(formData)
 	writer := multipart.NewWriter(formData)
 
 	// Hypothetical form fields

pkg/fuzz/dataformat/multipart.go

@@ -9,6 +9,7 @@ import (
 	"net/textproto"
 
 	mapsutil "github.com/projectdiscovery/utils/maps"
+	"github.com/valyala/bytebufferpool"
 )
 
 type MultiPartForm struct {
@@ -38,8 +39,9 @@ func (m *MultiPartForm) IsType(data string) bool {
 
 // Encode encodes the data into MultiPartForm format
 func (m *MultiPartForm) Encode(data KV) (string, error) {
-	var b bytes.Buffer
-	w := multipart.NewWriter(&b)
+	b := bytebufferpool.Get()
+	defer bytebufferpool.Put(b)
+	w := multipart.NewWriter(b)
 	if err := w.SetBoundary(m.boundary); err != nil {
 		return "", err
 	}
@@ -144,7 +146,8 @@ func (m *MultiPartForm) Decode(data string) (KV, error) {
 			}
 			defer file.Close()
 
-			buffer := new(bytes.Buffer)
+			buffer := bytebufferpool.Get()
+			defer bytebufferpool.Put(buffer)
 			if _, err := buffer.ReadFrom(file); err != nil {
 				return KV{}, err
 			}

pkg/fuzz/frequency/tracker.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/bluele/gcache"
 	"github.com/projectdiscovery/gologger"
+	"github.com/valyala/bytebufferpool"
 )
 
 // Tracker implements a frequency tracker for a given input
@@ -121,7 +122,8 @@ func (t *Tracker) UnmarkParameter(parameter, target, template string) {
 }
 
 func getFrequencyKey(parameter, target, template string) string {
-	var sb strings.Builder
+	sb := bytebufferpool.Get()
+	defer bytebufferpool.Put(sb)
 	sb.WriteString(target)
 	sb.WriteString(":")
 	sb.WriteString(template)