/
writer.go
109 lines (90 loc) · 2.21 KB
/
writer.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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
package httpu
import (
"io"
"net/http"
"github.com/pkg/errors"
)
type responseWriter struct {
http.ResponseWriter
wroteHeader, wrote bool
bytesWritten int
status int
tee []io.Writer
}
func NewResponseWriter(w http.ResponseWriter) *responseWriter {
return &responseWriter{ResponseWriter: w}
}
func (this *responseWriter) WriteHeader(s int) {
if !this.wroteHeader {
this.wroteHeader = true
this.ResponseWriter.WriteHeader(s)
this.status = s
}
}
func (this *responseWriter) Write(p []byte) (n int, err error) {
if !this.wrote {
this.WriteHeader(200)
this.wrote = true
}
n, err = this.ResponseWriter.Write(p)
this.bytesWritten += n
for i, tee := range this.tee {
_, err2 := tee.Write(p[:n])
// Prefer errors generated by the proxied writer.
if err == nil {
err = errors.Wrapf(err2, "tee[%d]", i)
}
}
return
}
func (this *responseWriter) WroteHeader() bool {
return this.wroteHeader
}
func (this *responseWriter) Wrote() bool {
return this.wrote
}
func (this *responseWriter) BytesWritten() int {
return this.bytesWritten
}
func (this *responseWriter) Status() int {
return this.status
}
func (this *responseWriter) Tee(w io.Writer) {
this.tee = append(this.tee, w)
}
func (this *responseWriter) Unwrap() http.ResponseWriter {
return this.ResponseWriter
}
type teeResponseWriter struct {
http.ResponseWriter
tee []io.Writer
}
func NewTeeResponseWriter(w http.ResponseWriter, tee ...io.Writer) *teeResponseWriter {
return &teeResponseWriter{ResponseWriter: w, tee: tee}
}
func (this *teeResponseWriter) Tee(w io.Writer) {
this.tee = append(this.tee, w)
}
func (this *teeResponseWriter) Write(p []byte) (n int, err error) {
n, err = this.ResponseWriter.Write(p)
for i, tee := range this.tee {
_, err2 := tee.Write(p[:n])
// Prefer errors generated by the proxied writer.
if err == nil {
err = errors.Wrapf(err2, "tee[%d]", i)
}
}
return
}
func ResponseWriterOf(w http.ResponseWriter) (wd ResponseWriter) {
if wd, ok := w.(ResponseWriter); ok {
return wd
}
return NewResponseWriter(w)
}
func TeeResponseWriterOf(w http.ResponseWriter) (wd TeeResponseWriter) {
if tw, ok := w.(TeeResponseWriter); ok {
return tw
}
return NewTeeResponseWriter(w)
}