-
Notifications
You must be signed in to change notification settings - Fork 6
/
respwriter.go
131 lines (101 loc) · 2.57 KB
/
respwriter.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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
package httputil
import (
"bufio"
"fmt"
"io"
"net"
"net/http"
)
// NewResponseWriterWrapper wraps an http.ResponseWriter with an enhanced proxy.
func NewResponseWriterWrapper(w http.ResponseWriter) ResponseWriterWrapper {
return &responseWriterWrapper{ResponseWriter: w}
}
// ResponseWriterWrapper is the interface defining the extendend functions of the proxy.
type ResponseWriterWrapper interface {
http.ResponseWriter
// Size returns the total number of bytes sent to the client.
Size() int
// Status returns the HTTP status of the request.
Status() int
// Tee sets a writer that will contain a copy of the bytes written to the response writer.
Tee(io.Writer)
}
type responseWriterWrapper struct {
http.ResponseWriter
headerWritten bool
size int
status int
tee io.Writer
}
func (b *responseWriterWrapper) Size() int {
return b.size
}
func (b *responseWriterWrapper) Flush() {
b.headerWritten = true
fl, ok := b.ResponseWriter.(http.Flusher)
if ok {
fl.Flush()
}
}
//nolint:wrapcheck
func (b *responseWriterWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
hj, ok := b.ResponseWriter.(http.Hijacker)
if !ok {
return nil, nil, fmt.Errorf("the Hijacker is not supported by the ResponseWriter")
}
return hj.Hijack()
}
//nolint:wrapcheck
func (b *responseWriterWrapper) Push(target string, opts *http.PushOptions) error {
pusher, ok := b.ResponseWriter.(http.Pusher)
if !ok {
return fmt.Errorf("the Pusher is not supported by the ResponseWriter")
}
return pusher.Push(target, opts)
}
//nolint:wrapcheck
func (b *responseWriterWrapper) ReadFrom(r io.Reader) (int64, error) {
if b.tee != nil {
n, err := io.Copy(b, r)
b.size += int(n)
return n, err
}
rf, ok := b.ResponseWriter.(io.ReaderFrom)
if !ok {
return 0, fmt.Errorf("the ReaderFrom is not supported by the ResponseWriter")
}
b.maybeWriteHeader()
n, err := rf.ReadFrom(r)
b.size += int(n)
return n, err
}
func (b *responseWriterWrapper) Status() int {
return b.status
}
func (b *responseWriterWrapper) Tee(w io.Writer) {
b.tee = w
}
func (b *responseWriterWrapper) Write(buf []byte) (int, error) {
b.maybeWriteHeader()
n, err := b.ResponseWriter.Write(buf)
if b.tee != nil {
_, teeErr := b.tee.Write(buf[:n])
if err == nil {
err = teeErr
}
}
b.size += n
return n, err
}
func (b *responseWriterWrapper) WriteHeader(code int) {
if !b.headerWritten {
b.status = code
b.headerWritten = true
b.ResponseWriter.WriteHeader(code)
}
}
func (b *responseWriterWrapper) maybeWriteHeader() {
if !b.headerWritten {
b.WriteHeader(http.StatusOK)
}
}