1
0
mirror of https://github.com/danog/gojekyll.git synced 2024-12-12 17:18:13 +01:00
gojekyll/server/taginjector.go

31 lines
934 B
Go

package server
import (
"bytes"
"io"
)
var closeHeadTag = []byte(`</head>`) // TODO also look for </HEAD>
// TagInjector wraps a writer and adds a script tag to its content.
// It depends on the fact that dynamic page rendering makes a single Write call,
// so that it's guaranteed to find the marker within a single invocation argument.
// It doesn't parse HTML, so it could be spoofed but probably only intentionally.
type TagInjector struct {
w io.Writer
insertion []byte
}
// Write injects a livereload script tag at the end of the HTML head, if present,
// else at the beginning of the document.
func (i TagInjector) Write(b []byte) (n int, err error) {
if !bytes.Contains(b, i.insertion) && bytes.Contains(b, closeHeadTag) {
r := append(i.insertion, closeHeadTag...)
b = bytes.Replace(b, closeHeadTag, r, 1)
}
if !bytes.Contains(b, i.insertion) {
b = append(i.insertion, b...)
}
return i.w.Write(b)
}