caddyhttp: Improve ResponseRecorder to buffer headers

2 files changed, 9 insertions, 11 deletions

diff --git a/modules/caddyhttp/templates/templates.go b/modules/caddyhttp/templates/templates.go
index 05a2f63..e9c1da8 100644
--- a/modules/caddyhttp/templates/templates.go
+++ b/modules/caddyhttp/templates/templates.go
@@ -17,7 +17,6 @@ package templates
 import (
 	"bytes"
 	"fmt"
-	"io"
 	"net/http"
 	"strconv"
 	"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,
 	// since generally we will not want to execute for images or CSS, etc.
-	shouldBuf := func(status int) bool {
-		ct := w.Header().Get("Content-Type")
+	shouldBuf := func(status int, header http.Header) bool {
+		ct := header.Get("Content-Type")
 		for _, mt := range t.MIMETypes {
 			if strings.Contains(ct, mt) {
 				return true
@@ -96,18 +95,17 @@ func (t *Templates) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddy
 		return err
 	}
 
-	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("Last-Modified") // useless for dynamic content since it's always changing
+	rec.Header().Set("Content-Length", strconv.Itoa(buf.Len()))
+	rec.Header().Del("Accept-Ranges") // we don't know ranges for dynamically-created content
+	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,
 	// but we can convert this to a weak etag to kind of indicate that
-	if etag := w.Header().Get("ETag"); etag != "" {
-		w.Header().Set("ETag", "W/"+etag)
+	if etag := rec.Header().Get("Etag"); etag != "" {
+		rec.Header().Set("Etag", "W/"+etag)
 	}
 
-	w.WriteHeader(rec.Status())
-	io.Copy(w, buf)
+	rec.WriteResponse()
 
 	return nil
 }
diff --git a/modules/caddyhttp/templates/tplcontext.go b/modules/caddyhttp/templates/tplcontext.go
index 5b74623..40d1370 100644
--- a/modules/caddyhttp/templates/tplcontext.go
+++ b/modules/caddyhttp/templates/tplcontext.go
@@ -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.
 func (c templateContext) HTTPInclude(uri string) (template.HTML, error) {
 	if c.Req.Header.Get(recursionPreventionHeader) == "1" {
-		return "", fmt.Errorf("virtual include cycle")
+		return "", fmt.Errorf("virtual request cycle")
 	}
 
 	buf := bufPool.Get().(*bytes.Buffer)