use ring buffer

pkg/buffers/buffer/buffer.go

@@ -4,27 +4,62 @@ package buffer
 
 import (
 	"bytes"
+	"io"
 	"time"
+
+	"github.com/trufflesecurity/trufflehog/v3/pkg/buffers/buffer/ring"
 )
 
-// Buffer is a wrapper around bytes.Buffer that includes a timestamp for tracking Buffer checkout duration.
-type Buffer struct {
-	*bytes.Buffer
+type Buffer interface {
+	io.ReadWriter
+	io.WriterTo
+	Len() int
+	Cap() int
+	Reset()
+	Grow(int)
+	Bytes() []byte
+	String() string
+}
+
+// CheckoutBuffer is a wrapper around bytes.Buffer that includes a timestamp for tracking CheckoutBuffer checkout duration.
+type CheckoutBuffer struct {
+	Buffer
 	checkedOutAt time.Time
 }
 
 const defaultBufferSize = 1 << 12 // 4KB
-// NewBuffer creates a new instance of Buffer.
-func NewBuffer() *Buffer { return &Buffer{Buffer: bytes.NewBuffer(make([]byte, 0, defaultBufferSize))} }
 
-func (b *Buffer) Grow(size int) {
+// WithBuffer is a functional option to specify a custom buffer.
+func WithBuffer(b Buffer) func(*CheckoutBuffer) {
+	return func(cb *CheckoutBuffer) {
+		cb.Buffer = b
+	}
+}
+
+// Option is a functional option to configure a CheckoutBuffer.
+type Option func(*CheckoutBuffer)
+
+// NewBuffer creates a new instance of CheckoutBuffer.
+func NewBuffer(opts ...Option) *CheckoutBuffer {
+	buf := new(CheckoutBuffer)
+	for _, opt := range opts {
+		opt(buf)
+	}
+	if buf.Buffer == nil {
+		buf.Buffer = ring.NewBuffer(defaultBufferSize)
+	}
+
+	return buf
+}
+
+func (b *CheckoutBuffer) Grow(size int) {
 	b.Buffer.Grow(size)
 	b.recordGrowth(size)
 }
 
-func (b *Buffer) ResetMetric() { b.checkedOutAt = time.Now() }
+func (b *CheckoutBuffer) ResetMetric() { b.checkedOutAt = time.Now() }
 
-func (b *Buffer) RecordMetric() {
+func (b *CheckoutBuffer) RecordMetric() {
 	dur := time.Since(b.checkedOutAt)
 	checkoutDuration.Observe(float64(dur.Microseconds()))
 	checkoutDurationTotal.Add(float64(dur.Microseconds()))
@@ -32,16 +67,16 @@ func (b *Buffer) RecordMetric() {
 	totalBufferLength.Add(float64(b.Len()))
 }
 
-func (b *Buffer) recordGrowth(size int) {
+func (b *CheckoutBuffer) recordGrowth(size int) {
 	growCount.Inc()
 	growAmount.Add(float64(size))
 }
 
-// Write date to the buffer.
-func (b *Buffer) Write(data []byte) (int, error) {
+// Write date to the Buffer.
+func (b *CheckoutBuffer) Write(data []byte) (int, error) {
 	if b.Buffer == nil {
 		// This case should ideally never occur if buffers are properly managed.
-		b.Buffer = bytes.NewBuffer(make([]byte, 0, defaultBufferSize))
+		b.Buffer = ring.NewBuffer(defaultBufferSize)
 		b.ResetMetric()
 	}
 
@@ -49,11 +84,11 @@ func (b *Buffer) Write(data []byte) (int, error) {
 	bufferLength := b.Buffer.Len()
 	totalSizeNeeded := bufferLength + size
 
-	// If the total size is within the threshold, write to the buffer.
+	// If the total size is within the threshold, write to the Buffer.
 	availableSpace := b.Buffer.Cap() - bufferLength
 	growSize := totalSizeNeeded - bufferLength
 	if growSize > availableSpace {
-		// We are manually growing the buffer so we can track the growth via metrics.
+		// We are manually growing the Buffer so we can track the growth via metrics.
 		// Knowing the exact data size, we directly resize to fit it, rather than exponential growth
 		// which may require multiple allocations and copies if the size required is much larger
 		// than double the capacity. Our approach aligns with default behavior when growth sizes
@@ -65,8 +100,8 @@ func (b *Buffer) Write(data []byte) (int, error) {
 }
 
 // readCloser is a custom implementation of io.ReadCloser. It wraps a bytes.Reader
-// for reading data from an in-memory buffer and includes an onClose callback.
-// The onClose callback is used to return the buffer to the pool, ensuring buffer re-usability.
+// for reading data from an in-memory Buffer and includes an onClose callback.
+// The onClose callback is used to return the Buffer to the pool, ensuring Buffer re-usability.
 type readCloser struct {
 	*bytes.Reader
 	onClose func()
@@ -77,14 +112,14 @@ func ReadCloser(data []byte, onClose func()) *readCloser {
 	return &readCloser{Reader: bytes.NewReader(data), onClose: onClose}
 }
 
-// Close implements the io.Closer interface. It calls the onClose callback to return the buffer
-// to the pool, enabling buffer reuse. This method should be called by the consumers of ReadCloser
+// Close implements the io.Closer interface. It calls the onClose callback to return the Buffer
+// to the pool, enabling Buffer reuse. This method should be called by the consumers of ReadCloser
 // once they have finished reading the data to ensure proper resource management.
 func (brc *readCloser) Close() error {
 	if brc.onClose == nil {
 		return nil
 	}
 
-	brc.onClose() // Return the buffer to the pool
+	brc.onClose() // Return the Buffer to the pool
 	return nil
 }

pkg/buffers/buffer/buffer_test.go

@@ -5,9 +5,11 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
+	"github.com/trufflesecurity/trufflehog/v3/pkg/buffers/buffer/ring"
 )
 
-func TestBufferWrite(t *testing.T) {
+func TestBytesBufferWrite(t *testing.T) {
 	t.Parallel()
 	tests := []struct {
 		name              string
@@ -17,7 +19,7 @@ func TestBufferWrite(t *testing.T) {
 		expectedCap       int
 	}{
 		{
-			name:            "Write to empty buffer",
+			name:            "Write to empty Buffer",
 			initialCapacity: defaultBufferSize,
 			writeDataSequence: [][]byte{
 				[]byte("hello"),
@@ -66,7 +68,81 @@ func TestBufferWrite(t *testing.T) {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
-			buf := &Buffer{Buffer: bytes.NewBuffer(make([]byte, 0, tc.initialCapacity))}
+			buf := &CheckoutBuffer{Buffer: bytes.NewBuffer(make([]byte, 0, tc.initialCapacity))}
+			totalWritten := 0
+			for _, data := range tc.writeDataSequence {
+				n, err := buf.Write(data)
+				assert.NoError(t, err)
+
+				totalWritten += n
+			}
+			assert.Equal(t, tc.expectedSize, totalWritten)
+			assert.Equal(t, tc.expectedSize, buf.Len())
+			assert.GreaterOrEqual(t, buf.Cap(), tc.expectedCap)
+		})
+	}
+}
+
+func TestRingBufferWrite(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name              string
+		initialCapacity   int
+		writeDataSequence [][]byte // Sequence of writes to simulate multiple writes
+		expectedSize      int
+		expectedCap       int
+	}{
+		{
+			name:            "Write to empty Buffer",
+			initialCapacity: defaultBufferSize,
+			writeDataSequence: [][]byte{
+				[]byte("hello"),
+			},
+			expectedSize: 5,
+			expectedCap:  defaultBufferSize, // No growth for small data
+		},
+		{
+			name:            "Write causing growth",
+			initialCapacity: 10, // Small initial capacity to force growth
+			writeDataSequence: [][]byte{
+				[]byte("this is a longer string exceeding initial capacity"),
+			},
+			expectedSize: 50,
+			expectedCap:  50,
+		},
+		{
+			name:              "Write nil data",
+			initialCapacity:   defaultBufferSize,
+			writeDataSequence: [][]byte{nil},
+			expectedCap:       defaultBufferSize,
+		},
+		{
+			name:            "Repeated writes, cumulative growth",
+			initialCapacity: 20, // Set an initial capacity to test growth over multiple writes
+			writeDataSequence: [][]byte{
+				[]byte("first write, "),
+				[]byte("second write, "),
+				[]byte("third write exceeding the initial capacity."),
+			},
+			expectedSize: 70,
+			expectedCap:  70, // Expect capacity to grow to accommodate all writes
+		},
+		{
+			name:            "Write large single data to test significant growth",
+			initialCapacity: 50, // Set an initial capacity smaller than the data to be written
+			writeDataSequence: [][]byte{
+				bytes.Repeat([]byte("a"), 1024), // 1KB data to significantly exceed initial capacity
+			},
+			expectedSize: 1024,
+			expectedCap:  1024,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			buf := &CheckoutBuffer{Buffer: ring.NewBuffer(tc.initialCapacity)}
 			totalWritten := 0
 			for _, data := range tc.writeDataSequence {
 				n, err := buf.Write(data)

pkg/buffers/buffer/metrics.go

@@ -26,14 +26,14 @@ var (
 		Namespace: common.MetricsNamespace,
 		Subsystem: common.MetricsSubsystem,
 		Name:      "checkout_duration_total_us",
-		Help:      "Total duration in microseconds of Buffer checkouts.",
+		Help:      "Total duration in microseconds of CheckoutBuffer checkouts.",
 	})
 
 	checkoutDuration = promauto.NewHistogram(prometheus.HistogramOpts{
 		Namespace: common.MetricsNamespace,
 		Subsystem: common.MetricsSubsystem,
 		Name:      "checkout_duration_us",
-		Help:      "Duration in microseconds of Buffer checkouts.",
+		Help:      "Duration in microseconds of CheckoutBuffer checkouts.",
 		Buckets:   prometheus.ExponentialBuckets(10, 10, 7),
 	})
 

pkg/buffers/buffer/ring/ring.go

@@ -0,0 +1,215 @@
+package ring
+
+import (
+	"io"
+)
+
+// Ring is a ring buffer implementation that implements the io.Writer and io.Reader interfaces.
+type Ring struct {
+	buf    []byte
+	size   int
+	rp     int // Read pointer
+	wp     int // Write pointer
+	isFull bool
+	// availableCap tracks the available capacity in the buffer.
+	// It avoids recalculating the available capacity on every write operation,
+	// which is needed in case the buffer needs to be resized.
+	availableCap int
+}
+
+const defaultBufferSize = 1 << 12 // 4KB
+
+// NewBuffer creates a new ring buffer with the given size.
+func NewBuffer(size int) *Ring {
+	return &Ring{
+		buf:          make([]byte, size),
+		size:         size,
+		availableCap: size,
+	}
+}
+
+// Write writes len(p) bytes from p to the underlying data buffer.
+func (r *Ring) Write(p []byte) (int, error) {
+	if len(p) == 0 {
+		return 0, nil
+	}
+
+	if len(p) > r.availableCap {
+		oldLen := r.Len()
+		newSize := r.size * 2
+		if newSize < oldLen+len(p) {
+			newSize = oldLen + len(p)
+			// Add extra capacity to ensure the buffer is not resized to fit only the new data.
+			newSize += newSize / 2
+		}
+		r.resize(newSize, oldLen) // Ensure there's enough space for new data
+	}
+	n := r.write(p)
+
+	return n, nil
+}
+
+// WriteTo writes the buffer's data to w.
+func (r *Ring) WriteTo(w io.Writer) (int64, error) {
+	if r.isEmpty() {
+		return 0, nil // Nothing to write
+	}
+
+	var totalWritten int64
+
+	// If the data wraps in the buffer, write in two parts: from read pointer to end, and from start to write pointer.
+	if r.wp <= r.rp {
+		// Data does not wrap, can directly write from r to w
+		n, err := w.Write(r.buf[r.rp:r.wp])
+		totalWritten += int64(n)
+		if err != nil {
+			return totalWritten, err
+		}
+		r.rp = 0
+	}
+
+	// Data does not wrap, or we are continuing from a previous write after a wrap.
+	if r.wp > 0 {
+		// Data wraps, write from r to end of buffer.
+		n, err := w.Write(r.buf[r.rp:])
+		totalWritten += int64(n)
+		if err != nil {
+			r.rp = (r.rp + n) % r.size
+			return totalWritten, err
+		}
+	}
+
+	// Update the ring buffer's read pointer and full flag after successful write.
+	// Since all data has been read, we can reset the ring buffer.
+	r.Reset()
+
+	return totalWritten, nil
+}
+
+func (r *Ring) write(p []byte) int {
+	endPos := (r.wp + len(p)) % r.size
+	n := len(p) // The number of bytes to write is determined by the length of p.
+
+	// Execute the copy operation based on the relative positions of wp and rp.
+	if r.wp >= r.rp {
+		// Write wraps the buffer when the buffer is being re-used.
+		if endPos < r.wp {
+			copy(r.buf[r.wp:], p[:r.size-r.wp])
+			copy(r.buf[:endPos], p[r.size-r.wp:])
+		} else {
+			copy(r.buf[r.wp:], p)
+		}
+	} else {
+		copy(r.buf[r.wp:endPos], p)
+	}
+
+	r.wp = endPos
+	// The buffer is full if after the write, the next write position (wp) would overwrite the read position (rp).
+	// Calculate the next write position after this write.
+	nextWp := (r.wp + 1) % r.size
+	// Initially check if the next position after wp is rp, indicating a full buffer on the next write.
+	r.isFull = nextWp == r.rp
+
+	// Exactly full if the next write position is the read position.
+	if r.wp == r.rp && n > 0 {
+		r.isFull = true
+	}
+
+	r.availableCap -= n
+	if r.availableCap < 0 {
+		r.availableCap = 0
+	}
+
+	return n
+}
+
+func (r *Ring) resize(newSize, oldLen int) {
+	newBuf := make([]byte, newSize)
+	r.read(newBuf)
+	r.buf = newBuf
+	r.size = newSize
+	r.rp = 0
+	r.wp = oldLen
+	r.isFull = false
+
+	r.availableCap = newSize - oldLen
+}
+
+// Read reads up to len(p) bytes into p from the underlying data buffer.
+func (r *Ring) Read(p []byte) (int, error) {
+	if r.isEmpty() {
+		return 0, nil
+	}
+
+	var n int
+	if r.wp > r.rp {
+		// Non-wrapped data: directly copy from rp to wp.
+		n = copy(p, r.buf[r.rp:r.wp])
+	} else {
+		// Wrapped data: copy in two steps.
+		// First, from rp to the end of the buffer.
+		n = copy(p, r.buf[r.rp:])
+		// If there's more space in p and data wrapped, copy the remaining data from the beginning.
+		if n < len(p) {
+			n += copy(p[n:], r.buf[:r.wp])
+		}
+	}
+
+	// Update the read pointer and re-calculate the available capacity.
+	r.rp = (r.rp + n) % r.size
+	r.isFull = false // The buffer cannot be full after a read operation.
+
+	r.availableCap = r.size - r.Len()
+
+	return n, nil
+}
+
+func (r *Ring) read(p []byte) int {
+	n := copy(p, r.buf[r.rp:])
+	if r.wp < r.rp {
+		n += copy(p[n:], r.buf[:r.wp])
+	}
+	return n
+}
+
+// Bytes returns a copy of the buffer's data.
+func (r *Ring) Bytes() []byte {
+	if r.isEmpty() {
+		return nil
+	}
+
+	return r.buf[r.rp:r.wp]
+}
+
+// String returns the buffer's data as a string.
+func (r *Ring) String() string { return string(r.Bytes()) }
+
+// isEmpty returns true if the buffer is empty.
+func (r *Ring) isEmpty() bool { return !r.isFull && r.rp == r.wp }
+
+// Len return the length of available read bytes.
+func (r *Ring) Len() int {
+	if r.isFull {
+		return r.size
+	}
+	if r.wp >= r.rp {
+		return r.wp - r.rp
+	}
+	return r.size - r.rp + r.wp
+}
+
+// Cap returns the size of the underlying buffer.
+func (r *Ring) Cap() int {
+	return r.size
+}
+
+// Reset resets the buffer.
+func (r *Ring) Reset() {
+	r.rp = 0
+	r.wp = 0
+	r.isFull = false
+	r.availableCap = r.size
+}
+
+// Grow is a no-op for ring buffer.
+func (*Ring) Grow(_ int) {}

pkg/buffers/buffer/ring/ring_test.go

@@ -0,0 +1,331 @@
+package ring
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRingWrite(t *testing.T) {
+	type writeTest struct {
+		name        string
+		initBufSize int    // Initial buffer size
+		initData    []byte // Initial data to preload the buffer
+		writeData   []byte // Data to write in the test
+		wantWritten int    // Expected number of bytes written
+		wantBufSize int    // Expected buffer size after write
+		wantBufData []byte // Expected buffer data after write
+	}
+
+	testCases := []writeTest{
+		{
+			name:        "WriteWithEmptyBuffer",
+			initBufSize: 10,
+			initData:    nil,
+			writeData:   []byte("hello"),
+			wantWritten: 5,
+			wantBufSize: 10,
+			wantBufData: []byte("hello"),
+		},
+		{
+			name:        "WriteWithExactFillNoInitialData",
+			initBufSize: 5,
+			initData:    nil,
+			writeData:   []byte("hello"),
+			wantWritten: 5,
+			wantBufSize: 5,
+			wantBufData: []byte("hello"),
+		},
+		{
+			name:        "WriteWithPartialFill",
+			initBufSize: 10,
+			initData:    []byte("hi"),
+			writeData:   []byte("there"),
+			wantWritten: 5,
+			wantBufSize: 10,
+			wantBufData: []byte("hithere"),
+		},
+		{
+			name:        "WriteWithExactFill",
+			initBufSize: 5,
+			initData:    []byte("hi"),
+			writeData:   []byte("123"),
+			wantWritten: 3,
+			wantBufSize: 5,
+			wantBufData: []byte("hi123"),
+		},
+		{
+			name:        "WriteWithOverflowRequiresResize",
+			initBufSize: 5,
+			initData:    []byte("hi"),
+			writeData:   []byte("12345"),
+			wantWritten: 5,
+			wantBufSize: 10,                            // Expected double the size of the initial buffer
+			wantBufData: []byte("hi12345\x00\x00\x00"), // Including the extra null bytes
+		},
+		{
+			name:        "WriteWithZeroLengthData",
+			initBufSize: 10,
+			initData:    []byte("data"),
+			writeData:   nil,
+			wantWritten: 0,
+			wantBufSize: 10,
+			wantBufData: []byte("data"),
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			r := NewBuffer(tc.initBufSize)
+			if len(tc.initData) > 0 {
+				_, _ = r.Write(tc.initData)
+			}
+
+			gotWritten, err := r.Write(tc.writeData)
+			assert.NoError(t, err)
+
+			assert.Equal(t, tc.wantWritten, gotWritten)
+			assert.Equal(t, tc.wantBufSize, len(r.buf))
+			if tc.wantWritten >= tc.initBufSize {
+				assert.Equal(t, tc.wantBufData, r.buf)
+			} else {
+				assert.Equal(t, tc.wantBufData, r.buf[:tc.wantWritten+len(tc.initData)])
+			}
+		})
+	}
+}
+
+func BenchmarkBufferWrite(b *testing.B) {
+	type benchCase struct {
+		name     string
+		dataSize int // Size of the data to write in bytes
+	}
+
+	benchmarks := []benchCase{
+		{"1KB", 1 << 10},     // 1KB
+		{"4KB", 4 << 10},     // 4KB
+		{"16KB", 16 << 10},   // 16KB
+		{"64KB", 64 << 10},   // 64KB
+		{"256KB", 256 << 10}, // 256KB
+		{"1MB", 1 << 20},     // 1MB
+		{"4MB", 4 << 20},     // 4MB
+		{"16MB", 16 << 20},   // 16MB
+		{"64MB", 64 << 20},   // 64MB
+	}
+
+	for _, bc := range benchmarks {
+		bc := bc
+		b.Run(bc.name, func(b *testing.B) {
+			data := generateData(bc.dataSize) // Generate pseudo-random data for this benchmark case
+			r := NewBuffer(defaultBufferSize)
+
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				_, err := r.Write(data)
+				assert.NoError(b, err)
+				r.Reset()
+			}
+		})
+	}
+}
+
+func generateData(size int) []byte {
+	rand.Seed(42)
+	data := make([]byte, size)
+	for i := range data {
+		data[i] = byte(rand.Intn(256))
+	}
+	return data
+}
+
+func TestRingBufferBytes(t *testing.T) {
+	tests := []struct {
+		name        string
+		initBufSize int
+		initData    []byte
+		wantBytes   []byte
+	}{
+		{
+			name:        "EmptyBuffer",
+			initBufSize: 10,
+		},
+		{
+			name:        "PartialFill",
+			initBufSize: 10,
+			initData:    []byte("hello"),
+			wantBytes:   []byte("hello"),
+		},
+		{
+			name:        "FullFill",
+			initBufSize: 5,
+			initData:    []byte("hello"),
+			wantBytes:   []byte("hello"),
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			r := NewBuffer(tc.initBufSize)
+			if tc.initData != nil {
+				_, _ = r.Write(tc.initData)
+			}
+
+			assert.Equal(t, tc.wantBytes, r.Bytes())
+		})
+	}
+}
+
+func TestRingBufferLen(t *testing.T) {
+	tests := []struct {
+		name        string
+		initBufSize int
+		initData    []byte
+		wantLen     int
+	}{
+		{
+			name:        "EmptyBuffer",
+			initBufSize: 10,
+		},
+		{
+			name:        "PartialFill",
+			initBufSize: 10,
+			initData:    []byte("hello"),
+			wantLen:     5,
+		},
+		{
+			name:        "FullFill",
+			initBufSize: 5,
+			initData:    []byte("hello"),
+			wantLen:     5,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			r := NewBuffer(tc.initBufSize)
+			if tc.initData != nil {
+				_, _ = r.Write(tc.initData)
+			}
+
+			assert.Equal(t, tc.wantLen, r.Len())
+		})
+	}
+}
+
+func TestRingBufferCap(t *testing.T) {
+	tests := []struct {
+		name        string
+		initBufSize int
+		wantCap     int
+	}{
+		{
+			name:        "InitialCapacity",
+			initBufSize: 10,
+			wantCap:     10,
+		},
+		{
+			name: "ZeroCapacity",
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			r := NewBuffer(tc.initBufSize)
+
+			assert.Equal(t, tc.wantCap, r.Cap())
+		})
+	}
+}
+
+func TestRingBufferReset(t *testing.T) {
+	tests := []struct {
+		name        string
+		initBufSize int
+		initData    []byte
+	}{
+		{
+			name:        "AfterWrite",
+			initBufSize: 10,
+			initData:    []byte("reset"),
+		},
+		{
+			name:        "EmptyBuffer",
+			initBufSize: 10,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			r := NewBuffer(tc.initBufSize)
+			if tc.initData != nil {
+				_, _ = r.Write(tc.initData)
+			}
+
+			r.Reset()
+			assert.Equal(t, 0, r.Len())
+			assert.Equal(t, tc.initBufSize, r.Cap())
+			assert.True(t, r.isEmpty())
+		})
+	}
+}
+
+func TestRingBufferRead(t *testing.T) {
+	tests := []struct {
+		name         string
+		initBufSize  int
+		writeData    []byte
+		readLen      int
+		wantReadData []byte
+	}{
+		{
+			name:         "ReadPartialData",
+			initBufSize:  10,
+			writeData:    []byte("hello"),
+			readLen:      5,
+			wantReadData: []byte("hello"),
+		},
+		{
+			name:         "ReadFullData",
+			initBufSize:  5,
+			writeData:    []byte("hello"),
+			readLen:      5,
+			wantReadData: []byte("hello"),
+		},
+		{
+			name:         "ReadEmptyBuffer",
+			initBufSize:  10,
+			writeData:    nil,
+			readLen:      5,
+			wantReadData: []byte{},
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			r := NewBuffer(tc.initBufSize)
+			if tc.writeData != nil {
+				_, _ = r.Write(tc.writeData)
+			}
+
+			readData := make([]byte, tc.readLen)
+			n, err := r.Read(readData)
+			assert.NoError(t, err)
+
+			assert.Equal(t, tc.wantReadData, readData[:n])
+		})
+	}
+}

pkg/buffers/pool/pool.go

@@ -1,7 +1,6 @@
 package pool
 
 import (
-	"bytes"
 	"sync"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/buffers/buffer"
@@ -20,18 +19,14 @@ func (poolMetrics) recordBufferRetrival() {
 	bufferCount.Inc()
 }
 
-func (poolMetrics) recordBufferReturn(buf *buffer.Buffer) {
+func (poolMetrics) recordBufferReturn(buf *buffer.CheckoutBuffer) {
 	activeBufferCount.Dec()
 	buf.RecordMetric()
 }
 
-// Opts is a function that configures a BufferPool.
-type Opts func(pool *Pool)
-
 // Pool of buffers.
 type Pool struct {
 	*sync.Pool
-	bufferSize uint32
 
 	metrics poolMetrics
 }
@@ -39,26 +34,22 @@ type Pool struct {
 // const defaultBufferSize = 1 << 12 // 4KB
 const defaultBufferSize = 1 << 16 // 64KB
 // NewBufferPool creates a new instance of BufferPool.
-func NewBufferPool(opts ...Opts) *Pool {
-	pool := &Pool{bufferSize: defaultBufferSize}
+func NewBufferPool(opts ...buffer.Option) *Pool {
+	pool := new(Pool)
 
-	for _, opt := range opts {
-		opt(pool)
-	}
+	buf := buffer.NewBuffer(opts...)
 	pool.Pool = &sync.Pool{
-		New: func() any {
-			return &buffer.Buffer{Buffer: bytes.NewBuffer(make([]byte, 0, pool.bufferSize))}
-		},
+		New: func() any { return buf },
 	}
 
 	return pool
 }
 
-// Get returns a Buffer from the pool.
-func (p *Pool) Get() *buffer.Buffer {
-	buf, ok := p.Pool.Get().(*buffer.Buffer)
+// Get returns a CheckoutBuffer from the pool.
+func (p *Pool) Get() *buffer.CheckoutBuffer {
+	buf, ok := p.Pool.Get().(*buffer.CheckoutBuffer)
 	if !ok {
-		buf = &buffer.Buffer{Buffer: bytes.NewBuffer(make([]byte, 0, p.bufferSize))}
+		buf = buffer.NewBuffer()
 	}
 	p.metrics.recordBufferRetrival()
 	buf.ResetMetric()
@@ -66,18 +57,18 @@ func (p *Pool) Get() *buffer.Buffer {
 	return buf
 }
 
-// Put returns a Buffer to the pool.
-func (p *Pool) Put(buf *buffer.Buffer) {
+// Put returns a CheckoutBuffer to the pool.
+func (p *Pool) Put(buf *buffer.CheckoutBuffer) {
 	p.metrics.recordBufferReturn(buf)
 
-	// If the Buffer is more than twice the default size, replace it with a new Buffer.
+	// If the CheckoutBuffer is more than twice the default size, replace it with a new CheckoutBuffer.
 	// This prevents us from returning very large buffers to the pool.
 	const maxAllowedCapacity = 2 * defaultBufferSize
 	if buf.Cap() > maxAllowedCapacity {
 		p.metrics.recordShrink(buf.Cap() - defaultBufferSize)
-		buf = &buffer.Buffer{Buffer: bytes.NewBuffer(make([]byte, 0, p.bufferSize))}
+		buf = buffer.NewBuffer()
 	} else {
-		// Reset the Buffer to clear any existing data.
+		// Reset the CheckoutBuffer to clear any existing data.
 		buf.Reset()
 	}
 

pkg/buffers/pool/pool_test.go

@@ -38,13 +38,13 @@ func TestBufferPoolGetPut(t *testing.T) {
 	t.Parallel()
 	tests := []struct {
 		name              string
-		preparePool       func(p *Pool) *buffer.Buffer // Prepare the pool and return an initial buffer to put if needed
-		expectedCapBefore int                          // Expected capacity before putting it back
-		expectedCapAfter  int                          // Expected capacity after retrieving it again
+		preparePool       func(p *Pool) *buffer.CheckoutBuffer // Prepare the pool and return an initial buffer to put if needed
+		expectedCapBefore int                                  // Expected capacity before putting it back
+		expectedCapAfter  int                                  // Expected capacity after retrieving it again
 	}{
 		{
 			name: "Get new buffer and put back without modification",
-			preparePool: func(_ *Pool) *buffer.Buffer {
+			preparePool: func(_ *Pool) *buffer.CheckoutBuffer {
 				return nil // No initial buffer to put
 			},
 			expectedCapBefore: int(defaultBufferSize),
@@ -52,8 +52,8 @@ func TestBufferPoolGetPut(t *testing.T) {
 		},
 		{
 			name: "Put oversized buffer, expect shrink",
-			preparePool: func(p *Pool) *buffer.Buffer {
-				buf := &buffer.Buffer{Buffer: bytes.NewBuffer(make([]byte, 0, 3*defaultBufferSize))}
+			preparePool: func(p *Pool) *buffer.CheckoutBuffer {
+				buf := &buffer.CheckoutBuffer{Buffer: bytes.NewBuffer(make([]byte, 0, 3*defaultBufferSize))}
 				return buf
 			},
 			expectedCapBefore: int(defaultBufferSize),

pkg/handlers/handlers_test.go

@@ -1,7 +1,6 @@
 package handlers
 
 import (
-	"io"
 	"net/http"
 	"os"
 	"strings"
@@ -40,12 +39,10 @@ func TestHandleFile(t *testing.T) {
 }
 
 func BenchmarkHandleFile(b *testing.B) {
-	file, err := os.Open("testdata/test.tgz")
-	assert.Nil(b, err)
-
-	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		sourceChan := make(chan *sources.Chunk, 1)
+		file, err := os.Open("testdata/test.tgz")
+		assert.Nil(b, err)
 
 		b.StartTimer()
 		go func() {
@@ -57,9 +54,6 @@ func BenchmarkHandleFile(b *testing.B) {
 		for range sourceChan {
 		}
 		b.StopTimer()
-
-		_, err = file.Seek(0, io.SeekStart)
-		assert.NoError(b, err)
 	}
 }
 

pkg/writers/buffer_writer/bufferwriter.go

@@ -20,7 +20,7 @@ func (metrics) recordDataProcessed(size int64, dur time.Duration) {
 
 func init() { bufferPool = pool.NewBufferPool() }
 
-// bufferPool is the shared Buffer pool used by all BufferedFileWriters.
+// bufferPool is the shared CheckoutBuffer pool used by all BufferedFileWriters.
 // This allows for efficient reuse of buffers across multiple writers.
 var bufferPool *pool.Pool
 
@@ -36,10 +36,10 @@ const (
 
 // BufferWriter implements contentWriter, using a shared buffer pool for memory management.
 type BufferWriter struct {
-	buf     *buffer.Buffer // The current buffer in use.
-	bufPool *pool.Pool     // The buffer pool used to manage the buffer.
-	size    int            // The total size of the content written to the buffer.
-	state   state          // The current state of the buffer.
+	buf     *buffer.CheckoutBuffer // The current buffer in use.
+	bufPool *pool.Pool             // The buffer pool used to manage the buffer.
+	size    int                    // The total size of the content written to the buffer.
+	state   state                  // The current state of the buffer.
 
 	metrics metrics
 }
@@ -62,6 +62,7 @@ func NewFromReader(ctx context.Context, r io.Reader) (*BufferWriter, error) {
 	return buf, nil
 }
 
+const defaultBufferSize = 1 << 12 // 4KB
 // Write delegates the writing operation to the underlying bytes.Buffer.
 func (b *BufferWriter) Write(data []byte) (int, error) {
 	if b.state != writeOnly {

pkg/writers/buffered_file_writer/bufferedfilewriter.go

@@ -30,7 +30,7 @@ func (bufferedFileWriterMetrics) recordDiskWrite(size int64) {
 
 func init() { bufferPool = pool.NewBufferPool() }
 
-// bufferPool is the shared Buffer pool used by all BufferedFileWriters.
+// bufferPool is the shared CheckoutBuffer pool used by all BufferedFileWriters.
 // This allows for efficient reuse of buffers across multiple writers.
 var bufferPool *pool.Pool
 
@@ -50,10 +50,10 @@ type BufferedFileWriter struct {
 	threshold uint64 // Threshold for switching to file writing.
 	size      uint64 // Total size of the data written.
 
-	bufPool  *pool.Pool     // Pool for storing buffers for reuse.
-	buf      *buffer.Buffer // Buffer for storing data under the threshold in memory.
-	filename string         // Name of the temporary file.
-	file     *os.File       // File for storing data over the threshold.
+	bufPool  *pool.Pool             // Pool for storing buffers for reuse.
+	buf      *buffer.CheckoutBuffer // CheckoutBuffer for storing data under the threshold in memory.
+	filename string                 // Name of the temporary file.
+	file     *os.File               // File for storing data over the threshold.
 
 	state state // Current state of the writer. (writeOnly or readOnly)