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
|
package senddata
import (
"net/http"
)
type sendDataResponseWriter struct {
rw http.ResponseWriter
status int
hijacked bool
req *http.Request
injecters []Injecter
}
func SendData(h http.Handler, injecters ...Injecter) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
s := sendDataResponseWriter{
rw: w,
req: r,
injecters: injecters,
}
defer s.Flush()
h.ServeHTTP(&s, r)
})
}
func (s *sendDataResponseWriter) Header() http.Header {
return s.rw.Header()
}
func (s *sendDataResponseWriter) Write(data []byte) (n int, err error) {
if s.status == 0 {
s.WriteHeader(http.StatusOK)
}
if s.hijacked {
return
}
return s.rw.Write(data)
}
func (s *sendDataResponseWriter) WriteHeader(status int) {
if s.status != 0 {
return
}
s.status = status
if s.status == http.StatusOK && s.tryInject() {
return
}
s.Header().Del(HeaderKey)
s.rw.WriteHeader(s.status)
}
func (s *sendDataResponseWriter) tryInject() bool {
header := s.Header().Get(HeaderKey)
s.Header().Del(HeaderKey)
if header == "" {
return false
}
for _, injecter := range s.injecters {
if injecter.Match(header) {
s.hijacked = true
injecter.Inject(s.rw, s.req, header)
return true
}
}
return false
}
func (s *sendDataResponseWriter) Flush() {
s.WriteHeader(http.StatusOK)
}
|
