Fix potential OOM in Apply (#75)

The apply logic trusted the line counts of the incoming patch when
allocating memory for preimages. This meant a patch could specify a very
large old position in a fragment and force the library to allocate
excessive memory before checking if the source actually contains that
many lines.

To fix this while keeping the current structure that reads the source
incrementally, read the preimage in chunks of 4096 lines. This bounds
the memory that might be allocated before we reach the end of the source
file. Most practical patches should only require one or two reads.

Also fix a related panic that could happen with an old position close to
math.MaxInt64.

Author: bluekeyes

gitdiff/apply_test.go

@@ -40,6 +40,20 @@ func TestApplyTextFragment(t *testing.T) {
 			},
 			Err: &Conflict{},
 		},
+		"errorShortSrcExtreme": {
+			Files: applyFiles{
+				Src:   "text_fragment_error.src",
+				Patch: "text_fragment_error_short_src_extreme.patch",
+			},
+			Err: &Conflict{},
+		},
+		"errorOverflow": {
+			Files: applyFiles{
+				Src:   "text_fragment_error.src",
+				Patch: "text_fragment_error_overflow.patch",
+			},
+			Err: "overflow",
+		},
 		"errorContextConflict": {
 			Files: applyFiles{
 				Src:   "text_fragment_error.src",

gitdiff/apply_text.go

@@ -3,6 +3,7 @@ package gitdiff
 import (
 	"errors"
 	"io"
+	"math"
 )
 
 // TextApplier applies changes described in text fragments to source data. If
@@ -66,6 +67,9 @@ func (a *TextApplier) ApplyFragment(f *TextFragment) error {
 	if fragStart < 0 {
 		fragStart = 0
 	}
+	if f.OldLines > math.MaxInt64-fragStart {
+		return applyError(errors.New("fragment bounds overflow"))
+	}
 	fragEnd := fragStart + f.OldLines
 
 	start := a.nextLine
@@ -83,15 +87,9 @@ func (a *TextApplier) ApplyFragment(f *TextFragment) error {
 		}
 	}
 
-	preimage := make([][]byte, fragEnd-start)
-	n, err := a.lineSrc.ReadLinesAt(preimage, start)
+	preimage, err := readPreimage(a.lineSrc, start, fragEnd-start)
 	if err != nil {
-		// an EOF indicates that source file is shorter than the patch expects,
-		// which should be reported as a conflict rather than a generic error
-		if errors.Is(err, io.EOF) {
-			err = &Conflict{"src has fewer lines than required by fragment"}
-		}
-		return applyError(err, lineNum(start+int64(n)))
+		return applyError(err)
 	}
 
 	// copy leading data before the fragment starts
@@ -131,6 +129,43 @@ func (a *TextApplier) ApplyFragment(f *TextFragment) error {
 	return nil
 }
 
+// readPreimage attempts to read lines from the reader in chunks to avoid
+// allocating too much memory if the expected line count is longer than the
+// actual input.
+func readPreimage(r LineReaderAt, start int64, lines int64) ([][]byte, error) {
+	// This chunk size is arbitrary, but is large enough that most preimages
+	// should read in a single chunk. It's generally safe to pick a large chunk
+	// size, as the chunk only allocates slice headers for the line content,
+	// with the actual content only allocated if it exists in the source. With
+	// a chunk size of 4096, we allocate at most ~96KB extra before detecting
+	// the short source in the worst case.
+	const chunkSize = 4096
+
+	chunks := ((lines - 1) / chunkSize) + 1
+	remaining := lines
+
+	var preimage [][]byte
+	for c := int64(0); c < chunks; c++ {
+		readSize := min(chunkSize, remaining)
+
+		i := int64(len(preimage))
+		preimage = append(preimage, make([][]byte, readSize)...)
+
+		n, err := r.ReadLinesAt(preimage[i:i+readSize], start)
+		if err != nil {
+			// an EOF indicates that source file is shorter than the patch expects,
+			// which should be reported as a conflict rather than a generic error
+			if errors.Is(err, io.EOF) {
+				err = &Conflict{"src has fewer lines than required by fragment"}
+			}
+			return nil, applyError(err, lineNum(start+int64(n)))
+		}
+		start += int64(n)
+		remaining -= int64(n)
+	}
+	return preimage, nil
+}
+
 func applyTextLine(dst io.Writer, line Line, preimage [][]byte, i int64) (err error) {
 	if line.Old() && string(preimage[i]) != line.Line {
 		return &Conflict{"fragment line does not match src line"}

gitdiff/assert_test.go

@@ -7,6 +7,8 @@ import (
 )
 
 func assertError(t *testing.T, expected any, actual error, action string) {
+	t.Helper()
+
 	if actual == nil {
 		t.Fatalf("expected error %s, but got nil", action)
 	}

gitdiff/testdata/apply/text_fragment_error_overflow.patch

@@ -0,0 +1,12 @@
+diff --git a/gitdiff/testdata/apply/text_fragment_error.src b/gitdiff/testdata/apply/text_fragment_error.src
+--- a/gitdiff/testdata/apply/text_fragment_error.src
++++ b/gitdiff/testdata/apply/text_fragment_error.src
+@@ -9223372036854775803,7 +9223372036854775803,7 @@ line 14
+ line 15
+ line 16
+ line 17
+-line 18
++new line a
+ line 19
+ line 20
+ line 21

gitdiff/testdata/apply/text_fragment_error_short_src_extreme.patch

@@ -0,0 +1,12 @@
+diff --git a/gitdiff/testdata/apply/text_fragment_error.src b/gitdiff/testdata/apply/text_fragment_error.src
+--- a/gitdiff/testdata/apply/text_fragment_error.src
++++ b/gitdiff/testdata/apply/text_fragment_error.src
+@@ -1152921504606846976,7 +1152921504606846976,7 @@ line 14
+ line 15
+ line 16
+ line 17
+-line 18
++new line a
+ line 19
+ line 20
+ line 21