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: 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 ++++++++++++++++++++++++++++++-----
 1 file changed, 30 insertions(+), 5 deletions(-)

(limited to 'modules/caddyhttp/responsewriter.go')

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.
-- 
cgit v1.2.3