net/context/ctxhttp: fix case where Body could leak and not be closed

Fixes golang/go#14065

Change-Id: Ic19a0f740cddced8fb782f65cea14da383b047b1
Reviewed-on: https://go-review.googlesource.com/18802
Reviewed-by: Olivier Poitrey <[email protected]>
Reviewed-by: Daniel Morsing <[email protected]>
Reviewed-by: Chris Broadfoot <[email protected]>
Run-TryBot: Brad Fitzpatrick <[email protected]>

Author: bradfitz

context/ctxhttp/ctxhttp.go

@@ -14,6 +14,14 @@ import (
 	"golang.org/x/net/context"
 )
 
+func nop() {}
+
+var (
+	testHookContextDoneBeforeHeaders = nop
+	testHookDoReturned               = nop
+	testHookDidBodyClose             = nop
+)
+
 // Do sends an HTTP request with the provided http.Client and returns an HTTP response.
 // If the client is nil, http.DefaultClient is used.
 // If the context is canceled or times out, ctx.Err() will be returned.
@@ -33,14 +41,23 @@ func Do(ctx context.Context, client *http.Client, req *http.Request) (*http.Resp
 
 	go func() {
 		resp, err := client.Do(req)
+		testHookDoReturned()
 		result <- responseAndError{resp, err}
 	}()
 
 	var resp *http.Response
 
 	select {
 	case <-ctx.Done():
+		testHookContextDoneBeforeHeaders()
 		cancel()
+		// Clean up after the goroutine calling client.Do:
+		go func() {
+			if r := <-result; r.resp != nil {
+				testHookDidBodyClose()
+				r.resp.Body.Close()
+			}
+		}()
 		return nil, ctx.Err()
 	case r := <-result:
 		var err error

context/ctxhttp/ctxhttp_test.go

@@ -8,8 +8,10 @@ package ctxhttp
 
 import (
 	"io/ioutil"
+	"net"
 	"net/http"
 	"net/http/httptest"
+	"sync"
 	"testing"
 	"time"
 
@@ -111,3 +113,64 @@ func doRequest(ctx context.Context) (*http.Response, error) {
 
 	return Get(ctx, nil, serv.URL)
 }
+
+// golang.org/issue/14065
+func TestClosesResponseBodyOnCancel(t *testing.T) {
+	defer func() { testHookContextDoneBeforeHeaders = nop }()
+	defer func() { testHookDoReturned = nop }()
+	defer func() { testHookDidBodyClose = nop }()
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	defer ts.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	// closed when Do enters select case <-ctx.Done()
+	enteredDonePath := make(chan struct{})
+
+	testHookContextDoneBeforeHeaders = func() {
+		close(enteredDonePath)
+	}
+
+	testHookDoReturned = func() {
+		// We now have the result (the Flush'd headers) at least,
+		// so we can cancel the request.
+		cancel()
+
+		// But block the client.Do goroutine from sending
+		// until Do enters into the <-ctx.Done() path, since
+		// otherwise if both channels are readable, select
+		// picks a random one.
+		<-enteredDonePath
+	}
+
+	sawBodyClose := make(chan struct{})
+	testHookDidBodyClose = func() { close(sawBodyClose) }
+
+	tr := &http.Transport{}
+	defer tr.CloseIdleConnections()
+	c := &http.Client{Transport: tr}
+	req, _ := http.NewRequest("GET", ts.URL, nil)
+	_, doErr := Do(ctx, c, req)
+
+	select {
+	case <-sawBodyClose:
+	case <-time.After(5 * time.Second):
+		t.Fatal("timeout waiting for body to close")
+	}
+
+	if doErr != ctx.Err() {
+		t.Errorf("Do error = %v; want %v", doErr, ctx.Err())
+	}
+}
+
+type noteCloseConn struct {
+	net.Conn
+	onceClose sync.Once
+	closefn   func()
+}
+
+func (c *noteCloseConn) Close() error {
+	c.onceClose.Do(c.closefn)
+	return c.Conn.Close()
+}