Fix buffer leak

trufflesecurity · Oct 18, 2024 · e7d5bbd · e7d5bbd
1 parent 871a2b0
commit e7d5bbd
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 11 deletions.
diff --git a/pkg/handlers/handlers.go b/pkg/handlers/handlers.go
@@ -83,7 +83,7 @@ func newMimeTypeReader(r io.Reader) (mimeTypeReader, error) {
 func newFileReader(r io.Reader) (fileReader, error) {
 	var fReader fileReader
 
-	fReader.BufferedReadSeeker = iobuf.NewBufferedReaderSeeker(r)
+	fReader.BufferedReadSeeker = iobuf.NewBufferedReadSeeker(r)
 
 	mime, err := mimetype.DetectReader(fReader)
 	if err != nil {

diff --git a/pkg/iobuf/bufferedreaderseeker.go b/pkg/iobuf/bufferedreaderseeker.go
@@ -70,10 +70,10 @@ func asSeeker(r io.Reader) io.Seeker {
 	return seeker
 }
 
-// NewBufferedReaderSeeker creates and initializes a BufferedReadSeeker.
+// NewBufferedReadSeeker creates and initializes a BufferedReadSeeker.
 // It takes an io.Reader and checks if it supports seeking.
 // If the reader supports seeking, it is stored in the seeker field.
-func NewBufferedReaderSeeker(r io.Reader) *BufferedReadSeeker {
+func NewBufferedReadSeeker(r io.Reader) *BufferedReadSeeker {
 	const defaultThreshold = 1 << 24 // 16MB threshold for switching to file buffering
 
 	var (
@@ -82,10 +82,6 @@ func NewBufferedReaderSeeker(r io.Reader) *BufferedReadSeeker {
 	)
 
 	seeker = asSeeker(r)
-	if seeker == nil {
-		buf = defaultBufferPool.Get()
-	}
-
 	return &BufferedReadSeeker{
 		reader:    r,
 		seeker:    seeker,
@@ -266,6 +262,10 @@ func (br *BufferedReadSeeker) readToEnd() error {
 }
 
 func (br *BufferedReadSeeker) writeData(data []byte) error {
+	if br.buf == nil {
+		br.buf = br.bufPool.Get()
+	}
+
 	_, err := br.buf.Write(data)
 	if err != nil {
 		return err

diff --git a/pkg/iobuf/bufferedreaderseeker_test.go b/pkg/iobuf/bufferedreaderseeker_test.go
@@ -121,7 +121,7 @@ func TestBufferedReaderSeekerRead(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 
-			brs := NewBufferedReaderSeeker(tt.reader)
+			brs := NewBufferedReadSeeker(tt.reader)
 
 			for i, readSize := range tt.reads {
 				buf := make([]byte, readSize)
@@ -241,7 +241,7 @@ func TestBufferedReaderSeekerSeek(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 
-			brs := NewBufferedReaderSeeker(tt.reader)
+			brs := NewBufferedReadSeeker(tt.reader)
 			pos, err := brs.Seek(tt.offset, tt.whence)
 			if tt.expectedErr {
 				assert.Error(t, err)
@@ -336,7 +336,7 @@ func TestBufferedReaderSeekerReadAt(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 
-			brs := NewBufferedReaderSeeker(tt.reader)
+			brs := NewBufferedReadSeeker(tt.reader)
 
 			out := make([]byte, tt.length)
 			n, err := brs.ReadAt(out, tt.offset)
@@ -436,7 +436,7 @@ func TestBufferedReadSeekerSize(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 
-			brs := NewBufferedReaderSeeker(tt.reader)
+			brs := NewBufferedReadSeeker(tt.reader)
 
 			if tt.setup != nil {
 				tt.setup(brs)