caddyhttp: Improve ResponseRecorder to buffer headers

This commit is contained in:
Matthew Holt 2019-10-15 14:07:10 -06:00
parent acf7dea68f
commit abf5ab340e
No known key found for this signature in database
GPG key ID: 2A349DD577D586A5
6 changed files with 116 additions and 47 deletions

View file

@ -593,6 +593,19 @@ func (ws WeakString) String() string {
return string(ws) return string(ws)
} }
// CopyHeader copies HTTP headers by completely
// replacing dest with src. (This allows deletions
// to be propagated, assuming src started as a
// consistent copy of dest.)
func CopyHeader(dest, src http.Header) {
for field := range dest {
delete(dest, field)
}
for field, val := range src {
dest[field] = val
}
}
// StatusCodeMatches returns true if a real HTTP status code matches // StatusCodeMatches returns true if a real HTTP status code matches
// the configured status code, which may be either a real HTTP status // the configured status code, which may be either a real HTTP status
// code or an integer representing a class of codes (e.g. 4 for all // code or an integer representing a class of codes (e.g. 4 for all

View file

@ -130,8 +130,7 @@ func (c *Cache) getter(ctx groupcache.Context, key string, dest groupcache.Sink)
// we need to record the response if we are to cache it; only cache if // we need to record the response if we are to cache it; only cache if
// request is successful (TODO: there's probably much more nuance needed here) // request is successful (TODO: there's probably much more nuance needed here)
var rr caddyhttp.ResponseRecorder rr := caddyhttp.NewResponseRecorder(combo.rw, buf, func(status int, header http.Header) bool {
rr = caddyhttp.NewResponseRecorder(combo.rw, buf, func(status int) bool {
shouldBuf := status < 300 shouldBuf := status < 300
if shouldBuf { if shouldBuf {
@ -141,7 +140,7 @@ func (c *Cache) getter(ctx groupcache.Context, key string, dest groupcache.Sink)
// the rest will be the body, which will be written // the rest will be the body, which will be written
// implicitly for us by the recorder // implicitly for us by the recorder
err := gob.NewEncoder(buf).Encode(headerAndStatus{ err := gob.NewEncoder(buf).Encode(headerAndStatus{
Header: rr.Header(), Header: header,
Status: status, Status: status,
}) })
if err != nil { if err != nil {

View file

@ -48,8 +48,8 @@ func (m Markdown) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyht
buf.Reset() buf.Reset()
defer bufPool.Put(buf) defer bufPool.Put(buf)
shouldBuf := func(status int) bool { shouldBuf := func(status int, header http.Header) bool {
return strings.HasPrefix(w.Header().Get("Content-Type"), "text/") return strings.HasPrefix(header.Get("Content-Type"), "text/")
} }
rec := caddyhttp.NewResponseRecorder(w, buf, shouldBuf) rec := caddyhttp.NewResponseRecorder(w, buf, shouldBuf)
@ -62,6 +62,8 @@ func (m Markdown) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyht
return nil return nil
} }
caddyhttp.CopyHeader(w.Header(), rec.Header())
output := blackfriday.Run(buf.Bytes()) output := blackfriday.Run(buf.Bytes())
w.Header().Set("Content-Length", strconv.Itoa(len(output))) w.Header().Set("Content-Length", strconv.Itoa(len(output)))

View file

@ -18,6 +18,7 @@ import (
"bufio" "bufio"
"bytes" "bytes"
"fmt" "fmt"
"io"
"net" "net"
"net/http" "net/http"
) )
@ -78,52 +79,89 @@ type responseRecorder struct {
wroteHeader bool wroteHeader bool
statusCode int statusCode int
buf *bytes.Buffer buf *bytes.Buffer
shouldBuffer func(status int) bool shouldBuffer ShouldBufferFunc
stream bool stream bool
size int size int
header http.Header
} }
// NewResponseRecorder returns a new ResponseRecorder that can be // NewResponseRecorder returns a new ResponseRecorder that can be
// used instead of a real http.ResponseWriter. The recorder is useful // used instead of a standard http.ResponseWriter. The recorder is
// for middlewares which need to buffer a responder's response and // useful for middlewares which need to buffer a response and
// process it in its entirety before actually allowing the response to // potentially process its entire body before actually writing the
// be written. Of course, this has a performance overhead, but // response to the underlying writer. Of course, buffering the entire
// sometimes there is no way to avoid buffering the whole response. // body has a memory overhead, but sometimes there is no way to avoid
// Still, if at all practical, middlewares should strive to stream // 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 // responses by wrapping Write and WriteHeader methods instead of
// buffering whole response bodies. // buffering whole response bodies.
// //
// Recorders optionally buffer the response. When the headers are // Buffering is actually optional. The shouldBuffer function will
// to be written, shouldBuffer will be called with the status // be called just before the headers are written. If it returns
// code that is being written. The rest of the headers can be read // true, the headers and body will be buffered by this recorder
// from w.Header(). If shouldBuffer returns true, the response // and not written to the underlying writer; if false, the headers
// will be buffered. You can know the response was buffered if // will be written immediately and the body will be streamed out
// the Buffered() method returns true. If the response was not // directly to the underlying writer. If shouldBuffer is nil,
// buffered, Buffered() will return false and that means the // the response will never be buffered and will always be streamed
// response bypassed the recorder and was written directly to the // directly to the writer.
// 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.
// //
// Before calling this function in a middleware handler, make a // You can know if shouldBuffer returned true by calling Buffered().
// 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.
// //
// The returned recorder can be used in place of w when calling // The provided buffer buf should be obtained from a pool for best
// the next handler in the chain. When that handler returns, you // performance (see the sync.Pool type).
// can read the status code from the recorder's Status() method. //
// The response body fills buf if it was buffered, and the headers // Proper usage of a recorder looks like this:
// are available via w.Header(). //
func NewResponseRecorder(w http.ResponseWriter, buf *bytes.Buffer, shouldBuffer func(status int) bool) ResponseRecorder { // 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{ return &responseRecorder{
ResponseWriterWrapper: &ResponseWriterWrapper{ResponseWriter: w}, ResponseWriterWrapper: &ResponseWriterWrapper{ResponseWriter: w},
buf: buf, buf: buf,
shouldBuffer: shouldBuffer, shouldBuffer: shouldBuffer,
header: hdr,
} }
} }
func (rr *responseRecorder) Header() http.Header {
return rr.header
}
func (rr *responseRecorder) WriteHeader(statusCode int) { func (rr *responseRecorder) WriteHeader(statusCode int) {
if rr.wroteHeader { if rr.wroteHeader {
return return
@ -135,9 +173,12 @@ func (rr *responseRecorder) WriteHeader(statusCode int) {
if rr.shouldBuffer == nil { if rr.shouldBuffer == nil {
rr.stream = true rr.stream = true
} else { } else {
rr.stream = !rr.shouldBuffer(rr.statusCode) rr.stream = !rr.shouldBuffer(rr.statusCode, rr.header)
} }
// if not buffered, immediately write header
if rr.stream { if rr.stream {
CopyHeader(rr.ResponseWriterWrapper.Header(), rr.header)
rr.ResponseWriterWrapper.WriteHeader(rr.statusCode) rr.ResponseWriterWrapper.WriteHeader(rr.statusCode)
} }
} }
@ -179,16 +220,32 @@ func (rr *responseRecorder) Buffered() bool {
return !rr.stream 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 // 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 { type ResponseRecorder interface {
HTTPInterfaces HTTPInterfaces
Status() int Status() int
Buffer() *bytes.Buffer Buffer() *bytes.Buffer
Buffered() bool Buffered() bool
Size() int 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 // Interface guards
var ( var (
_ HTTPInterfaces = (*ResponseWriterWrapper)(nil) _ HTTPInterfaces = (*ResponseWriterWrapper)(nil)

View file

@ -17,7 +17,6 @@ package templates
import ( import (
"bytes" "bytes"
"fmt" "fmt"
"io"
"net/http" "net/http"
"strconv" "strconv"
"strings" "strings"
@ -71,8 +70,8 @@ func (t *Templates) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddy
// shouldBuf determines whether to execute templates on this response, // shouldBuf determines whether to execute templates on this response,
// since generally we will not want to execute for images or CSS, etc. // since generally we will not want to execute for images or CSS, etc.
shouldBuf := func(status int) bool { shouldBuf := func(status int, header http.Header) bool {
ct := w.Header().Get("Content-Type") ct := header.Get("Content-Type")
for _, mt := range t.MIMETypes { for _, mt := range t.MIMETypes {
if strings.Contains(ct, mt) { if strings.Contains(ct, mt) {
return true return true
@ -96,18 +95,17 @@ func (t *Templates) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddy
return err return err
} }
w.Header().Set("Content-Length", strconv.Itoa(buf.Len())) rec.Header().Set("Content-Length", strconv.Itoa(buf.Len()))
w.Header().Del("Accept-Ranges") // we don't know ranges for dynamically-created content rec.Header().Del("Accept-Ranges") // we don't know ranges for dynamically-created content
w.Header().Del("Last-Modified") // useless for dynamic content since it's always changing rec.Header().Del("Last-Modified") // useless for dynamic content since it's always changing
// we don't know a way to guickly generate etag for dynamic content, // we don't know a way to guickly generate etag for dynamic content,
// but we can convert this to a weak etag to kind of indicate that // but we can convert this to a weak etag to kind of indicate that
if etag := w.Header().Get("ETag"); etag != "" { if etag := rec.Header().Get("Etag"); etag != "" {
w.Header().Set("ETag", "W/"+etag) rec.Header().Set("Etag", "W/"+etag)
} }
w.WriteHeader(rec.Status()) rec.WriteResponse()
io.Copy(w, buf)
return nil return nil
} }

View file

@ -80,7 +80,7 @@ func (c templateContext) Include(filename string, args ...interface{}) (template
// If it is not trusted, be sure to use escaping functions yourself. // If it is not trusted, be sure to use escaping functions yourself.
func (c templateContext) HTTPInclude(uri string) (template.HTML, error) { func (c templateContext) HTTPInclude(uri string) (template.HTML, error) {
if c.Req.Header.Get(recursionPreventionHeader) == "1" { if c.Req.Header.Get(recursionPreventionHeader) == "1" {
return "", fmt.Errorf("virtual include cycle") return "", fmt.Errorf("virtual request cycle")
} }
buf := bufPool.Get().(*bytes.Buffer) buf := bufPool.Get().(*bytes.Buffer)