Fast and extensible multi-platform HTTP/1-2-3 web server with automatic HTTPS
Mirrored from https://github.com/Slazee/caddy.git
Synced Nov 13, 05:30 UTC
reverseproxy: Log status code and byte count for websockets (#5140)
* log response size for websocket request * record size when using hijack bufio.Writer
WeidiDeng committed
Feb 7, 2023 at 07:14 UTC
c77a6bea66b30ea88cd991671bf67669c80ffcdd
2 files changed
+36
-14
modules/caddyhttp/responsewriter.go
+30
-5
@@ -211,11 +211,7 @@ func (rr *responseRecorder) ReadFrom(r io.Reader) (int64, error) {
211
var n int64
212
var err error
213
if rr.stream {
214
- if rf, ok := rr.ResponseWriter.(io.ReaderFrom); ok {
215
- n, err = rf.ReadFrom(r)
216
- } else {
217
- n, err = io.Copy(rr.ResponseWriter, r)
218
- }
214
+ n, err = rr.ResponseWriterWrapper.ReadFrom(r)
215
} else {
216
n, err = rr.buf.ReadFrom(r)
217
}
@@ -260,6 +256,35 @@ func (rr *responseRecorder) WriteResponse() error {
256
return err
257
}
258
259
+func (rr *responseRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
260
+ conn, brw, err := rr.ResponseWriterWrapper.Hijack()
261
+ if err != nil {
262
+ return nil, nil, err
263
+ }
264
+ // Per http documentation, returned bufio.Writer is empty, but bufio.Read maybe not
265
+ conn = &hijackedConn{conn, rr}
266
+ brw.Writer.Reset(conn)
267
+ return conn, brw, nil
268
+}
269
+
270
+// used to track the size of hijacked response writers
271
+type hijackedConn struct {
272
+ net.Conn
273
+ rr *responseRecorder
274
+}
275
+
276
+func (hc *hijackedConn) Write(p []byte) (int, error) {
277
+ n, err := hc.Conn.Write(p)
278
+ hc.rr.size += n
279
+ return n, err
280
+}
281
+
282
+func (hc *hijackedConn) ReadFrom(r io.Reader) (int64, error) {
283
+ n, err := io.Copy(hc.Conn, r)
284
+ hc.rr.size += int(n)
285
+ return n, err
286
+}
287
+
288
// ResponseRecorder is a http.ResponseWriter that records
289
// responses instead of writing them to the client. See
290
// docs for NewResponseRecorder for proper usage.
modules/caddyhttp/reverseproxy/streaming.go
+6
-9
@@ -74,27 +74,24 @@ func (h Handler) handleUpgradeResponse(logger *zap.Logger, rw http.ResponseWrite
74
}()
75
defer close(backConnCloseCh)
76
77
+ // write header first, response headers should not be counted in size
78
+ // like the rest of handler chain.
79
+ copyHeader(rw.Header(), res.Header)
80
+ rw.WriteHeader(res.StatusCode)
81
+
82
logger.Debug("upgrading connection")
83
conn, brw, err := hj.Hijack()
84
if err != nil {
85
h.logger.Error("hijack failed on protocol switch", zap.Error(err))
86
return
87
}
83
- defer conn.Close()
88
89
start := time.Now()
90
defer func() {
91
+ conn.Close()
92
logger.Debug("connection closed", zap.Duration("duration", time.Since(start)))
93
}()
94
90
- copyHeader(rw.Header(), res.Header)
91
-
92
- res.Header = rw.Header()
93
- res.Body = nil // so res.Write only writes the headers; we have res.Body in backConn above
94
- if err := res.Write(brw); err != nil {
95
- h.logger.Debug("response write", zap.Error(err))
96
- return
97
- }
95
if err := brw.Flush(); err != nil {
96
h.logger.Debug("response flush", zap.Error(err))
97
return