From 1dff3b7ece4a959f59f0f213ac8b3258ba79a67e Mon Sep 17 00:00:00 2001 From: Dwi Siswanto Date: Wed, 26 Nov 2025 14:02:57 +0700 Subject: [PATCH 1/2] fix(httputil): racy in `ResponseChain` and optimize memory usage. This commit addresses a critical regression and improves memory efficiency in the `httputil` pkg. Changes: * Previously, it returned a slice referencing a buffer that was immediately returned to the pool, leading to data corruption (use-after-free). It now allocates a safe, independent copy. * Make `Close()` idempotent. Calling `Close()` multiple times no longer causes a panic. * Optimize `FullResponseString` to use zero-copy conversion from the safe byte slice. * Optimize `limitedBuffer` to use a `sync.Pool` for temp 32KB (via `chunkSize` = `DefaultChunkSize`) read chunks. Signed-off-by: Dwi Siswanto --- http/normalization.go | 41 ++++++++++++++++++++++++-- http/respChain.go | 25 ++++++++++------ http/respChain_test.go | 67 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 122 insertions(+), 11 deletions(-) diff --git a/http/normalization.go b/http/normalization.go index 831c5e8..b210d9f 100644 --- a/http/normalization.go +++ b/http/normalization.go @@ -6,6 +6,7 @@ import ( "io" "net/http" "strings" + "sync" "github.com/dsnet/compress/brotli" "github.com/klauspost/compress/gzip" @@ -18,6 +19,41 @@ import ( stringsutil "github.com/projectdiscovery/utils/strings" ) +const ( + // DefaultChunkSize defines the default chunk size for reading response + // bodies. + // + // Use [SetChunkSize] to adjust the size. + DefaultChunkSize = 32 * 1024 // 32KB +) + +var ( + chunkSize = DefaultChunkSize + chunkPool = sync.Pool{ + New: func() any { + b := make([]byte, chunkSize) + return &b + }, + } +) + +// SetChunkSize sets the chunk size for reading response bodies. +// +// If size is less than or equal to zero, it resets to the default chunk size. +func SetChunkSize(size int) { + if size <= 0 { + size = DefaultChunkSize + } + + chunkSize = size + chunkPool = sync.Pool{ + New: func() any { + b := make([]byte, chunkSize) + return &b + }, + } +} + // limitedBuffer wraps [bytes.Buffer] to prevent capacity growth beyond maxCap. // This prevents bytes.Buffer.ReadFrom() from over-allocating when it doesn't // know the final size. @@ -27,8 +63,9 @@ type limitedBuffer struct { } func (lb *limitedBuffer) ReadFrom(r io.Reader) (n int64, err error) { - const chunkSize = 32 * 1024 // 32KB chunks - chunk := make([]byte, chunkSize) + chunkPtr := chunkPool.Get().(*[]byte) + defer chunkPool.Put(chunkPtr) + chunk := *chunkPtr for { available := lb.buf.Cap() - lb.buf.Len() diff --git a/http/respChain.go b/http/respChain.go index 5ea297e..18be06c 100644 --- a/http/respChain.go +++ b/http/respChain.go @@ -306,12 +306,14 @@ func (r *ResponseChain) FullResponse() *bytes.Buffer { // FullResponseBytes returns the current response (headers+body) as byte slice. // // The returned slice is valid only until Close() is called. -// Note: This creates a new buffer internally which is returned to the pool. func (r *ResponseChain) FullResponseBytes() []byte { - buf := r.FullResponse() - defer putBuffer(buf) + size := r.headers.Len() + r.body.Len() + buf := make([]byte, size) - return buf.Bytes() + copy(buf, r.headers.Bytes()) + copy(buf[r.headers.Len():], r.body.Bytes()) + + return buf } // FullResponseString returns the current response as string in the chain. @@ -319,7 +321,7 @@ func (r *ResponseChain) FullResponseBytes() []byte { // The returned string is valid only until Close() is called. // This is a zero-copy operation for performance. func (r *ResponseChain) FullResponseString() string { - return conversion.String(r.FullResponse().Bytes()) + return conversion.String(r.FullResponseBytes()) } // previous updates response pointer to previous response @@ -372,10 +374,15 @@ func (r *ResponseChain) Fill() error { // Close the response chain and releases the buffers. func (r *ResponseChain) Close() { - putBuffer(r.headers) - putBuffer(r.body) - r.headers = nil - r.body = nil + if r.headers != nil { + putBuffer(r.headers) + r.headers = nil + } + + if r.body != nil { + putBuffer(r.body) + r.body = nil + } } // Has returns true if the response chain has a response diff --git a/http/respChain_test.go b/http/respChain_test.go index 5890eda..fdfe001 100644 --- a/http/respChain_test.go +++ b/http/respChain_test.go @@ -1115,3 +1115,70 @@ func TestResponseChain_BurstWithPoolExhaustion(t *testing.T) { require.NoError(t, err) } } + +func TestResponseChain_FullResponseBytes_Race(t *testing.T) { + resp := &http.Response{ + StatusCode: 200, + Status: "200 OK", + } + chain := NewResponseChain(resp, 1024) + chain.headers.WriteString("Header: Value\r\n") + chain.body.WriteString("Body Content") + + data := chain.FullResponseBytes() + initialContent := string(data) + + // Trigger buffer reuse + // We need to get a buffer from the pool. + + found := false + for i := 0; i < 100; i++ { + b := getBuffer() + b.WriteString("OVERWRITTEN_DATA_XXXXXXXXXXXXXXXX") + + if string(data) != initialContent { + found = true + t.Logf("Iteration %d: Content changed to %q", i, string(data)) + + break + } + + putBuffer(b) + } + + if found { + t.Fatalf("Race detected! Content changed from %q", initialContent) + } +} + +func TestResponseChain_Close_Idempotency(t *testing.T) { + resp := &http.Response{ + StatusCode: 200, + } + rc := NewResponseChain(resp, 1024) + + rc.Close() + + defer func() { + if r := recover(); r != nil { + t.Errorf("Close() panicked on second call: %v", r) + } + }() + + rc.Close() +} + +func TestLimitedBuffer_Pool(t *testing.T) { + buf := getBuffer() + defer putBuffer(buf) + + lb := &limitedBuffer{buf: buf, maxCap: 1024 * 1024} + data := bytes.Repeat([]byte("A"), 100*1024) // 100KB + r := bytes.NewReader(data) + + n, err := lb.ReadFrom(r) + require.NoError(t, err) + require.Equal(t, int64(len(data)), n) + require.Equal(t, len(data), buf.Len()) + require.Equal(t, data, buf.Bytes()) +} From 4a19d93a5fb6e3794aaaee3d542dbac78141c8dc Mon Sep 17 00:00:00 2001 From: Dwi Siswanto Date: Wed, 26 Nov 2025 22:30:18 +0700 Subject: [PATCH 2/2] docs(httputil): update `FullResponse{Bytes,String} doc comments` Signed-off-by: Dwi Siswanto --- http/respChain.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/http/respChain.go b/http/respChain.go index 18be06c..4b07b0d 100644 --- a/http/respChain.go +++ b/http/respChain.go @@ -305,7 +305,7 @@ func (r *ResponseChain) FullResponse() *bytes.Buffer { // FullResponseBytes returns the current response (headers+body) as byte slice. // -// The returned slice is valid only until Close() is called. +// The returned slice is a copy and remains valid even after Close() is called. func (r *ResponseChain) FullResponseBytes() []byte { size := r.headers.Len() + r.body.Len() buf := make([]byte, size) @@ -318,8 +318,8 @@ func (r *ResponseChain) FullResponseBytes() []byte { // FullResponseString returns the current response as string in the chain. // -// The returned string is valid only until Close() is called. -// This is a zero-copy operation for performance. +// The returned string is a copy and remains valid even after Close() is called. +// This is a zero-copy operation from the byte slice. func (r *ResponseChain) FullResponseString() string { return conversion.String(r.FullResponseBytes()) }