diff options
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. |