-
Notifications
You must be signed in to change notification settings - Fork 32
/
response_writer.go
105 lines (88 loc) · 2.03 KB
/
response_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
package middleware
import (
"bufio"
"fmt"
"io"
"net"
"net/http"
)
// ResponseWriter is the proxy responseWriter
type ResponseWriter interface {
http.ResponseWriter
http.Flusher
http.Hijacker
Status() int
BytesWritten() int
Tee(io.Writer)
Unwrap() http.ResponseWriter
}
// WrapResponseWriter wraps an http.ResponseWriter and returns a ResponseWriter
func wrapResponseWriter(w http.ResponseWriter) ResponseWriter {
return &basicWriter{ResponseWriter: w}
}
var _ ResponseWriter = (*basicWriter)(nil)
var _ http.ResponseWriter = (*basicWriter)(nil)
type basicWriter struct {
http.ResponseWriter
code int
bytes int
tee io.Writer
}
func (b *basicWriter) Header() http.Header {
return b.ResponseWriter.Header()
}
func (b *basicWriter) WriteHeader(code int) {
if !b.Written() {
b.ResponseWriter.WriteHeader(code)
b.code = code
}
}
func (b *basicWriter) Write(data []byte) (int, error) {
if !b.Written() {
b.WriteHeader(http.StatusOK)
}
size, err := b.ResponseWriter.Write(data)
b.bytes += size
return size, err
}
func (b *basicWriter) Written() bool {
return b.Status() != 0
}
func (b *basicWriter) BytesWritten() int {
return b.bytes
}
func (b *basicWriter) Status() int {
return b.code
}
func (b *basicWriter) Tee(w io.Writer) {
b.tee = w
}
func (b *basicWriter) Unwrap() http.ResponseWriter {
return b.ResponseWriter
}
func (b *basicWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
hijacker, ok := b.ResponseWriter.(http.Hijacker)
if !ok {
return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface")
}
return hijacker.Hijack()
}
func (b *basicWriter) CloseNotify() <-chan bool {
return b.ResponseWriter.(http.CloseNotifier).CloseNotify()
}
func (b *basicWriter) Flush() {
fl, ok := b.ResponseWriter.(http.Flusher)
if ok {
fl.Flush()
}
}
func (b *basicWriter) ReadFrom(r io.Reader) (int64, error) {
if b.tee != nil {
return io.Copy(b, r)
}
rf := b.ResponseWriter.(io.ReaderFrom)
if !b.Written() {
b.ResponseWriter.WriteHeader(http.StatusOK)
}
return rf.ReadFrom(r)
}