From 590862a96230361f36e84257725fdca9450a45a9 Mon Sep 17 00:00:00 2001
From: Tw <tw19881113@gmail.com>
Date: Fri, 19 Aug 2016 13:09:22 +0800
Subject: [PATCH] replacer: capture request body normally

fix issue #1015

Signed-off-by: Tw <tw19881113@gmail.com>
---
 caddyhttp/httpserver/replacer.go      | 63 +++++++++++++++++----------
 caddyhttp/httpserver/replacer_test.go | 27 ------------
 2 files changed, 40 insertions(+), 50 deletions(-)

diff --git a/caddyhttp/httpserver/replacer.go b/caddyhttp/httpserver/replacer.go
index c77230e86..b31011a92 100644
--- a/caddyhttp/httpserver/replacer.go
+++ b/caddyhttp/httpserver/replacer.go
@@ -44,6 +44,36 @@ type replacer struct {
 	emptyValue         string
 	responseRecorder   *ResponseRecorder
 	request            *http.Request
+	requestBody        *limitWriter
+}
+
+type limitWriter struct {
+	w      bytes.Buffer
+	remain int
+}
+
+func newLimitWriter(max int) *limitWriter {
+	return &limitWriter{
+		w:      bytes.Buffer{},
+		remain: max,
+	}
+}
+
+func (lw *limitWriter) Write(p []byte) (int, error) {
+	// skip if we are full
+	if lw.remain <= 0 {
+		return len(p), nil
+	}
+	if n := len(p); n > lw.remain {
+		p = p[:lw.remain]
+	}
+	n, err := lw.w.Write(p)
+	lw.remain -= n
+	return n, err
+}
+
+func (lw *limitWriter) String() string {
+	return lw.w.String()
 }
 
 // NewReplacer makes a new replacer based on r and rr which
@@ -54,8 +84,16 @@ type replacer struct {
 // emptyValue should be the string that is used in place
 // of empty string (can still be empty string).
 func NewReplacer(r *http.Request, rr *ResponseRecorder, emptyValue string) Replacer {
+	rb := newLimitWriter(MaxLogBodySize)
+	if r.Body != nil {
+		r.Body = struct {
+			io.Reader
+			io.Closer
+		}{io.TeeReader(r.Body, rb), io.Closer(r.Body)}
+	}
 	rep := &replacer{
 		request:            r,
+		requestBody:        rb,
 		responseRecorder:   rr,
 		customReplacements: make(map[string]string),
 		emptyValue:         emptyValue,
@@ -81,27 +119,6 @@ func canLogRequest(r *http.Request) bool {
 	return false
 }
 
-// readRequestBody reads the request body and sets a
-// new io.ReadCloser that has not yet been read.
-func readRequestBody(r *http.Request, n int64) ([]byte, error) {
-	defer r.Body.Close()
-
-	body, err := ioutil.ReadAll(io.LimitReader(r.Body, n))
-	if err != nil {
-		return nil, err
-	}
-
-	// Read the remaining bytes
-	remaining, err := ioutil.ReadAll(r.Body)
-	if err != nil {
-		return nil, err
-	}
-
-	buf := bytes.NewBuffer(append(body, remaining...))
-	r.Body = ioutil.NopCloser(buf)
-	return body, nil
-}
-
 // Replace performs a replacement of values on s and returns
 // the string with the replaced values.
 func (r *replacer) Replace(s string) string {
@@ -249,11 +266,11 @@ func (r *replacer) getSubstitution(key string) string {
 		if !canLogRequest(r.request) {
 			return r.emptyValue
 		}
-		body, err := readRequestBody(r.request, maxLogBodySize)
+		_, err := ioutil.ReadAll(r.request.Body)
 		if err != nil {
 			return r.emptyValue
 		}
-		return requestReplacer.Replace(string(body))
+		return requestReplacer.Replace(r.requestBody.String())
 	case "{status}":
 		if r.responseRecorder == nil {
 			return r.emptyValue
diff --git a/caddyhttp/httpserver/replacer_test.go b/caddyhttp/httpserver/replacer_test.go
index 25d39d07b..9fcff8c90 100644
--- a/caddyhttp/httpserver/replacer_test.go
+++ b/caddyhttp/httpserver/replacer_test.go
@@ -1,8 +1,6 @@
 package httpserver
 
 import (
-	"bytes"
-	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -149,28 +147,3 @@ func TestRound(t *testing.T) {
 		}
 	}
 }
-
-func TestReadRequestBody(t *testing.T) {
-	payload := []byte(`{ "foo": "bar" }`)
-	var readSize int64 = 5
-	r, err := http.NewRequest("POST", "/", bytes.NewReader(payload))
-	if err != nil {
-		t.Error(err)
-	}
-	defer r.Body.Close()
-
-	logBody, err := readRequestBody(r, readSize)
-	if err != nil {
-		t.Error("readRequestBody failed", err)
-	} else if !bytes.EqualFold(payload[0:readSize], logBody) {
-		t.Error("Expected log comparison failed")
-	}
-
-	// Ensure the Request body is the same as the original.
-	reqBody, err := ioutil.ReadAll(r.Body)
-	if err != nil {
-		t.Error("Unable to read request body", err)
-	} else if !bytes.EqualFold(payload, reqBody) {
-		t.Error("Expected request body comparison failed")
-	}
-}