Fix potential OOM/panic in Apply

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