[Heartbeat] Read entire body before closing connection

This fixes an issue where the connection would be closed after only a partial body read. The RoundTripper would close the conn before returning the response which included a partially buffered body. This would actually work for short responses, since the backing bufio would do a partial read, but would fail on all but the shortest responses.

Normally connection lifecycle is handled outside the realm of the `RoundTripper`, but for our purposes we don't want to re-use connections. Since it is a requirement that all response bodies be closed, we can piggy-back on top of that to ensure the connection is closed.

Fixes #8588

heartbeat/hbtest/hbtestutil.go

@@ -26,6 +26,7 @@ import (
 	"net/url"
 	"os"
 	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -51,6 +52,20 @@ func HelloWorldHandler(status int) http.HandlerFunc {
 	)
 }
 
+func LargeResponseHandler(bytes int) http.HandlerFunc {
+	var body strings.Builder
+	for i := 0; i < bytes; i++ {
+		body.WriteString("x")
+	}
+
+	return http.HandlerFunc(
+		func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(200)
+			io.WriteString(w, body.String())
+		},
+	)
+}
+
 // ServerPort takes an httptest.Server and returns its port as a uint16.
 func ServerPort(server *httptest.Server) (uint16, error) {
 	u, err := url.Parse(server.URL)

heartbeat/monitors/active/http/http_test.go

@@ -210,6 +210,41 @@ func TestDownStatuses(t *testing.T) {
 	}
 }
 
+func TestLargeResponse(t *testing.T) {
+	server := httptest.NewServer(hbtest.LargeResponseHandler(1024 * 1024))
+	defer server.Close()
+
+	configSrc := map[string]interface{}{
+		"urls":                server.URL,
+		"timeout":             "1s",
+		"check.response.body": "x",
+	}
+
+	config, err := common.NewConfigFrom(configSrc)
+	require.NoError(t, err)
+
+	jobs, err := create("largeresp", config)
+	require.NoError(t, err)
+
+	job := jobs[0]
+
+	event, _, err := job.Run()
+	require.NoError(t, err)
+
+	port, err := hbtest.ServerPort(server)
+	require.NoError(t, err)
+
+	mapvaltest.Test(
+		t,
+		mapval.Strict(mapval.Compose(
+			hbtest.MonitorChecks("http@"+server.URL, server.URL, "127.0.0.1", "http", "up"),
+			hbtest.RespondingTCPChecks(port),
+			respondingHTTPChecks(server.URL, 200),
+		)),
+		event.Fields,
+	)
+}
+
 func TestHTTPSServer(t *testing.T) {
 	server := httptest.NewTLSServer(hbtest.HelloWorldHandler(http.StatusOK))
 	port, err := hbtest.ServerPort(server)

heartbeat/monitors/active/http/simple_transp.go

@@ -85,7 +85,6 @@ func (t *SimpleTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer conn.Close()
 
 	requestedGzip := false
 	if t.DisableCompression &&
@@ -147,13 +146,30 @@ func (t *SimpleTransport) writeRequest(conn net.Conn, req *http.Request) error {
 	return err
 }
 
+// comboConnReadCloser wraps a ReadCloser that is backed by
+// on a net.Conn. It will close the net.Conn when the ReadCloser is closed.
+type comboConnReadCloser struct {
+	conn net.Conn
+	rc   io.ReadCloser
+}
+
+func (c comboConnReadCloser) Read(p []byte) (n int, err error) {
+	return c.rc.Read(p)
+}
+
+func (c comboConnReadCloser) Close() error {
+	defer c.conn.Close()
+	return c.rc.Close()
+}
+
 func (t *SimpleTransport) readResponse(
 	conn net.Conn,
 	req *http.Request,
 	requestedGzip bool,
 ) (*http.Response, error) {
 	reader := bufio.NewReader(conn)
 	resp, err := http.ReadResponse(reader, req)
+	resp.Body = comboConnReadCloser{conn, resp.Body}
 	if err != nil {
 		return nil, err
 	}