1 files changed, 60 insertions, 9 deletions

diff --git a/modules/caddyhttp/reverseproxy/reverseproxy.go b/modules/caddyhttp/reverseproxy/reverseproxy.go
index 3a2457f..76506ca 100644
--- a/modules/caddyhttp/reverseproxy/reverseproxy.go
+++ b/modules/caddyhttp/reverseproxy/reverseproxy.go
@@ -21,7 +21,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"io/ioutil"
 	"net"
 	"net/http"
 	"regexp"
@@ -93,9 +92,20 @@ type Handler struct {
 
 	// If true, the entire request body will be read and buffered
 	// in memory before being proxied to the backend. This should
-	// be avoided if at all possible for performance reasons.
+	// be avoided if at all possible for performance reasons, but
+	// could be useful if the backend is intolerant of read latency.
 	BufferRequests bool `json:"buffer_requests,omitempty"`
 
+	// If true, the entire response body will be read and buffered
+	// in memory before being proxied to the client. This should
+	// be avoided if at all possible for performance reasons, but
+	// could be useful if the backend has tighter memory constraints.
+	BufferResponses bool `json:"buffer_responses,omitempty"`
+
+	// If body buffering is enabled, the maximum size of the buffers
+	// used for the requests and responses (in bytes).
+	MaxBufferSize int64 `json:"max_buffer_size,omitempty"`
+
 	// List of handlers and their associated matchers to evaluate
 	// after successful roundtrips. The first handler that matches
 	// the response from a backend will be invoked. The response
@@ -337,12 +347,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyht
 	// required, if read timeouts are set,
 	// and if body size is limited
 	if h.BufferRequests {
-		buf := bufPool.Get().(*bytes.Buffer)
-		buf.Reset()
-		defer bufPool.Put(buf)
-		_, _ = io.Copy(buf, r.Body)
-		r.Body.Close()
-		r.Body = ioutil.NopCloser(buf)
+		r.Body = h.bufferedBody(r.Body)
 	}
 
 	// prepare the request for proxying; this is needed only once
@@ -563,6 +568,11 @@ func (h *Handler) reverseProxy(rw http.ResponseWriter, req *http.Request, di Dia
 		}
 	}
 
+	// if enabled, buffer the response body
+	if h.BufferResponses {
+		res.Body = h.bufferedBody(res.Body)
+	}
+
 	// see if any response handler is configured for this response from the backend
 	for i, rh := range h.HandleResponse {
 		if rh.Match != nil && !rh.Match.Match(res.StatusCode, res.Header) {
@@ -599,7 +609,7 @@ func (h *Handler) reverseProxy(rw http.ResponseWriter, req *http.Request, di Dia
 		}
 	}
 
-	// Deal with 101 Switching Protocols responses: (WebSocket, h2c, etc)
+	// deal with 101 Switching Protocols responses: (WebSocket, h2c, etc)
 	if res.StatusCode == http.StatusSwitchingProtocols {
 		h.handleUpgradeResponse(rw, req, res)
 		return nil
@@ -735,6 +745,30 @@ func (h Handler) directRequest(req *http.Request, di DialInfo) {
 	req.URL.Host = reqHost
 }
 
+// bufferedBody reads originalBody into a buffer, then returns a reader for the buffer.
+// Always close the return value when done with it, just like if it was the original body!
+func (h Handler) bufferedBody(originalBody io.ReadCloser) io.ReadCloser {
+	buf := bufPool.Get().(*bytes.Buffer)
+	buf.Reset()
+	if h.MaxBufferSize > 0 {
+		n, err := io.CopyN(buf, originalBody, h.MaxBufferSize)
+		if err != nil || n == h.MaxBufferSize {
+			return bodyReadCloser{
+				Reader: io.MultiReader(buf, originalBody),
+				buf:    buf,
+				body:   originalBody,
+			}
+		}
+	} else {
+		_, _ = io.Copy(buf, originalBody)
+	}
+	originalBody.Close() // no point in keeping it open
+	return bodyReadCloser{
+		Reader: buf,
+		buf:    buf,
+	}
+}
+
 func copyHeader(dst, src http.Header) {
 	for k, vv := range src {
 		for _, v := range vv {
@@ -858,6 +892,23 @@ type TLSTransport interface {
 // roundtrip succeeded, but an error occurred after-the-fact.
 type roundtripSucceeded struct{ error }
 
+// bodyReadCloser is a reader that, upon closing, will return
+// its buffer to the pool and close the underlying body reader.
+type bodyReadCloser struct {
+	io.Reader
+	buf  *bytes.Buffer
+	body io.ReadCloser
+}
+
+func (brc bodyReadCloser) Close() error {
+	bufPool.Put(brc.buf)
+	if brc.body != nil {
+		return brc.body.Close()
+	}
+	return nil
+}
+
+// bufPool is used for buffering requests and responses.
 var bufPool = sync.Pool{
 	New: func() interface{} {
 		return new(bytes.Buffer)