From e7498f9468a2d4e700e381ffea503f943c287ea3 Mon Sep 17 00:00:00 2001
From: Matthew Riley <mattdr@pulumi.com>
Date: Fri, 1 Jun 2018 16:58:12 -0700
Subject: [PATCH] Send request body on retries

Not sending request body on retry is bad enough, but we keep the original
`Content-Length` and the server ends up expecting a body we never send.
This results in `400 Bad Request` errors from the ALB.
---
 pkg/util/httputil/http.go      | 12 +++++
 pkg/util/httputil/http_test.go | 94 ++++++++++++++++++++++++++++++++++
 2 files changed, 106 insertions(+)
 create mode 100644 pkg/util/httputil/http_test.go

diff --git a/pkg/util/httputil/http.go b/pkg/util/httputil/http.go
index b1a688f6d..800c45f0f 100644
--- a/pkg/util/httputil/http.go
+++ b/pkg/util/httputil/http.go
@@ -28,12 +28,24 @@ const maxRetryCount = 5
 
 // DoWithRetry calls client.Do, and in the case of an error, retries the operation again after a slight delay.
 func DoWithRetry(req *http.Request, client *http.Client) (*http.Response, error) {
+	contract.Assertf(req.ContentLength == 0 || req.GetBody != nil,
+		"Retryable request must have no body or rewindable body")
+
 	inRange := func(test, lower, upper int) bool {
 		return lower <= test && test <= upper
 	}
 
 	_, res, err := retry.Until(context.Background(), retry.Acceptor{
 		Accept: func(try int, nextRetryTime time.Duration) (bool, interface{}, error) {
+			if try > 0 && req.GetBody != nil {
+				// Reset request body, if present, for retries.
+				rc, bodyErr := req.GetBody()
+				if bodyErr != nil {
+					return false, nil, bodyErr
+				}
+				req.Body = rc
+			}
+
 			res, resErr := client.Do(req)
 			if resErr == nil && !inRange(res.StatusCode, 500, 599) {
 				return true, res, nil
diff --git a/pkg/util/httputil/http_test.go b/pkg/util/httputil/http_test.go
new file mode 100644
index 000000000..f05dfad91
--- /dev/null
+++ b/pkg/util/httputil/http_test.go
@@ -0,0 +1,94 @@
+// Copyright 2016-2018, Pulumi Corporation.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package httputil
+
+import (
+	"crypto/tls"
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"strings"
+	"testing"
+
+	"golang.org/x/net/http2"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func http2ServerAndClient(handler http.Handler) (*httptest.Server, *http.Client) {
+	// Create an HTTP/2 test server.
+	// httptest.StartTLS will set NextProtos to ["http/1.1"] if it's unset, so we need to add
+	// HTTP/2 eagerly before starting the server.
+	server := httptest.NewUnstartedServer(handler)
+	server.TLS = &tls.Config{
+		NextProtos: []string{http2.NextProtoTLS},
+	}
+	server.StartTLS()
+
+	// Create a client for the test server that will use HTTP/2.
+	// We need a client that will (a) upgrade to HTTP/2 and (b) trust the test server's certs.
+	// In order to satisfy (b), httptest sets Transport to an `http.Transport`, breaking (a),
+	// so we have to manually create an `http2.Transport` and copy over the `tls.Config`.
+	tlsConfig := server.Client().Transport.(*http.Transport).TLSClientConfig
+	client := &http.Client{
+		Transport: &http2.Transport{
+			TLSClientConfig: tlsConfig,
+		},
+	}
+
+	return server, client
+}
+
+// Test that DoWithRetry rewinds and resends the request body when retrying POSTs over HTTP/2.
+func TestRetryPostHTTP2(t *testing.T) {
+	tries := 0
+	handler := func(w http.ResponseWriter, r *http.Request) {
+		tries++
+		t.Logf("try %d", tries)
+
+		assert.Equal(t, "HTTP/2.0", r.Proto)
+
+		// Check that the body's content length matches the sent data.
+		defer r.Body.Close()
+		content, err := ioutil.ReadAll(r.Body)
+		assert.NoError(t, err)
+		assert.Equal(t, strconv.Itoa(len(content)), r.Header.Get("Content-Length"))
+
+		// Check the message matches.
+		assert.Equal(t, string(content), "hello, server")
+
+		// Fail the first try with 500, which will prompt a retry.
+		switch tries {
+		case 1:
+			w.WriteHeader(500)
+		default:
+			w.WriteHeader(200)
+		}
+	}
+
+	server, client := http2ServerAndClient(http.HandlerFunc(handler))
+
+	req, err := http.NewRequest("POST", server.URL, strings.NewReader("hello, server"))
+	assert.NoError(t, err)
+
+	res, err := DoWithRetry(req, client)
+	assert.NoError(t, err)
+	defer res.Body.Close()
+
+	// Check that the request succeeded on the second try.
+	assert.Equal(t, 2, tries)
+	assert.Equal(t, 200, res.StatusCode)
+}