/
middleware.go
81 lines (71 loc) · 1.93 KB
/
middleware.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
package livereload
import (
"bytes"
"compress/gzip"
"compress/zlib"
"fmt"
"io"
"net/http"
"regexp"
"strings"
)
type writer struct {
http.ResponseWriter
host string
}
func (self *writer) addJavaScript(data []byte) []byte {
javascript := fmt.Sprintf(`<script defer src="//%s%s"></script>
</head>`, self.host, URL.JavaScript)
return regexp.MustCompile(`</head>`).ReplaceAll(data, []byte(javascript))
}
func (self *writer) Write(data []byte) (size int, e error) {
if strings.Contains(self.Header().Get("Content-Type"), "html") {
var encoding = self.Header().Get("Content-Encoding")
if encoding == "" {
data = self.addJavaScript(data)
} else {
var reader io.ReadCloser
var buffer *bytes.Buffer = new(bytes.Buffer)
if encoding == "gzip" {
// decode to add javascript reference.
reader, _ = gzip.NewReader(bytes.NewReader(data))
io.Copy(buffer, reader)
output := self.addJavaScript(buffer.Bytes())
reader.Close()
buffer.Reset()
// encode back to HTML with added javascript reference.
writer := gzip.NewWriter(buffer)
writer.Write(output)
writer.Close()
data = buffer.Bytes()
} else if encoding == "deflate" {
// decode to add javascript reference.
reader, _ = zlib.NewReader(bytes.NewReader(data))
io.Copy(buffer, reader)
output := self.addJavaScript(buffer.Bytes())
reader.Close()
buffer.Reset()
// encode back to HTML with added javascript reference.
writer := zlib.NewWriter(buffer)
writer.Write(output)
writer.Close()
data = buffer.Bytes()
}
}
}
return self.ResponseWriter.Write(data)
}
func Middleware(next http.Handler) http.Handler {
Start()
fn := func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == URL.WebSocket {
ServeWebSocket(w, r)
} else if r.URL.Path == URL.JavaScript {
ServeJavaScript(w, r)
} else {
writer := &writer{w, r.Host}
next.ServeHTTP(writer, r)
}
}
return http.HandlerFunc(fn)
}