From c7d843a8cc905cbd61d34e2d45cdb3301b9c4124 Mon Sep 17 00:00:00 2001
From: Derek Gonyeo <derek@gonyeo.com>
Date: Tue, 22 May 2018 14:46:36 -0700
Subject: [PATCH] internal: don't cancel http contexts until the body has been
 read

Contexts for HTTP requests should not be cancelled until after the HTTP
body has been successfully read, or errors can be encountered.
---
 internal/resource/http.go | 19 ++++++++++---------
 internal/resource/url.go  |  8 ++++++--
 2 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/internal/resource/http.go b/internal/resource/http.go
index ad3170da8..985415aa9 100644
--- a/internal/resource/http.go
+++ b/internal/resource/http.go
@@ -194,13 +194,14 @@ func (f *Fetcher) newHttpClient() {
 	}
 }
 
-// getReaderWithHeader performs an HTTP GET on the provided URL with the provided request header
-// and returns the response body Reader, HTTP status code, and error (if any). By
+// getReaderWithHeader performs an HTTP GET on the provided URL with the
+// provided request header and returns the response body Reader, HTTP status
+// code, a cancel function for the result's context, and error (if any). By
 // default, User-Agent is added to the header but this can be overridden.
-func (c HttpClient) getReaderWithHeader(ctx context.Context, url string, header http.Header) (io.ReadCloser, int, error) {
+func (c HttpClient) getReaderWithHeader(url string, header http.Header) (io.ReadCloser, int, context.CancelFunc, error) {
 	req, err := http.NewRequest("GET", url, nil)
 	if err != nil {
-		return nil, 0, err
+		return nil, 0, nil, err
 	}
 
 	req.Header.Set("User-Agent", "Ignition/"+version.Raw)
@@ -212,10 +213,10 @@ func (c HttpClient) getReaderWithHeader(ctx context.Context, url string, header
 		}
 	}
 
+	ctx, cancelFn := context.WithCancel(context.Background())
 	if c.timeout != 0 {
-		ctxTo, cancel := context.WithTimeout(ctx, c.timeout)
-		ctx = ctxTo
-		defer cancel()
+		cancelFn()
+		ctx, cancelFn = context.WithTimeout(context.Background(), c.timeout)
 	}
 
 	duration := initialBackoff
@@ -226,7 +227,7 @@ func (c HttpClient) getReaderWithHeader(ctx context.Context, url string, header
 		if err == nil {
 			c.logger.Info("GET result: %s", http.StatusText(resp.StatusCode))
 			if resp.StatusCode < 500 {
-				return resp.Body, resp.StatusCode, nil
+				return resp.Body, resp.StatusCode, cancelFn, nil
 			}
 			resp.Body.Close()
 		} else {
@@ -242,7 +243,7 @@ func (c HttpClient) getReaderWithHeader(ctx context.Context, url string, header
 		select {
 		case <-time.After(duration):
 		case <-ctx.Done():
-			return nil, 0, ErrTimeout
+			return nil, 0, cancelFn, ErrTimeout
 		}
 	}
 }
diff --git a/internal/resource/url.go b/internal/resource/url.go
index 49cc59206..943f35793 100644
--- a/internal/resource/url.go
+++ b/internal/resource/url.go
@@ -218,11 +218,15 @@ func (f *Fetcher) FetchFromTFTP(u url.URL, dest *os.File, opts FetchOptions) err
 // FetchFromHTTP fetches a resource from u via HTTP(S) into dest, returning an
 // error if one is encountered.
 func (f *Fetcher) FetchFromHTTP(u url.URL, dest *os.File, opts FetchOptions) error {
-	ctx := context.Background()
 	if f.client == nil {
 		f.newHttpClient()
 	}
-	dataReader, status, err := f.client.getReaderWithHeader(ctx, u.String(), opts.Headers)
+	dataReader, status, ctxCancel, err := f.client.getReaderWithHeader(u.String(), opts.Headers)
+	if ctxCancel != nil {
+		// whatever context getReaderWithHeader created for the request should
+		// be cancelled once we're done reading the response
+		defer ctxCancel()
+	}
 	if err != nil {
 		return err
 	}