diff --git a/probes/http/http.go b/probes/http/http.go index 9d199fc1..0a542fe9 100644 --- a/probes/http/http.go +++ b/probes/http/http.go @@ -16,6 +16,7 @@ package http import ( + "bytes" "context" "crypto/tls" "fmt" @@ -43,6 +44,7 @@ import ( const ( maxResponseSizeForMetrics = 128 targetsUpdateInterval = 1 * time.Minute + largeBodyThreshold = bytes.MinRead // 512. ) // Probe holds aggregate information about all probe runs, per-target. @@ -76,6 +78,8 @@ type Probe struct { // statsExportInterval / p.opts.Interval. Metrics are exported when // (runCnt % statsExportFrequency) == 0 statsExportFrequency int64 + + requestBody []byte } type probeResult struct { @@ -108,6 +112,8 @@ func (p *Probe) Init(name string, opts *options.Options) error { return fmt.Errorf("Invalid Relative URL: %s, must begin with '/'", p.url) } + p.requestBody = []byte(p.c.GetBody()) + // Create a transport for our use. This is mostly based on // http.DefaultTransport with some timeouts changed. // TODO(manugarg): Considering cloning DefaultTransport once @@ -215,6 +221,11 @@ func isClientTimeout(err error) bool { // httpRequest executes an HTTP request and updates the provided result struct. func (p *Probe) doHTTPRequest(req *http.Request, result *probeResult, resultMu *sync.Mutex) { + if len(p.requestBody) >= largeBodyThreshold { + req = req.Clone(req.Context()) + req.Body = ioutil.NopCloser(bytes.NewReader(p.requestBody)) + } + if p.c.GetKeepAlive() { trace := &httptrace.ClientTrace{ ConnectDone: func(_, addr string, err error) { diff --git a/probes/http/http_test.go b/probes/http/http_test.go index 31d46f62..9042abf3 100644 --- a/probes/http/http_test.go +++ b/probes/http/http_test.go @@ -36,7 +36,8 @@ import ( // The Transport is mocked instead of the Client because Client is not an // interface, but RoundTripper (which Transport implements) is. type testTransport struct { - noBody io.ReadCloser + noBody io.ReadCloser + lastProcessedRequestBody []byte } func newTestTransport() *testTransport { @@ -70,6 +71,7 @@ func (tt *testTransport) RoundTrip(req *http.Request) (*http.Response, error) { return nil, err } req.Body.Close() + tt.lastProcessedRequestBody = b return &http.Response{ Body: &testReadCloser{ @@ -193,6 +195,52 @@ func TestProbeWithBody(t *testing.T) { } } +func TestProbeWithLargeBody(t *testing.T) { + for _, size := range []int{largeBodyThreshold - 1, largeBodyThreshold, largeBodyThreshold + 1, largeBodyThreshold * 2} { + t.Run(fmt.Sprintf("size:%d", size), func(t *testing.T) { + testProbeWithLargeBody(t, size) + }) + } +} + +func testProbeWithLargeBody(t *testing.T, bodySize int) { + testBody := strings.Repeat("a", bodySize) + testTarget := "test-large-body.com" + + p := &Probe{} + err := p.Init("http_test", &options.Options{ + Targets: targets.StaticTargets(testTarget), + Interval: 2 * time.Second, + ProbeConf: &configpb.ProbeConf{ + Body: &testBody, + // Can't use ExportResponseAsMetrics for large bodies, + // since maxResponseSizeForMetrics is small + ExportResponseAsMetrics: proto.Bool(false), + }, + }) + + if err != nil { + t.Errorf("Error while initializing probe: %v", err) + } + testTransport := newTestTransport() + p.client.Transport = testTransport + + // Probe 1st run + p.runProbe(context.Background()) + + got := string(testTransport.lastProcessedRequestBody) + if got != testBody { + t.Errorf("response body length: got=%d, expected=%d", len(got), len(testBody)) + } + + // Probe 2nd run (we should get the same request body). + p.runProbe(context.Background()) + got = string(testTransport.lastProcessedRequestBody) + if got != testBody { + t.Errorf("response body length: got=%d, expected=%d", len(got), len(testBody)) + } +} + func TestMultipleTargetsMultipleRequests(t *testing.T) { testTargets := []string{"test.com", "fail-test.com"} reqPerProbe := int64(6) diff --git a/probes/http/request.go b/probes/http/request.go index 1718f6fa..3e6e7d6f 100644 --- a/probes/http/request.go +++ b/probes/http/request.go @@ -97,8 +97,8 @@ func (p *Probe) httpRequestForTarget(target endpoint.Endpoint, resolveF resolveF // Prepare request body var body io.Reader - if p.c.GetBody() != "" { - body = &requestBody{[]byte(p.c.GetBody())} + if len(p.requestBody) > 0 { + body = &requestBody{p.requestBody} } req, err := http.NewRequest(p.method, url, body) if err != nil {