-
-
Notifications
You must be signed in to change notification settings - Fork 39
/
response_writer.go
133 lines (112 loc) · 3.5 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
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
132
133
// Copyright 2021 Flamego. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.
package flamego
import (
"bufio"
"net"
"net/http"
"sync"
"sync/atomic"
"github.com/pkg/errors"
)
// ResponseWriter is a wrapper of http.ResponseWriter that provides extra
// information about the response. It is recommended that middleware handlers
// use this construct to wrap a ResponseWriter if the functionality calls for
// it.
type ResponseWriter interface {
http.ResponseWriter
http.Flusher
http.Pusher
// Status returns the status code of the response or 0 if the response has not
// been written.
Status() int
// Written returns whether the ResponseWriter has been written.
Written() bool
// Size returns the written size of the response body.
Size() int
// Before allows for a function to be called before the ResponseWriter has been
// written to. This is useful for setting headers or any other operations that
// must happen before a response has been written. Multiple calls to this method
// will stack up functions, and functions will be called in the FILO manner.
Before(BeforeFunc)
}
type responseWriter struct {
http.ResponseWriter
method string // The HTTP method of the coming request.
status int32 // The written status of the response.
size int // The written size of the response.
beforeFuncs []BeforeFunc // The list of functions to be called before written to the response.
writeHeaderOnce sync.Once
}
// BeforeFunc is a function that is called before the ResponseWriter is written.
type BeforeFunc func(ResponseWriter)
// NewResponseWriter returns a wrapper of http.ResponseWriter.
func NewResponseWriter(method string, w http.ResponseWriter) ResponseWriter {
return &responseWriter{
ResponseWriter: w,
method: method,
}
}
func (w *responseWriter) callBefore() {
for i := len(w.beforeFuncs) - 1; i >= 0; i-- {
w.beforeFuncs[i](w)
}
}
func (w *responseWriter) WriteHeader(s int) {
w.writeHeaderOnce.Do(func() {
if w.Written() {
return
}
w.callBefore()
w.ResponseWriter.WriteHeader(s)
atomic.StoreInt32(&w.status, int32(s))
})
}
func (w *responseWriter) Write(b []byte) (size int, err error) {
if !w.Written() {
// The status will be StatusOK if WriteHeader has not been called yet.
w.WriteHeader(http.StatusOK)
}
if w.method != http.MethodHead {
size, err = w.ResponseWriter.Write(b)
w.size += size
}
return size, err
}
func (w *responseWriter) Status() int {
return int(atomic.LoadInt32(&w.status))
}
func (w *responseWriter) Size() int {
return w.size
}
func (w *responseWriter) Written() bool {
return w.Status() != 0
}
func (w *responseWriter) Before(before BeforeFunc) {
w.beforeFuncs = append(w.beforeFuncs, before)
}
func (w *responseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
hijacker, ok := w.ResponseWriter.(http.Hijacker)
if !ok {
return nil, nil, errors.New("the ResponseWriter doesn't support the Hijacker interface")
}
return hijacker.Hijack()
}
func (w *responseWriter) Flush() {
if !w.Written() {
// The status will be StatusOK if WriteHeader has not been called yet.
w.WriteHeader(http.StatusOK)
}
flusher, ok := w.ResponseWriter.(http.Flusher)
if ok {
flusher.Flush()
}
}
func (w *responseWriter) Push(target string, opts *http.PushOptions) error {
pusher, ok := w.ResponseWriter.(http.Pusher)
if !ok {
return errors.New("the ResponseWriter doesn't support the Pusher interface")
}
return pusher.Push(target, opts)
}