From c77a6bea66b30ea88cd991671bf67669c80ffcdd Mon Sep 17 00:00:00 2001
From: WeidiDeng <weidi_deng@icloud.com>
Date: Tue, 7 Feb 2023 07:14:59 +0800
Subject: [PATCH] reverseproxy: Log status code and byte count for websockets
 (#5140)

* log response size for websocket request

* record size when using hijack bufio.Writer
---
 modules/caddyhttp/responsewriter.go         | 35 ++++++++++++++++++---
 modules/caddyhttp/reverseproxy/streaming.go | 15 ++++-----
 2 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/modules/caddyhttp/responsewriter.go b/modules/caddyhttp/responsewriter.go
index 0631229d..1b28cf0b 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 834cb9e1..1db107af 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