Add fail_if_body_too_large http probe option

This can be used in conjunction with `body_size_limit` when it is
expected that endpoints return a body larger than the specified limit,
e.g. streaming endpoints that potentially return infinite data.

Signed-off-by: Justin Kromlinger <[email protected]>

Author: hashworks

CONFIGURATION.md

@@ -89,6 +89,9 @@ modules:
   # Probe fails if SSL is not present.
   [ fail_if_not_ssl: <boolean> | default = false ]
 
+  # Probe fails if a defined body_size_limit is exceeded.
+  [ fail_if_body_too_large: <boolean> | default = true ]
+
   # Probe fails if response body matches regex.
   fail_if_body_matches_regexp:
     [ - <regex>, ... ]

config/config.go

@@ -212,6 +212,7 @@ type HTTPProbe struct {
 	NoFollowRedirects            *bool                   `yaml:"no_follow_redirects,omitempty"`
 	FailIfSSL                    bool                    `yaml:"fail_if_ssl,omitempty"`
 	FailIfNotSSL                 bool                    `yaml:"fail_if_not_ssl,omitempty"`
+	FailIfBodyTooLarge           *bool                   `yaml:"fail_if_body_too_large,omitempty"`
 	Method                       string                  `yaml:"method,omitempty"`
 	Headers                      map[string]string       `yaml:"headers,omitempty"`
 	FailIfBodyMatchesRegexp      []Regexp                `yaml:"fail_if_body_matches_regexp,omitempty"`
@@ -450,6 +451,13 @@ func (s *HeaderMatch) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	return nil
 }
 
+func (h HTTPProbe) GetFailIfBodyTooLarge() bool {
+	if h.FailIfBodyTooLarge == nil {
+		return true
+	}
+	return *h.FailIfBodyTooLarge
+}
+
 // isCompressionAcceptEncodingValid validates the compression +
 // Accept-Encoding combination.
 //

prober/http.go

@@ -550,7 +550,7 @@ func ProbeHTTP(ctx context.Context, target string, module config.Module, registr
 
 		if !requestErrored {
 			_, err = io.Copy(io.Discard, byteCounter)
-			if err != nil {
+			if err != nil && (errors.Is(err, &http.MaxBytesError{}) || httpConfig.GetFailIfBodyTooLarge()) {
 				level.Info(logger).Log("msg", "Failed to read HTTP response body", "err", err)
 				success = false
 			}

prober/http_test.go

@@ -34,6 +34,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/alecthomas/units"
 	"github.com/andybalholm/brotli"
 	"github.com/go-kit/log"
 	"github.com/prometheus/client_golang/prometheus"
@@ -837,6 +838,32 @@ func TestFailIfNotSSL(t *testing.T) {
 	checkRegistryResults(expectedResults, mfs, t)
 }
 
+func TestFailIfBodySizeTooLarge(t *testing.T) {
+	bodySizeLimit := units.Base2Bytes(1)
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		resp := bytes.Repeat([]byte{'A'}, int(bodySizeLimit)+1)
+
+		w.Header().Set("Content-Length", strconv.Itoa(len(resp)))
+		w.WriteHeader(http.StatusOK)
+		w.Write(resp)
+	}))
+	defer ts.Close()
+
+	for _, failIfBodyTooLarge := range []bool{true, false} {
+		registry := prometheus.NewRegistry()
+		testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+		result := ProbeHTTP(testCTX, ts.URL,
+			config.Module{Timeout: time.Second, HTTP: config.HTTPProbe{IPProtocolFallback: true, BodySizeLimit: bodySizeLimit, FailIfBodyTooLarge: &failIfBodyTooLarge}}, registry, log.NewNopLogger())
+		if result && failIfBodyTooLarge {
+			t.Fatal("Fail if body size too large succeeded unexpectedly")
+		} else if !result && !failIfBodyTooLarge {
+			t.Fatal("Dont't fail if body too large failed unexpectedly")
+		}
+	}
+}
+
 type logRecorder struct {
 	msgs map[string]bool
 }