-
Notifications
You must be signed in to change notification settings - Fork 0
/
injector.go
87 lines (70 loc) · 2.27 KB
/
injector.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
84
85
86
87
package handlers
import (
"bytes"
"io"
"log"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"github.com/PuerkitoBio/goquery"
)
// https://justinas.org/writing-http-middleware-in-go/
// TransformationFunc is the function used to transform the original response
type TransformationFunc func(r io.Reader) (string, error)
// InjectMiddleware intercepts the response and
// modifies it in some way
type InjectMiddleware struct {
handler http.Handler
transform TransformationFunc
}
func (m *InjectMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
// Intercepts only requests with text/html in the Accept request header
// We need to do this so that we inject the script only into html pages
// TODO: Prevent multiple injections for same page (ex: embedded pages, xhr templates, etc)
if !strings.Contains(r.Header.Get("Accept"), "text/html") {
m.handler.ServeHTTP(w, r)
return
}
rec := httptest.NewRecorder()
// passing a ResponseRecorder instead of the original RW
m.handler.ServeHTTP(rec, r)
// Original response body
body := rec.Body.Bytes()
// Modify response only if a valid HTTP response
if rec.Code >= 200 || rec.Code < 300 {
// If we get here, we have to intercept the request and inject our script
log.Println("Intercepting: " + r.URL.String())
content, err := m.transform(bytes.NewReader(body))
if err == nil && content != "" {
body = []byte(content)
}
}
// we copy the original headers first
for k, v := range rec.Header() {
w.Header()[k] = v
}
// Set the correct Content-Length now that we are injecting
// extra data to the response
contentLength := len(body)
w.Header().Set("Content-Length", strconv.Itoa(contentLength))
w.WriteHeader(rec.Code)
// Write out the modified response
w.Write(body)
}
// NewInjectingHandler returns a http.Handler that intercepts
// the original response and modifies it
func NewInjectingHandler(h http.Handler) http.Handler {
return &InjectMiddleware{h, transformResponse}
}
func transformResponse(r io.Reader) (string, error) {
doc, err := goquery.NewDocumentFromReader(r)
if err != nil {
return "", err
}
// Even if body is empty, goquery inserts the missing tags
// automatically when calling doc.Html()
body := doc.Find("body").First()
body.AppendHtml(ScriptContent)
return doc.Html()
}