diff --git a/heartbeat/hbtest/hbtestutil.go b/heartbeat/hbtest/hbtestutil.go index ac6e5b0f3ee7..9144ed92b934 100644 --- a/heartbeat/hbtest/hbtestutil.go +++ b/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) diff --git a/heartbeat/monitors/active/http/http_test.go b/heartbeat/monitors/active/http/http_test.go index 8899cdf0fc5f..37d11f483bbd 100644 --- a/heartbeat/monitors/active/http/http_test.go +++ b/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) diff --git a/heartbeat/monitors/active/http/simple_transp.go b/heartbeat/monitors/active/http/simple_transp.go index 9f43c4fec4cf..81f37d48e009 100644 --- a/heartbeat/monitors/active/http/simple_transp.go +++ b/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,6 +146,22 @@ 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, @@ -154,6 +169,7 @@ func (t *SimpleTransport) readResponse( ) (*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 }