Add fail_if_body_too_large http probe option

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

@@ -50,6 +50,7 @@ var (
 	DefaultHTTPProbe = HTTPProbe{
 		IPProtocolFallback: true,
 		HTTPClientConfig:   config.DefaultHTTPClientConfig,
+		FailIfBodyTooLarge: true,
 	}
 
 	// DefaultGRPCProbe set default value for HTTPProbe
@@ -211,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"`

example.yml

@@ -98,6 +98,15 @@ modules:
       compression: gzip
       headers:
         Accept-Encoding: gzip
+  http_endless_stream_example:
+    prober: http
+    http:
+      method: GET
+      fail_if_header_not_matches:
+        - header: content-type
+          regexp: "audio/.*"
+      body_size_limit: 512B
+      fail_if_body_too_large: false
   tls_connect:
     prober: tcp
     timeout: 5s

prober/http.go

@@ -550,8 +550,11 @@ func ProbeHTTP(ctx context.Context, target string, module config.Module, registr
 		if !requestErrored {
 			_, err = io.Copy(io.Discard, byteCounter)
 			if err != nil {
-				logger.Info("Failed to read HTTP response body", "err", err)
-				success = false
+				// If FailIfBodyTooLarge is false then MaxBytesError is not a failure
+				if errors.Is(err, &http.MaxBytesError{}) || httpConfig.FailIfBodyTooLarge {
+					logger.Info("Failed to read HTTP response body", "err", err)
+					success = false
+				}
 			}
 
 			respBodyBytes = byteCounter.n

prober/http_test.go

@@ -35,6 +35,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/alecthomas/units"
 	"github.com/andybalholm/brotli"
 	"github.com/prometheus/client_golang/prometheus"
 	pconfig "github.com/prometheus/common/config"
@@ -838,6 +839,51 @@ 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 title, tc := range map[string]struct {
+		Config          config.Module
+		URL             string
+		Success         bool
+		MessageExpected bool
+	}{
+		"Read failure and message due to exeeded body size expected": {
+			Config:          config.Module{Timeout: time.Second, HTTP: config.HTTPProbe{IPProtocolFallback: true, BodySizeLimit: bodySizeLimit, FailIfBodyTooLarge: true}},
+			Success:         false,
+			MessageExpected: true,
+		},
+		"No read failure or message due to exeeded body size expected": {
+			Config:          config.Module{Timeout: time.Second, HTTP: config.HTTPProbe{IPProtocolFallback: true, BodySizeLimit: bodySizeLimit, FailIfBodyTooLarge: false}},
+			Success:         true,
+			MessageExpected: false,
+		},
+	} {
+		t.Run(title, func(t *testing.T) {
+			recorder := logRecorder{next: promslog.NewNopLogger()}
+			registry := prometheus.NewRegistry()
+			testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
+			result := ProbeHTTP(testCTX, ts.URL, tc.Config, registry, slog.New(&recorder))
+			if result != tc.Success {
+				t.Fatalf("Expected success=%v, got=%v", tc.Success, result)
+			}
+			if seen := recorder.msgs["Failed to read HTTP response body"]; seen != tc.MessageExpected {
+				t.Fatalf("Read failure message expected=%v, seen=%v", tc.MessageExpected, seen)
+			}
+		})
+	}
+}
+
 type logRecorder struct {
 	msgs map[string]bool
 	next *slog.Logger