Fast and extensible multi-platform HTTP/1-2-3 web server with automatic HTTPS
Mirrored from https://github.com/Slazee/caddy.git
Synced Nov 14, 05:15 UTC
caddyhttp: Make responseRecorder capable of counting body size
Matthew Holt committed
Oct 10, 2019 at 15:36 UTC
5300949e0def70411fde307afad4c15b6cc22dfd
1 file changed
+23
-5
modules/caddyhttp/responsewriter.go
+23
-5
@@ -80,6 +80,7 @@ type responseRecorder struct {
80
buf *bytes.Buffer
81
shouldBuffer func(status int) bool
82
stream bool
83
+ size int
84
}
85
86
// NewResponseRecorder returns a new ResponseRecorder that can be
@@ -100,7 +101,9 @@ type responseRecorder struct {
101
// the Buffered() method returns true. If the response was not
102
// buffered, Buffered() will return false and that means the
103
// response bypassed the recorder and was written directly to the
103
-// underlying writer.
104
+// underlying writer. If shouldBuffer is nil, the response will
105
+// never be buffered (it will always be streamed directly), and
106
+// buf can also safely be nil.
107
//
108
// Before calling this function in a middleware handler, make a
109
// new buffer or obtain one from a pool (use the sync.Pool) type.
@@ -130,9 +133,10 @@ func (rr *responseRecorder) WriteHeader(statusCode int) {
133
134
// decide whether we should buffer the response
135
if rr.shouldBuffer == nil {
133
- return
136
+ rr.stream = true
137
+ } else {
138
+ rr.stream = !rr.shouldBuffer(rr.statusCode)
139
}
135
- rr.stream = !rr.shouldBuffer(rr.statusCode)
140
if rr.stream {
141
rr.ResponseWriterWrapper.WriteHeader(rr.statusCode)
142
}
@@ -140,10 +144,17 @@ func (rr *responseRecorder) WriteHeader(statusCode int) {
144
145
func (rr *responseRecorder) Write(data []byte) (int, error) {
146
rr.WriteHeader(http.StatusOK)
147
+ var n int
148
+ var err error
149
if rr.stream {
144
- return rr.ResponseWriterWrapper.Write(data)
150
+ n, err = rr.ResponseWriterWrapper.Write(data)
151
+ } else {
152
+ n, err = rr.buf.Write(data)
153
+ }
154
+ if err == nil {
155
+ rr.size += n
156
}
146
- return rr.buf.Write(data)
157
+ return n, err
158
}
159
160
// Status returns the status code that was written, if any.
@@ -151,6 +162,12 @@ func (rr *responseRecorder) Status() int {
162
return rr.statusCode
163
}
164
165
+// Size returns the number of bytes written,
166
+// not including the response headers.
167
+func (rr *responseRecorder) Size() int {
168
+ return rr.size
169
+}
170
+
171
// Buffer returns the body buffer that rr was created with.
172
// You should still have your original pointer, though.
173
func (rr *responseRecorder) Buffer() *bytes.Buffer {
@@ -169,6 +186,7 @@ type ResponseRecorder interface {
186
Status() int
187
Buffer() *bytes.Buffer
188
Buffered() bool
189
+ Size() int
190
}
191
192
// Interface guards