From 54f0ac4c46184ae44486a31ca2705076abcc5321 Mon Sep 17 00:00:00 2001
From: ale <ale@incal.net>
Date: Sun, 30 Jun 2019 09:30:15 +0100
Subject: [PATCH] Prevent premature cancellation of the HTTP request context

The Context in the http.Request is bound to the one we give it when
calling balancedBackend.do(), so we were accessing it out of scope
when calling the final json.Decode on the response body.

This was not a problem for most *small* requests, as the response Body
already contains all the data due to having read it along with the
headers. However, larger response bodies would cause json.Decode to
call a Read with what at that point is a canceled Context, so the Call
function would return a mysterious "context canceled" error.

(Note that this change introduces a minor, probably less annoying,
issue, where we can't reset "resp" between successive calls so we may
be invoking json.Decode on a non-pristine object if the first call
fails in some circumstances).
---
 clientutil/backend_test.go |  4 ++--
 clientutil/balancer.go     | 37 +++++++++++++++++++------------------
 2 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/clientutil/backend_test.go b/clientutil/backend_test.go
index 561f5be..6b4e3d8 100644
--- a/clientutil/backend_test.go
+++ b/clientutil/backend_test.go
@@ -90,7 +90,7 @@ func newErrorHTTPServer(statusCode int) *httpServer {
 func newJSONHTTPServer() *httpServer {
 	return &httpServer{httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "application/json")
-		io.WriteString(w, "{\"value\": 42}")
+		io.WriteString(w, "{\"value\": 42}") // nolint
 	}))}
 }
 
@@ -99,7 +99,7 @@ func newHostCountingJSONHTTPServer() (*httpServer, map[string]int) {
 	return &httpServer{httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		counters[r.Host]++
 		w.Header().Set("Content-Type", "application/json")
-		io.WriteString(w, "{\"value\": 42}")
+		io.WriteString(w, "{\"value\": 42}") // nolint
 	}))}, counters
 }
 
diff --git a/clientutil/balancer.go b/clientutil/balancer.go
index 84633ac..d2ca827 100644
--- a/clientutil/balancer.go
+++ b/clientutil/balancer.go
@@ -118,30 +118,31 @@ func (b *balancedBackend) Call(ctx context.Context, shard, path string, req, res
 
 	// Call the backends in the sequence until one succeeds, with an
 	// exponential backoff policy controlled by the outer Context.
-	var httpResp *http.Response
-	err = backoff.Retry(func() error {
+	return backoff.Retry(func() error {
 		req, rerr := b.newJSONRequest(path, shard, data)
 		if rerr != nil {
 			return rerr
 		}
 		innerCtx, cancel := context.WithTimeout(ctx, innerTimeout)
-		httpResp, rerr = b.do(innerCtx, seq, req)
-		cancel()
-		return rerr
-	}, backoff.WithContext(newExponentialBackOff(), ctx))
-	if err != nil {
-		return err
-	}
-	defer httpResp.Body.Close() // nolint
+		defer cancel()
 
-	// Decode the response.
-	if httpResp.Header.Get("Content-Type") != "application/json" {
-		return errors.New("not a JSON response")
-	}
-	if resp == nil {
-		return nil
-	}
-	return json.NewDecoder(httpResp.Body).Decode(resp)
+		// When do() returns successfully, we already know that the
+		// response had an HTTP status of 200.
+		httpResp, rerr := b.do(innerCtx, seq, req)
+		if rerr != nil {
+			return rerr
+		}
+		defer httpResp.Body.Close() // nolint
+
+		// Decode the response, unless the 'resp' output is nil.
+		if httpResp.Header.Get("Content-Type") != "application/json" {
+			return errors.New("not a JSON response")
+		}
+		if resp == nil {
+			return nil
+		}
+		return json.NewDecoder(httpResp.Body).Decode(resp)
+	}, backoff.WithContext(newExponentialBackOff(), ctx))
 }
 
 // Initialize a new target sequence.
-- 
GitLab