summaryrefslogtreecommitdiff
path: root/modules/caddyhttp
diff options
context:
space:
mode:
Diffstat (limited to 'modules/caddyhttp')
-rw-r--r--modules/caddyhttp/responsewriter.go35
-rw-r--r--modules/caddyhttp/reverseproxy/streaming.go15
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