proxy: clone request Body for retry (fixes #1229)

caddyhttp/proxy/clone.go

@@ -0,0 +1,61 @@
+package proxy
+
+import (
+	"bytes"
+	"io"
+	"io/ioutil"
+	"sync"
+)
+
+// clone returns an origin and cloned io.ReadCloser from src. origin should
+// be used in place of src. Close on origin closes src. Read on cloned returns
+// io.EOF only when origin is closed. origin and cloned can be read together.
+func clone(src io.ReadCloser) (origin io.ReadCloser, cloned io.ReadCloser) {
+	buf := new(bytes.Buffer)
+	r := io.TeeReader(src, buf)
+	ch := make(chan struct{})
+	var once sync.Once
+	origin = &originBody{
+		Reader: r,
+		closeFunc: func() error {
+			// ensure all data is read into buf
+			_, _ = io.Copy(ioutil.Discard, r)
+			err := src.Close()
+			once.Do(func() {
+				close(ch)
+			})
+			return err
+		},
+	}
+	cloned = clonedBody(func(p []byte) (int, error) {
+		n, err := buf.Read(p)
+		select {
+		case <-ch:
+		default:
+			if err == io.EOF {
+				err = nil
+			}
+		}
+		return n, err
+	})
+	return origin, cloned
+}
+
+type clonedBody func([]byte) (int, error)
+
+func (c clonedBody) Read(p []byte) (int, error) {
+	return c(p)
+}
+
+func (clonedBody) Close() error {
+	return nil
+}
+
+type originBody struct {
+	io.Reader
+	closeFunc func() error
+}
+
+func (o *originBody) Close() error {
+	return o.closeFunc()
+}

caddyhttp/proxy/clone_test.go

@@ -0,0 +1,93 @@
+package proxy
+
+import (
+	"bytes"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+type testcase struct {
+	io.Reader
+	closedCalled bool
+}
+
+func (t *testcase) Close() error {
+	t.closedCalled = true
+	return nil
+}
+
+func TestClone(t *testing.T) {
+	expected := "test content"
+	tc := &testcase{Reader: bytes.NewBufferString(expected)}
+	o, c := clone(tc)
+	p := make([]byte, len(expected))
+	if _, err := o.Read(p); err != nil {
+		t.Fatal(err)
+	}
+	if string(p) != expected {
+		t.Fatalf("string(p) = %s, want %s", string(p), expected)
+	}
+	q := make([]byte, len(expected))
+	if _, err := c.Read(q); err != nil {
+		t.Fatal(err)
+	}
+	if string(q) != expected {
+		t.Fatalf("string(q) = %s, want %s", string(q), expected)
+	}
+	if _, err := c.Read(q); err != nil {
+		t.Fatalf("err = %v, want nil until o is closed", err)
+	}
+	if err := o.Close(); err != nil {
+		t.Fatal(err)
+	}
+	if tc.closedCalled == false {
+		t.Fatalf("tc.closedCalled = false, want true")
+	}
+	if _, err := c.Read(q); err != io.EOF {
+		t.Fatalf("err = %v, want io.EOF", err)
+	}
+	if err := o.Close(); err != nil {
+		t.Fatalf("err = %v, want nil on redundant close", err)
+	}
+}
+
+func TestCloneRetry(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		io.Copy(w, r.Body)
+		r.Body.Close()
+	}))
+	defer ts.Close()
+
+	testcase := "test content"
+	req, err := http.NewRequest(http.MethodPost, ts.URL, bytes.NewBufferString(testcase))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	o, c := clone(req.Body)
+
+	// simulate fail request
+	req.Body = o
+	host := req.URL.Host
+	req.URL.Host = "example"
+	_, _ = http.DefaultTransport.RoundTrip(req)
+
+	// retry request
+	req.Body = c
+	req.URL.Host = host
+	resp, err := http.DefaultTransport.RoundTrip(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer resp.Body.Close()
+	result, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if string(result) != testcase {
+		t.Fatalf("result = %s, want %s", result, testcase)
+	}
+}

caddyhttp/proxy/proxy.go

@@ -222,6 +222,9 @@ func createUpstreamRequest(r *http.Request) *http.Request {
 	// For server requests the Request Body is always non-nil.
 	if r.ContentLength == 0 {
 		outreq.Body = nil
+	} else {
+		// clone request Body for reuse
+		outreq.Body, r.Body = clone(r.Body)
 	}
 
 	// Restore URL Path if it has been modified