summaryrefslogtreecommitdiff
path: root/modules/caddyhttp
diff options
context:
space:
mode:
authorWeidiDeng <weidi_deng@icloud.com>2023-02-07 07:14:59 +0800
committerGitHub <noreply@github.com>2023-02-06 16:14:59 -0700
commitc77a6bea66b30ea88cd991671bf67669c80ffcdd (patch)
tree2f908a42f559bd6526034712298e7ca425879c2a /modules/caddyhttp
parent12bcbe2c4924ecbf6730fc340a7a4250bddcc9be (diff)
reverseproxy: Log status code and byte count for websockets (#5140)
* log response size for websocket request * record size when using hijack bufio.Writer
Diffstat (limited to 'modules/caddyhttp')
-rw-r--r--modules/caddyhttp/responsewriter.go35
-rw-r--r--modules/caddyhttp/reverseproxy/streaming.go15
2 files changed, 36 insertions, 14 deletions
diff --git a/modules/caddyhttp/responsewriter.go b/modules/caddyhttp/responsewriter.go
index 0631229..1b28cf0 100644
--- a/modules/caddyhttp/responsewriter.go
+++ b/modules/caddyhttp/responsewriter.go
@@ -211,11 +211,7 @@ func (rr *responseRecorder) ReadFrom(r io.Reader) (int64, error) {
var n int64
var err error
if rr.stream {
- if rf, ok := rr.ResponseWriter.(io.ReaderFrom); ok {
- n, err = rf.ReadFrom(r)
- } else {
- n, err = io.Copy(rr.ResponseWriter, r)
- }
+ n, err = rr.ResponseWriterWrapper.ReadFrom(r)
} else {
n, err = rr.buf.ReadFrom(r)
}
@@ -260,6 +256,35 @@ func (rr *responseRecorder) WriteResponse() error {
return err
}
+func (rr *responseRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+ conn, brw, err := rr.ResponseWriterWrapper.Hijack()
+ if err != nil {
+ return nil, nil, err
+ }
+ // Per http documentation, returned bufio.Writer is empty, but bufio.Read maybe not
+ conn = &hijackedConn{conn, rr}
+ brw.Writer.Reset(conn)
+ return conn, brw, nil
+}
+
+// used to track the size of hijacked response writers
+type hijackedConn struct {
+ net.Conn
+ rr *responseRecorder
+}
+
+func (hc *hijackedConn) Write(p []byte) (int, error) {
+ n, err := hc.Conn.Write(p)
+ hc.rr.size += n
+ return n, err
+}
+
+func (hc *hijackedConn) ReadFrom(r io.Reader) (int64, error) {
+ n, err := io.Copy(hc.Conn, r)
+ hc.rr.size += int(n)
+ return n, err
+}
+
// ResponseRecorder is a http.ResponseWriter that records
// responses instead of writing them to the client. See
// docs for NewResponseRecorder for proper usage.
diff --git a/modules/caddyhttp/reverseproxy/streaming.go b/modules/caddyhttp/reverseproxy/streaming.go
index 834cb9e..1db107a 100644
--- a/modules/caddyhttp/reverseproxy/streaming.go
+++ b/modules/caddyhttp/reverseproxy/streaming.go
@@ -74,27 +74,24 @@ func (h Handler) handleUpgradeResponse(logger *zap.Logger, rw http.ResponseWrite
}()
defer close(backConnCloseCh)
+ // write header first, response headers should not be counted in size
+ // like the rest of handler chain.
+ copyHeader(rw.Header(), res.Header)
+ rw.WriteHeader(res.StatusCode)
+
logger.Debug("upgrading connection")
conn, brw, err := hj.Hijack()
if err != nil {
h.logger.Error("hijack failed on protocol switch", zap.Error(err))
return
}
- defer conn.Close()
start := time.Now()
defer func() {
+ conn.Close()
logger.Debug("connection closed", zap.Duration("duration", time.Since(start)))
}()
- copyHeader(rw.Header(), res.Header)
-
- res.Header = rw.Header()
- res.Body = nil // so res.Write only writes the headers; we have res.Body in backConn above
- if err := res.Write(brw); err != nil {
- h.logger.Debug("response write", zap.Error(err))
- return
- }
if err := brw.Flush(); err != nil {
h.logger.Debug("response flush", zap.Error(err))
return