Merge pull request #708 from projectdiscovery/dwisiswant0/fix/httputil/racy-in-ResponseChain

fix(httputil): racy in `ResponseChain`

Author: Mzack9999

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()

http/respChain.go

@@ -305,21 +305,23 @@ 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.
+// The returned slice is a copy and remains valid even after Close() is called.
 func (r *ResponseChain) FullResponseBytes() []byte {
-	buf := r.FullResponse()
-	defer putBuffer(buf)
+	size := r.headers.Len() + r.body.Len()
+	buf := make([]byte, size)
+
+	copy(buf, r.headers.Bytes())
+	copy(buf[r.headers.Len():], r.body.Bytes())
 
-	return buf.Bytes()
+	return buf
 }
 
 // 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.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

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())
+}