[bug]- Invalid Seek for Non-Seekable Readers (#3095)

* inital work * fix and add tests * uncomment * fix seek end * use buffer pool * revert timeout * make linter happy * More linting :()
trufflesecurity · Jul 25, 2024 · ebfbd21 · ebfbd21
1 parent 4a8b213
commit ebfbd21
Show file tree

Hide file tree

Showing 9 changed files with 337 additions and 127 deletions.
diff --git a/pkg/handlers/ar_test.go b/pkg/handlers/ar_test.go
@@ -20,6 +20,7 @@ func TestHandleARFile(t *testing.T) {
 
 	rdr, err := newFileReader(file)
 	assert.NoError(t, err)
+	defer rdr.Close()
 
 	handler := newARHandler()
 	archiveChan, err := handler.HandleFile(context.AddLogger(ctx), rdr)

diff --git a/pkg/handlers/archive.go b/pkg/handlers/archive.go
@@ -111,6 +111,7 @@ func (h *archiveHandler) openArchive(ctx logContext.Context, depth int, reader f
 			}
 			return fmt.Errorf("error creating custom reader: %w", err)
 		}
+		defer rdr.Close()
 
 		return h.openArchive(ctx, depth+1, rdr, archiveChan)
 	case archiver.Extractor:
@@ -194,6 +195,7 @@ func (h *archiveHandler) extractorHandler(archiveChan chan []byte) func(context.
 			}
 			return fmt.Errorf("error creating custom reader: %w", err)
 		}
+		defer rdr.Close()
 
 		h.metrics.incFilesProcessed()
 		h.metrics.observeFileSize(fileSize)

