From 888f33feb1c37e594c31db0b59817d6618c807ab Mon Sep 17 00:00:00 2001
From: Arpit Jain <arpitjain099@gmail.com>
Date: Tue, 11 Aug 2026 07:31:48 +0900
Subject: [PATCH] fix: stop dropping the bytes between the chunk boundary and
 the read size

ReadChunk consumes up to size+maxPeekSize bytes from the reader, then
generateChunk stops at the first safe "\n\n" boundary it finds past the
chunk size and returns only that prefix. Everything the read consumed after
the boundary is thrown away, and the next call starts from wherever the
reader now sits, so that content is never scanned.

Peek instead of Read and Discard exactly the bytes the chunk covers, which
is what the maxPeekSize name implies. Both call sites already size the
bufio.Reader as size+maxPeekSize, and ErrBufferFull is tolerated so a
smaller reader degrades to a short look-ahead instead of failing.

Signed-off-by: Arpit Jain <arpitjain099@gmail.com>
---
 engine/chunk/chunk.go      | 37 +++++++++++++++++---------------
 engine/chunk/chunk_test.go | 44 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 17 deletions(-)

diff --git a/engine/chunk/chunk.go b/engine/chunk/chunk.go
index 4011d7d3..c705a3bc 100644
--- a/engine/chunk/chunk.go
+++ b/engine/chunk/chunk.go
@@ -7,6 +7,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"io"
 	"sync"
 	"unicode"
 
@@ -134,28 +135,30 @@ func (c *Chunk) GetFileThreshold() int64 {
 
 // ReadChunk reads the next chunk of data from file
 func (c *Chunk) ReadChunk(reader *bufio.Reader, totalLines int) (string, error) {
-	// borrow a []bytes from the pool and seed it with raw data from file (up to chunk size + peek size)
-	rawData, ok := c.GetPeekedBuf()
-	if !ok {
-		return "", fmt.Errorf("expected *bytes.Buffer, got %T", rawData)
-	}
-	defer c.PutPeekedBuf(rawData)
-	n, err := reader.Read(*rawData)
-
-	var chunkStr string
-	// "Callers should always process the n > 0 bytes returned before considering the error err."
-	// https://pkg.go.dev/io#Reader
-	if n > 0 {
-		// only check the filetype at the start of file
-		if totalLines == 0 && ShouldSkipFile((*rawData)[:n]) {
-			return "", fmt.Errorf("skipping file: %w", ErrUnsupportedFileType)
-		}
+	// look ahead without consuming (up to chunk size + peek size): whatever sits past the
+	// chunk boundary has to stay in the reader, otherwise the next call never sees it
+	rawData, err := reader.Peek(c.size + c.maxPeekSize)
+	if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, bufio.ErrBufferFull) {
+		return "", err
+	}
+	if len(rawData) == 0 {
+		return "", io.EOF
+	}
 
-		chunkStr, err = c.generateChunk((*rawData)[:n])
+	// only check the filetype at the start of file
+	if totalLines == 0 && ShouldSkipFile(rawData) {
+		return "", fmt.Errorf("skipping file: %w", ErrUnsupportedFileType)
 	}
+
+	chunkStr, err := c.generateChunk(rawData)
 	if err != nil {
 		return "", err
 	}
+
+	// consume exactly what the chunk covers
+	if _, err := reader.Discard(len(chunkStr)); err != nil {
+		return "", err
+	}
 	return chunkStr, nil
 }
 
diff --git a/engine/chunk/chunk_test.go b/engine/chunk/chunk_test.go
index aaa17e5a..22cffa02 100644
--- a/engine/chunk/chunk_test.go
+++ b/engine/chunk/chunk_test.go
@@ -160,3 +160,47 @@ func TestGenerateChunk(t *testing.T) {
 		})
 	}
 }
+
+func TestReadChunkKeepsEveryByte(t *testing.T) {
+	// Arrange
+	testCases := []struct {
+		name  string
+		input string
+	}{
+		{
+			name:  "boundary inside the peek window",
+			input: "abc\ndef\n\n\n\n\nghi\njkl",
+		},
+		{
+			name:  "boundary right after the chunk size",
+			input: "0123456789\n\nSECRET\n",
+		},
+		{
+			name:  "no boundary at all",
+			input: "abc\ndef\nghi\njkl\nmno\npqr\nstu\nvwx\nyz",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			c := New(WithSize(chunkSize), WithMaxPeekSize(maxPeekSize), WithSmallFileThreshold(smallFileThreshold))
+			reader := bufio.NewReaderSize(strings.NewReader(tc.input), chunkSize+maxPeekSize)
+
+			// Act
+			var scanned strings.Builder
+			totalLines := 0
+			for {
+				chunkStr, err := c.ReadChunk(reader, totalLines)
+				if err == io.EOF {
+					break
+				}
+				require.NoError(t, err)
+				totalLines += strings.Count(chunkStr, "\n")
+				scanned.WriteString(chunkStr)
+			}
+
+			// Assert
+			require.Equal(t, tc.input, scanned.String())
+		})
+	}
+}