caddyfile: avoid quadratic heredoc marker scanning (#8203)

(cherry picked from commit 7328273d6f)

caddyconfig/caddyfile/lexer.go

@@ -108,6 +108,8 @@ func (l *lexer) next() (bool, error) {
 	var val []rune
 	var comment, quoted, btQuoted, inHeredoc, heredocEscaped, escaped bool
 	var heredocMarker string
+	var heredocPrefix []int
+	var heredocMatched int
 
 	makeToken := func(quoted rune) bool {
 		l.token.Text = string(val)
@@ -168,6 +170,20 @@ func (l *lexer) next() (bool, error) {
 					return false, fmt.Errorf("heredoc marker on line #%d must contain only alphanumeric characters, dashes and underscores; got '%s'", l.line, heredocMarker)
 				}
 
+				// Build a prefix table for incremental marker matching (KMP).
+				// Comparing the whole suffix after every rune costs O(body * marker)
+				// and repeatedly converts large rune slices to strings.
+				heredocPrefix = make([]int, len(heredocMarker))
+				for i, matched := 1, 0; i < len(heredocMarker); i++ {
+					for matched > 0 && heredocMarker[i] != heredocMarker[matched] {
+						matched = heredocPrefix[matched-1]
+					}
+					if heredocMarker[i] == heredocMarker[matched] {
+						matched++
+					}
+					heredocPrefix[i] = matched
+				}
+
 				inHeredoc = true
 				l.skippedLines++
 				val = nil
@@ -185,8 +201,16 @@ func (l *lexer) next() (bool, error) {
 				l.skippedLines++
 			}
 
-			// check if we're done, i.e. that the last few characters are the marker
-			if len(val) >= len(heredocMarker) && heredocMarker == string(val[len(val)-len(heredocMarker):]) {
+			// Markers contain only ASCII, so compare each rune directly with
+			// the next marker byte. Prefix fallback also handles overlapping
+			// partial matches without rescanning the body.
+			for heredocMatched > 0 && ch != rune(heredocMarker[heredocMatched]) {
+				heredocMatched = heredocPrefix[heredocMatched-1]
+			}
+			if ch == rune(heredocMarker[heredocMatched]) {
+				heredocMatched++
+			}
+			if heredocMatched == len(heredocMarker) {
 				// set the final value
 				val, err = l.finalizeHeredoc(val, heredocMarker)
 				if err != nil {

caddyconfig/caddyfile/lexer_test.go

@@ -15,6 +15,7 @@
 package caddyfile
 
 import (
+	"strings"
 	"testing"
 )
 
@@ -539,3 +540,49 @@ func lexerCompare(t *testing.T, n int, expected, actual []Token) {
 		}
 	}
 }
+
+// TestLexerLongHeredoc exercises long and overlapping marker prefixes. Checking
+// the entire marker after every rune made these inputs take seconds to tokenize.
+func TestLexerLongHeredoc(t *testing.T) {
+	for _, tc := range []struct {
+		name, marker, body string
+		incomplete         bool
+	}{
+		{name: "long marker", marker: strings.Repeat("a", 8000) + "b", body: strings.Repeat("a", 500000)},
+		{name: "overlapping prefixes", marker: "ababaca", body: "ababababac"},
+		{name: "unicode body", marker: "END", body: "你好\xffEN"},
+		{name: "incomplete", marker: strings.Repeat("a", 8000) + "b", body: strings.Repeat("a", 500000), incomplete: true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			input := "<<" + tc.marker + "\n" + tc.body
+			if !tc.incomplete {
+				input += "\n" + tc.marker
+			}
+			tokens, err := Tokenize([]byte(input), "Caddyfile")
+			if tc.incomplete {
+				if err == nil {
+					t.Fatal("expected incomplete heredoc error")
+				}
+				return
+			}
+			if err != nil {
+				t.Fatal(err)
+			}
+			if len(tokens) != 1 || tokens[0].Text != string([]rune(tc.body)) {
+				t.Fatalf("unexpected heredoc contents: %v", tokens)
+			}
+		})
+	}
+}
+
+func BenchmarkLexerLongHeredoc(b *testing.B) {
+	marker := strings.Repeat("a", 8000) + "b"
+	input := []byte("<<" + marker + "\n" + strings.Repeat("a", 500000) + "\n" + marker)
+	b.SetBytes(int64(len(input)))
+	b.ReportAllocs()
+	for b.Loop() {
+		if _, err := Tokenize(input, "Caddyfile"); err != nil {
+			b.Fatal(err)
+		}
+	}
+}