From abf5ab340ed76792214ae80c62df7abe0ad1b8a8 Mon Sep 17 00:00:00 2001 From: Matthew Holt Date: Tue, 15 Oct 2019 14:07:10 -0600 Subject: caddyhttp: Improve ResponseRecorder to buffer headers --- modules/caddyhttp/responsewriter.go | 119 ++++++++++++++++++++++++++---------- 1 file changed, 88 insertions(+), 31 deletions(-) (limited to 'modules/caddyhttp/responsewriter.go') diff --git a/modules/caddyhttp/responsewriter.go b/modules/caddyhttp/responsewriter.go index 344298f..5beb40e 100644 --- a/modules/caddyhttp/responsewriter.go +++ b/modules/caddyhttp/responsewriter.go @@ -18,6 +18,7 @@ import ( "bufio" "bytes" "fmt" + "io" "net" "net/http" ) @@ -78,52 +79,89 @@ type responseRecorder struct { wroteHeader bool statusCode int buf *bytes.Buffer - shouldBuffer func(status int) bool + shouldBuffer ShouldBufferFunc stream bool size int + header http.Header } // NewResponseRecorder returns a new ResponseRecorder that can be -// used instead of a real http.ResponseWriter. The recorder is useful -// for middlewares which need to buffer a responder's response and -// process it in its entirety before actually allowing the response to -// be written. Of course, this has a performance overhead, but -// sometimes there is no way to avoid buffering the whole response. -// Still, if at all practical, middlewares should strive to stream +// used instead of a standard http.ResponseWriter. The recorder is +// useful for middlewares which need to buffer a response and +// potentially process its entire body before actually writing the +// response to the underlying writer. Of course, buffering the entire +// body has a memory overhead, but sometimes there is no way to avoid +// buffering the whole response, hence the existence of this type. +// Still, if at all practical, handlers should strive to stream // responses by wrapping Write and WriteHeader methods instead of // buffering whole response bodies. // -// Recorders optionally buffer the response. When the headers are -// to be written, shouldBuffer will be called with the status -// code that is being written. The rest of the headers can be read -// from w.Header(). If shouldBuffer returns true, the response -// will be buffered. You can know the response was buffered if -// the Buffered() method returns true. If the response was not -// buffered, Buffered() will return false and that means the -// response bypassed the recorder and was written directly to the -// underlying writer. If shouldBuffer is nil, the response will -// never be buffered (it will always be streamed directly), and -// buf can also safely be nil. +// Buffering is actually optional. The shouldBuffer function will +// be called just before the headers are written. If it returns +// true, the headers and body will be buffered by this recorder +// and not written to the underlying writer; if false, the headers +// will be written immediately and the body will be streamed out +// directly to the underlying writer. If shouldBuffer is nil, +// the response will never be buffered and will always be streamed +// directly to the writer. // -// Before calling this function in a middleware handler, make a -// new buffer or obtain one from a pool (use the sync.Pool) type. -// Using a pool is generally recommended for performance gains; -// do profiling to ensure this is the case. If using a pool, be -// sure to reset the buffer before using it. +// You can know if shouldBuffer returned true by calling Buffered(). // -// The returned recorder can be used in place of w when calling -// the next handler in the chain. When that handler returns, you -// can read the status code from the recorder's Status() method. -// The response body fills buf if it was buffered, and the headers -// are available via w.Header(). -func NewResponseRecorder(w http.ResponseWriter, buf *bytes.Buffer, shouldBuffer func(status int) bool) ResponseRecorder { +// The provided buffer buf should be obtained from a pool for best +// performance (see the sync.Pool type). +// +// Proper usage of a recorder looks like this: +// +// rec := caddyhttp.NewResponseRecorder(w, buf, shouldBuffer) +// err := next.ServeHTTP(rec, req) +// if err != nil { +// return err +// } +// if !rec.Buffered() { +// return nil +// } +// // process the buffered response here +// +// After a response has been buffered, remember that any upstream header +// manipulations are only manifest in the recorder's Header(), not the +// Header() of the underlying ResponseWriter. Thus if you wish to inspect +// or change response headers, you either need to use rec.Header(), or +// copy rec.Header() into w.Header() first (see caddyhttp.CopyHeader). +// +// Once you are ready to write the response, there are two ways you can do +// it. The easier way is to have the recorder do it: +// +// rec.WriteResponse() +// +// This writes the recorded response headers as well as the buffered body. +// Or, you may wish to do it yourself, especially if you manipulated the +// buffered body. First you will need to copy the recorded headers, then +// write the headers with the recorded status code, then write the body +// (this example writes the recorder's body buffer, but you might have +// your own body to write instead): +// +// caddyhttp.CopyHeader(w.Header(), rec.Header()) +// w.WriteHeader(rec.Status()) +// io.Copy(w, rec.Buffer()) +// +func NewResponseRecorder(w http.ResponseWriter, buf *bytes.Buffer, shouldBuffer ShouldBufferFunc) ResponseRecorder { + // copy the current response header into this buffer so + // that any header manipulations on the buffered header + // are consistent with what would be written out + hdr := make(http.Header) + CopyHeader(hdr, w.Header()) return &responseRecorder{ ResponseWriterWrapper: &ResponseWriterWrapper{ResponseWriter: w}, buf: buf, shouldBuffer: shouldBuffer, + header: hdr, } } +func (rr *responseRecorder) Header() http.Header { + return rr.header +} + func (rr *responseRecorder) WriteHeader(statusCode int) { if rr.wroteHeader { return @@ -135,9 +173,12 @@ func (rr *responseRecorder) WriteHeader(statusCode int) { if rr.shouldBuffer == nil { rr.stream = true } else { - rr.stream = !rr.shouldBuffer(rr.statusCode) + rr.stream = !rr.shouldBuffer(rr.statusCode, rr.header) } + + // if not buffered, immediately write header if rr.stream { + CopyHeader(rr.ResponseWriterWrapper.Header(), rr.header) rr.ResponseWriterWrapper.WriteHeader(rr.statusCode) } } @@ -179,16 +220,32 @@ func (rr *responseRecorder) Buffered() bool { return !rr.stream } +func (rr *responseRecorder) WriteResponse() error { + if rr.stream { + return nil + } + CopyHeader(rr.ResponseWriterWrapper.Header(), rr.header) + rr.ResponseWriterWrapper.WriteHeader(rr.statusCode) + _, err := io.Copy(rr.ResponseWriterWrapper, rr.buf) + return err +} // ResponseRecorder is a http.ResponseWriter that records -// responses instead of writing them to the client. +// responses instead of writing them to the client. See +// docs for NewResponseRecorder for proper usage. type ResponseRecorder interface { HTTPInterfaces Status() int Buffer() *bytes.Buffer Buffered() bool Size() int + WriteResponse() error } +// ShouldBufferFunc is a function that returns true if the +// response should be buffered, given the pending HTTP status +// code and response headers. +type ShouldBufferFunc func(status int, header http.Header) bool + // Interface guards var ( _ HTTPInterfaces = (*ResponseWriterWrapper)(nil) -- cgit v1.2.3