diff --git a/caddyconfig/caddyfile/lexer.go b/caddyconfig/caddyfile/lexer.go index 2866039ed..6b60627e6 100644 --- a/caddyconfig/caddyfile/lexer.go +++ b/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 { diff --git a/caddyconfig/caddyfile/lexer_test.go b/caddyconfig/caddyfile/lexer_test.go index 89dde2d9f..d4b60ec55 100644 --- a/caddyconfig/caddyfile/lexer_test.go +++ b/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) + } + } +}