diff --git a/pkg/handlers/archive_test.go b/pkg/handlers/archive_test.go
@@ -89,6 +89,8 @@ func TestArchiveHandler(t *testing.T) {
 			if err != nil {
 				t.Errorf("error creating reusable reader: %s", err)
 			}
+			defer newReader.Close()
+
 			archiveChan, err := handler.HandleFile(logContext.Background(), newReader)
 			if testCase.expectErr {
 				assert.NoError(t, err)
@@ -119,6 +121,7 @@ func TestOpenInvalidArchive(t *testing.T) {
 
 	rdr, err := newFileReader(io.NopCloser(reader))
 	assert.NoError(t, err)
+	defer rdr.Close()
 
 	archiveChan := make(chan []byte)
 

diff --git a/pkg/handlers/default_test.go b/pkg/handlers/default_test.go
@@ -20,6 +20,7 @@ func TestHandleNonArchiveFile(t *testing.T) {
 
 	rdr, err := newFileReader(file)
 	assert.NoError(t, err)
+	defer rdr.Close()
 
 	handler := newDefaultHandler(defaultHandlerType)
 	archiveChan, err := handler.HandleFile(context.AddLogger(ctx), rdr)

diff --git a/pkg/handlers/handlers.go b/pkg/handlers/handlers.go
@@ -59,7 +59,7 @@ func newMimeTypeReaderFromFileReader(r fileReader) mimeTypeReader {
 
 // newMimeTypeReader creates a new mimeTypeReader from an io.Reader.
 // It uses a bufio.Reader to perform MIME type detection on the input reader
-// without consuming it, by peeking into the first 512 bytes of the input.
+// without consuming it, by peeking into the first 3072 bytes of the input.
 // This encapsulates both the original reader and the detected MIME type information.
 // This function is particularly useful for specialized archive handlers
 // that need to pass extracted content to the default handler without modifying the original reader.
@@ -84,10 +84,6 @@ func newFileReader(r io.Reader) (fileReader, error) {
 
 	fReader.BufferedReadSeeker = iobuf.NewBufferedReaderSeeker(r)
 
-	// Disable buffering after initial reads.
-	// This optimization ensures we don't continue writing to the buffer after the initial reads.
-	defer fReader.DisableBuffering()
-
 	mime, err := mimetype.DetectReader(fReader)
 	if err != nil {
 		return fReader, fmt.Errorf("unable to detect MIME type: %w", err)
@@ -281,6 +277,7 @@ func HandleFile(
 		}
 		return fmt.Errorf("error creating custom reader: %w", err)
 	}
+	defer rdr.Close()
 
 	mimeT := mimeType(rdr.mime.String())
 	config := newFileHandlingConfig(options...)

diff --git a/pkg/handlers/handlers_test.go b/pkg/handlers/handlers_test.go
@@ -27,21 +27,101 @@ func TestHandleFileCancelledContext(t *testing.T) {
 }
 
 func TestHandleFile(t *testing.T) {
-	reporter := sources.ChanReporter{Ch: make(chan *sources.Chunk, 2)}
+	reporter := sources.ChanReporter{Ch: make(chan *sources.Chunk, 513)}
 
 	// Only one chunk is sent on the channel.
 	// TODO: Embed a zip without making an HTTP request.
 	resp, err := http.Get("https://raw.githubusercontent.com/bill-rich/bad-secrets/master/aws-canary-creds.zip")
 	assert.NoError(t, err)
-	if resp != nil && resp.Body != nil {
-		defer resp.Body.Close()
-	}
+	defer func() {
+		if resp != nil && resp.Body != nil {
+			resp.Body.Close()
+		}
+	}()
 
 	assert.Equal(t, 0, len(reporter.Ch))
 	assert.NoError(t, HandleFile(context.Background(), resp.Body, &sources.Chunk{}, reporter))
 	assert.Equal(t, 1, len(reporter.Ch))
 }
 
+func TestHandleHTTPJson(t *testing.T) {
+	resp, err := http.Get("https://raw.githubusercontent.com/ahrav/nothing-to-see-here/main/sm_random_data.json")
+	assert.NoError(t, err)
+	defer func() {
+		if resp != nil && resp.Body != nil {
+			resp.Body.Close()
+		}
+	}()
+
+	chunkCh := make(chan *sources.Chunk, 1)
+	go func() {
+		defer close(chunkCh)
+		err := HandleFile(logContext.Background(), resp.Body, &sources.Chunk{}, sources.ChanReporter{Ch: chunkCh})
+		assert.NoError(t, err)
+	}()
+
+	wantCount := 513
+	count := 0
+	for range chunkCh {
+		count++
+	}
+	assert.Equal(t, wantCount, count)
+}
+
+func TestHandleHTTPJsonZip(t *testing.T) {
+	resp, err := http.Get("https://raw.githubusercontent.com/ahrav/nothing-to-see-here/main/sm.zip")
+	assert.NoError(t, err)
+	defer func() {
+		if resp != nil && resp.Body != nil {
+			resp.Body.Close()
+		}
+	}()
+
+	chunkCh := make(chan *sources.Chunk, 1)
+	go func() {
+		defer close(chunkCh)
+		err := HandleFile(logContext.Background(), resp.Body, &sources.Chunk{}, sources.ChanReporter{Ch: chunkCh})
+		assert.NoError(t, err)
+	}()
+
+	wantCount := 513
+	count := 0
+	for range chunkCh {
+		count++
+	}
+	assert.Equal(t, wantCount, count)
+}
+
+func BenchmarkHandleHTTPJsonZip(b *testing.B) {
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		func() {
+			resp, err := http.Get("https://raw.githubusercontent.com/ahrav/nothing-to-see-here/main/sm.zip")
+			assert.NoError(b, err)
+
+			defer func() {
+				if resp != nil && resp.Body != nil {
+					resp.Body.Close()
+				}
+			}()
+
+			chunkCh := make(chan *sources.Chunk, 1)
+
+			b.StartTimer()
+			go func() {
+				defer close(chunkCh)
+				err := HandleFile(logContext.Background(), resp.Body, &sources.Chunk{}, sources.ChanReporter{Ch: chunkCh})
+				assert.NoError(b, err)
+			}()
+
+			for range chunkCh {
+			}
+
+			b.StopTimer()
+		}()
+	}
+}
+
 func BenchmarkHandleFile(b *testing.B) {
 	file, err := os.Open("testdata/test.tgz")
 	assert.Nil(b, err)

diff --git a/pkg/handlers/rpm_test.go b/pkg/handlers/rpm_test.go
@@ -20,6 +20,7 @@ func TestHandleRPMFile(t *testing.T) {
 
 	rdr, err := newFileReader(file)
 	assert.NoError(t, err)
+	defer rdr.Close()
 
 	handler := newRPMHandler()
 	archiveChan, err := handler.HandleFile(context.AddLogger(ctx), rdr)