-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathhandler.go
83 lines (73 loc) · 2.5 KB
/
handler.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
package rewritehtml
import (
"errors"
"golang.org/x/net/html"
"io"
"net/http"
)
// EditorFunc is called to rewriteFn each token of the document until done.
//
// raw is the current content and token is the parsed representation.
// EditorFunc must not modify or retain raw.
// If data is not nil it will be written to the stream instead of raw.
// When the function returns done the rest of the document will be written
// without being parsed.
type EditorFunc func(raw []byte, token *html.Token) (data []byte, done bool)
func Attrs(replacements map[string]string) EditorFunc {
return func(raw []byte, token *html.Token) ([]byte, bool) {
if token.Type == html.StartTagToken {
var found bool
for n := range token.Attr {
if val, ok := replacements[token.Attr[n].Key]; ok {
token.Attr[n].Val = val
found = true
}
}
if found {
return []byte(token.String()), false
}
}
return raw, false
}
}
func AfterTag(tag, data string, once bool) EditorFunc {
return func(raw []byte, token *html.Token) ([]byte, bool) {
if token.Type == html.StartTagToken {
if token.Data == tag {
combined := make([]byte, 0, len(raw)+len(data))
combined = append(combined, raw...)
combined = append(combined, data...)
return combined, once
}
}
return raw, false
}
}
// AfterHead returns an EditorFunc that will inject data after the first <head> tag.
func AfterHead(data string) EditorFunc {
return AfterTag("head", data, true)
}
// Handle will rewriteFn any text/html documents that are served by next.
//
// On each request, Handle will call processRequest to provide an EditorFunc.
// The response will be intercepted and EditorFunc will be applied if the Content-Type is text/html.
// If the Content-Type is not text/html EditorFunc will not be called.
// Handle does not spawn any additional goroutines and will attempt to use the smallest buffer possible
// to read and edit the document.
func Handle(next http.Handler, processRequest func(r *http.Request) (EditorFunc, error)) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
r.Header.Set("Accept-Encoding", "identity")
fn, err := processRequest(r)
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
editor := NewResponseEditor(w, fn)
next.ServeHTTP(editor, r)
// the response has been partially sent and cannot be recovered
// the only thing left to do is to panic
if err := editor.Close(); err != nil && !errors.Is(err, io.EOF) {
panic(err)
}
})
}