diff options
| -rw-r--r-- | modules/caddyhttp/responsewriter.go | 35 | ||||
| -rw-r--r-- | 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 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 |