diff options
author | WeidiDeng <weidi_deng@icloud.com> | 2023-02-07 07:14:59 +0800 |
---|---|---|
committer | GitHub <noreply@github.com> | 2023-02-06 16:14:59 -0700 |
commit | c77a6bea66b30ea88cd991671bf67669c80ffcdd (patch) | |
tree | 2f908a42f559bd6526034712298e7ca425879c2a /modules/caddyhttp/responsewriter.go | |
parent | 12bcbe2c4924ecbf6730fc340a7a4250bddcc9be (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/responsewriter.go')
-rw-r--r-- | modules/caddyhttp/responsewriter.go | 35 |
1 files changed, 30 insertions, 5 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. |