caddyhttp: ResponseRecorder type for middlewares to buffer responses

Unfortunately, templates and markdown require buffering the full
response before it can be processed and written to the client

3 files changed, 110 insertions, 59 deletions

diff --git a/modules/caddyhttp/markdown/markdown.go b/modules/caddyhttp/markdown/markdown.go
index 90d615a..574039e 100644
--- a/modules/caddyhttp/markdown/markdown.go
+++ b/modules/caddyhttp/markdown/markdown.go
@@ -1,8 +1,10 @@
 package markdown
 
 import (
+	"bytes"
 	"net/http"
 	"strconv"
+	"sync"
 
 	"gopkg.in/russross/blackfriday.v2"
 
@@ -22,31 +24,35 @@ type Markdown struct {
 }
 
 func (m Markdown) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
-	mrw := &markdownResponseWriter{
-		ResponseWriterWrapper: &caddyhttp.ResponseWriterWrapper{ResponseWriter: w},
+	buf := bufPool.Get().(*bytes.Buffer)
+	buf.Reset()
+	defer bufPool.Put(buf)
+
+	rr := caddyhttp.NewResponseRecorder(w, buf)
+
+	err := next.ServeHTTP(rr, r)
+	if err != nil {
+		return err
 	}
-	return next.ServeHTTP(mrw, r)
-}
 
-type markdownResponseWriter struct {
-	*caddyhttp.ResponseWriterWrapper
-	statusCode  int
-	wroteHeader bool
-}
+	output := blackfriday.Run(buf.Bytes())
 
-func (mrw *markdownResponseWriter) WriteHeader(code int) {
-	mrw.statusCode = code
+	w.Header().Set("Content-Length", strconv.Itoa(len(output)))
+	w.Header().Set("Content-Type", "text/html; charset=utf-8")
+	w.Header().Del("Accept-Ranges") // we don't know ranges for dynamically-created content
+	w.Header().Del("Etag")          // don't know a way to quickly generate etag for dynamic content
+	w.Header().Del("Last-Modified") // useless for dynamic content since it's always changing
+
+	w.WriteHeader(rr.Status())
+	w.Write(output)
+
+	return nil
 }
 
-func (mrw *markdownResponseWriter) Write(d []byte) (int, error) {
-	output := blackfriday.Run(d)
-	if !mrw.wroteHeader {
-		mrw.Header().Set("Content-Length", strconv.Itoa(len(output)))
-		mrw.Header().Set("Content-Type", "text/html; charset=utf-8")
-		mrw.WriteHeader(mrw.statusCode)
-		mrw.wroteHeader = true
-	}
-	return mrw.ResponseWriter.Write(output)
+var bufPool = sync.Pool{
+	New: func() interface{} {
+		return new(bytes.Buffer)
+	},
 }
 
 // Interface guard
diff --git a/modules/caddyhttp/responsewriter.go b/modules/caddyhttp/responsewriter.go
index 3bf3965..7fb45f3 100644
--- a/modules/caddyhttp/responsewriter.go
+++ b/modules/caddyhttp/responsewriter.go
@@ -2,6 +2,7 @@ package caddyhttp
 
 import (
 	"bufio"
+	"bytes"
 	"fmt"
 	"net"
 	"net/http"
@@ -58,5 +59,75 @@ type HTTPInterfaces interface {
 // ResponseWriter does not implement the required method.
 var ErrNotImplemented = fmt.Errorf("method not implemented")
 
+type responseRecorder struct {
+	*ResponseWriterWrapper
+	wroteHeader bool
+	statusCode  int
+	buf         *bytes.Buffer
+}
+
+// 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
+// responses by wrapping Write and WriteHeader methods instead of
+// buffering whole response bodies.
+//
+// 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.
+//
+// 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, and the headers are available in
+// w.Header().
+func NewResponseRecorder(w http.ResponseWriter, buf *bytes.Buffer) ResponseRecorder {
+	return &responseRecorder{
+		ResponseWriterWrapper: &ResponseWriterWrapper{ResponseWriter: w},
+		buf:                   buf,
+	}
+}
+
+func (rr *responseRecorder) WriteHeader(statusCode int) {
+	if rr.wroteHeader {
+		return
+	}
+	rr.statusCode = statusCode
+	rr.wroteHeader = true
+}
+
+func (rr *responseRecorder) Write(data []byte) (int, error) {
+	rr.WriteHeader(http.StatusOK)
+	return rr.buf.Write(data)
+}
+
+// Status returns the status code that was written, if any.
+func (rr *responseRecorder) Status() int {
+	return rr.statusCode
+}
+
+// Buffer returns the body buffer that rr was created with.
+// You should still have your original pointer, though.
+func (rr *responseRecorder) Buffer() *bytes.Buffer {
+	return rr.buf
+}
+
+// ResponseRecorder is a http.ResponseWriter that records
+// responses instead of writing them to the client.
+type ResponseRecorder interface {
+	HTTPInterfaces
+	Status() int
+	Buffer() *bytes.Buffer
+}
+
 // Interface guards
-var _ HTTPInterfaces = (*ResponseWriterWrapper)(nil)
+var (
+	_ HTTPInterfaces   = (*ResponseWriterWrapper)(nil)
+	_ ResponseRecorder = (*responseRecorder)(nil)
+)
diff --git a/modules/caddyhttp/templates/templates.go b/modules/caddyhttp/templates/templates.go
index a0fb87f..c659f6b 100644
--- a/modules/caddyhttp/templates/templates.go
+++ b/modules/caddyhttp/templates/templates.go
@@ -37,35 +37,31 @@ func (t *Templates) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddy
 	buf.Reset()
 	defer bufPool.Put(buf)
 
-	wb := &responseBuffer{
-		ResponseWriterWrapper: &caddyhttp.ResponseWriterWrapper{ResponseWriter: w},
-		buf:                   buf,
-	}
+	rr := caddyhttp.NewResponseRecorder(w, buf)
 
-	err := next.ServeHTTP(wb, r)
+	err := next.ServeHTTP(rr, r)
 	if err != nil {
 		return err
 	}
 
-	err = t.executeTemplate(wb, r)
+	err = t.executeTemplate(rr, r)
 	if err != nil {
 		return err
 	}
 
-	w.Header().Set("Content-Length", strconv.Itoa(wb.buf.Len()))
+	w.Header().Set("Content-Length", strconv.Itoa(buf.Len()))
 	w.Header().Del("Accept-Ranges") // we don't know ranges for dynamically-created content
 	w.Header().Del("Etag")          // don't know a way to quickly generate etag for dynamic content
 	w.Header().Del("Last-Modified") // useless for dynamic content since it's always changing
 
-	w.WriteHeader(wb.statusCode)
-	io.Copy(w, wb.buf)
+	w.WriteHeader(rr.Status())
+	io.Copy(w, buf)
 
 	return nil
 }
 
-// executeTemplate executes the template contianed
-// in wb.buf and replaces it with the results.
-func (t *Templates) executeTemplate(wb *responseBuffer, r *http.Request) error {
+// executeTemplate executes the template contained in wb.buf and replaces it with the results.
+func (t *Templates) executeTemplate(rr caddyhttp.ResponseRecorder, r *http.Request) error {
 	var fs http.FileSystem
 	if t.FileRoot != "" {
 		fs = http.Dir(t.FileRoot)
@@ -74,11 +70,11 @@ func (t *Templates) executeTemplate(wb *responseBuffer, r *http.Request) error {
 	ctx := &templateContext{
 		Root:       fs,
 		Req:        r,
-		RespHeader: tplWrappedHeader{wb.Header()},
+		RespHeader: tplWrappedHeader{rr.Header()},
 		config:     t,
 	}
 
-	err := ctx.executeTemplateInBuffer(r.URL.Path, wb.buf)
+	err := ctx.executeTemplateInBuffer(r.URL.Path, rr.Buffer())
 	if err != nil {
 		return caddyhttp.Error(http.StatusInternalServerError, err)
 	}
@@ -86,29 +82,8 @@ func (t *Templates) executeTemplate(wb *responseBuffer, r *http.Request) error {
 	return nil
 }
 
-// responseBuffer buffers the response so that it can be
-// executed as a template.
-type responseBuffer struct {
-	*caddyhttp.ResponseWriterWrapper
-	wroteHeader bool
-	statusCode  int
-	buf         *bytes.Buffer
-}
-
-func (rb *responseBuffer) WriteHeader(statusCode int) {
-	if rb.wroteHeader {
-		return
-	}
-	rb.statusCode = statusCode
-	rb.wroteHeader = true
-}
-
-func (rb *responseBuffer) Write(data []byte) (int, error) {
-	rb.WriteHeader(http.StatusOK)
-	return rb.buf.Write(data)
-}
-
-// virtualResponseWriter is used in virtualized HTTP requests.
+// virtualResponseWriter is used in virtualized HTTP requests
+// that templates may execute.
 type virtualResponseWriter struct {
 	status int
 	header http.Header
@@ -131,5 +106,4 @@ func (vrw *virtualResponseWriter) Write(data []byte) (int, error) {
 var (
 	_ caddy.Validator             = (*Templates)(nil)
 	_ caddyhttp.MiddlewareHandler = (*Templates)(nil)
-	_ caddyhttp.HTTPInterfaces    = (*responseBuffer)(nil)
 )