livereloadinject: Use more robust injection method

transform/livereloadinject/livereloadinject.go

@@ -14,10 +14,10 @@
 package livereloadinject
 
 import (
-	"bytes"
 	"fmt"
 	"html"
 	"net/url"
+	"regexp"
 	"strings"
 
 	"github.com/gohugoio/hugo/common/loggers"
@@ -25,42 +25,27 @@ import (
 	"github.com/gohugoio/hugo/transform"
 )
 
-const warnMessage = `"head" or "body" tag is required in html to append livereload script. ` +
-	"As a fallback, Hugo injects it somewhere but it might not work properly."
-
-var warnScript = fmt.Sprintf(`<script data-no-instant defer>console.warn('%s');</script>`, warnMessage)
-
-type tag struct {
-	markup       []byte
-	appendScript bool
-	warnRequired bool
+var ignoredSyntax = regexp.MustCompile(`(?s)^(?:\s+|<!--.*?-->|<\?.*?\?>)*`)
+var tagsBeforeHead = []*regexp.Regexp{
+	regexp.MustCompile(`(?is)^<!doctype\s[^>]*>`),
+	regexp.MustCompile(`(?is)^<html(?:\s[^>]*)?>`),
+	regexp.MustCompile(`(?is)^<head(?:\s[^>]*)?>`),
 }
 
-var tags = []tag{
-	{markup: []byte("<head"), appendScript: true},
-	{markup: []byte("<HEAD"), appendScript: true},
-	{markup: []byte("</body>")},
-	{markup: []byte("</BODY>")},
-	{markup: []byte("<html"), appendScript: true, warnRequired: true},
-	{markup: []byte("<HTML"), appendScript: true, warnRequired: true},
-}
-
-// New creates a function that can be used
-// to inject a script tag for the livereload JavaScript in a HTML document.
+// New creates a function that can be used to inject a script tag for
+// the livereload JavaScript at the start of an HTML document's head.
 func New(baseURL url.URL) transform.Transformer {
 	return func(ft transform.FromTo) error {
 		b := ft.From().Bytes()
-		idx := -1
-		var match tag
-		// We used to insert the livereload script right before the closing body.
-		// This does not work when combined with tools such as Turbolinks.
-		// So we try to inject the script as early as possible.
-		for _, t := range tags {
-			idx = bytes.Index(b, t.markup)
-			if idx != -1 {
-				match = t
-				break
-			}
+
+		// We find the start of the head by reading past (in order)
+		// the doctype declaration, HTML start tag and head start tag,
+		// all of which are optional, and any whitespace, comments, or
+		// XML instructions in-between.
+		idx := 0
+		for _, tag := range tagsBeforeHead {
+			idx += len(ignoredSyntax.Find(b[idx:]))
+			idx += len(tag.Find(b[idx:]))
 		}
 
 		path := strings.TrimSuffix(baseURL.Path, "/")
@@ -72,23 +57,9 @@ func New(baseURL url.URL) transform.Transformer {
 		c := make([]byte, len(b))
 		copy(c, b)
 
-		if idx == -1 {
-			idx = len(b)
-			match = tag{warnRequired: true}
-		}
-
 		script := []byte(fmt.Sprintf(`<script src="%s" data-no-instant defer></script>`, html.EscapeString(src)))
 
-		i := idx
-		if match.appendScript {
-			i += bytes.Index(b[i:], []byte(">")) + 1
-		}
-
-		if match.warnRequired {
-			script = append(script, []byte(warnScript)...)
-		}
-
-		c = append(c[:i], append(script, c[i:]...)...)
+		c = append(c[:idx], append(script, c[idx:]...)...)
 
 		if _, err := ft.To().Write(c); err != nil {
 			loggers.Log().Warnf("Failed to inject LiveReload script:", err)