transform/livereloadinject: Inject livereload script right after head if possible

bep · bep · commit 873624469c66 · 2020-01-29T19:41:42.000+01:00
This makes it work with Turbolinks. Fixes gohugoio#6821
diff --git a/transform/livereloadinject/livereloadinject.go b/transform/livereloadinject/livereloadinject.go
@@ -21,25 +21,54 @@ import (
 	"github.com/gohugoio/hugo/transform"
 )
 
+type tag struct {
+	markup       []byte
+	appendScript bool
+}
+
+var tags = []tag{
+	tag{markup: []byte("<head>"), appendScript: true},
+	tag{markup: []byte("<HEAD>"), appendScript: true},
+	tag{markup: []byte("</body>")},
+	tag{markup: []byte("</BODY>")},
+}
+
 // New creates a function that can be used
 // to inject a script tag for the livereload JavaScript in a HTML document.
 func New(port int) transform.Transformer {
 	return func(ft transform.FromTo) error {
 		b := ft.From().Bytes()
-		endBodyTag := "</body>"
-		match := []byte(endBodyTag)
-		replaceTemplate := `<script data-no-instant>document.write('<script src="/livereload.js?port=%d&mindelay=10&v=2"></' + 'script>')</script>%s`
-		replace := []byte(fmt.Sprintf(replaceTemplate, port, endBodyTag))
-
-		newcontent := bytes.Replace(b, match, replace, 1)
-		if len(newcontent) == len(b) {
-			endBodyTag = "</BODY>"
-			replace := []byte(fmt.Sprintf(replaceTemplate, port, endBodyTag))
-			match := []byte(endBodyTag)
-			newcontent = bytes.Replace(b, match, replace, 1)
+		var 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
+			}
 		}
 
-		if _, err := ft.To().Write(newcontent); err != nil {
+		c := make([]byte, len(b))
+		copy(c, b)
+
+		if idx == -1 {
+			_, err := ft.To().Write(c)
+			return err
+		}
+
+		script := []byte(fmt.Sprintf(`<script data-no-instant>document.write('<script src="/livereload.js?port=%d&mindelay=10&v=2"></' + 'script>')</script>`, port))
+
+		i := idx
+		if match.appendScript {
+			i += len(match.markup)
+		}
+
+		c = append(c[:i], append(script, c[i:]...)...)
+
+		if _, err := ft.To().Write(c); err != nil {
 			helpers.DistinctWarnLog.Println("Failed to inject LiveReload script:", err)
 		}
 		return nil
diff --git a/transform/livereloadinject/livereloadinject_test.go b/transform/livereloadinject/livereloadinject_test.go
@@ -15,27 +15,45 @@ package livereloadinject
 
 import (
 	"bytes"
-	"fmt"
 	"strings"
 	"testing"
 
+	qt "github.com/frankban/quicktest"
 	"github.com/gohugoio/hugo/transform"
 )
 
 func TestLiveReloadInject(t *testing.T) {
-	doTestLiveReloadInject(t, "</body>")
-	doTestLiveReloadInject(t, "</BODY>")
-}
+	c := qt.New(t)
 
-func doTestLiveReloadInject(t *testing.T, bodyEndTag string) {
-	out := new(bytes.Buffer)
-	in := strings.NewReader(bodyEndTag)
+	expectBase := `<script data-no-instant>document.write('<script src="/livereload.js?port=1313&mindelay=10&v=2"></' + 'script>')</script>`
+	apply := func(s string) string {
+		out := new(bytes.Buffer)
+		in := strings.NewReader(s)
 
-	tr := transform.New(New(1313))
-	tr.Apply(out, in)
+		tr := transform.New(New(1313))
+		tr.Apply(out, in)
 
-	expected := fmt.Sprintf(`<script data-no-instant>document.write('<script src="/livereload.js?port=1313&mindelay=10&v=2"></' + 'script>')</script>%s`, bodyEndTag)
-	if out.String() != expected {
-		t.Errorf("Expected %s got %s", expected, out.String())
+		return out.String()
 	}
+
+	c.Run("Head lower", func(c *qt.C) {
+		c.Assert(apply("<html><head>foo"), qt.Equals, "<html><head>"+expectBase+"foo")
+	})
+
+	c.Run("Head upper", func(c *qt.C) {
+		c.Assert(apply("<html><HEAD>foo"), qt.Equals, "<html><HEAD>"+expectBase+"foo")
+	})
+
+	c.Run("Body lower", func(c *qt.C) {
+		c.Assert(apply("foo</body>"), qt.Equals, "foo"+expectBase+"</body>")
+	})
+
+	c.Run("Body upper", func(c *qt.C) {
+		c.Assert(apply("foo</BODY>"), qt.Equals, "foo"+expectBase+"</BODY>")
+	})
+
+	c.Run("No match", func(c *qt.C) {
+		c.Assert(apply("<h1>No match</h1>"), qt.Equals, "<h1>No match</h1>")
+	})
+
 }
