Add body_size_limit option to http module

CONFIGURATION.md

@@ -49,6 +49,14 @@ The other placeholders are specified separately.
   headers:
     [ <string>: <string> ... ]
 
+  # The maximum uncompressed body length in bytes that will be processed. A value of 0 means no limit.
+  #
+  # If the response includes a Content-Length header, it is NOT validated against this value. This
+  # setting is only meant to limit the amount of data that you are willing to read from the server.
+  #
+  # Example: 10MB
+  [ body_size_limit: <size> | default = 0 ]
+
   # The compression algorithm to use to decompress the response (gzip, br, deflate, identity).
   #
   # If an "Accept-Encoding" header is specified, it MUST be such that the compression algorithm

config/config.go

@@ -28,6 +28,7 @@ import (
 
 	yaml "gopkg.in/yaml.v3"
 
+	"github.com/alecthomas/units"
 	"github.com/go-kit/kit/log"
 	"github.com/go-kit/kit/log/level"
 	"github.com/miekg/dns"
@@ -207,6 +208,7 @@ type HTTPProbe struct {
 	Body                         string                  `yaml:"body,omitempty"`
 	HTTPClientConfig             config.HTTPClientConfig `yaml:"http_client_config,inline"`
 	Compression                  string                  `yaml:"compression,omitempty"`
+	BodySizeLimit                units.Base2Bytes        `yaml:"body_size_limit,omitempty"`
 }
 
 type HeaderMatch struct {
@@ -287,6 +289,11 @@ func (s *HTTPProbe) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	if err := unmarshal((*plain)(s)); err != nil {
 		return err
 	}
+
+	if s.BodySizeLimit < 0 {
+		s.BodySizeLimit = math.MaxInt64
+	}
+
 	if err := s.HTTPClientConfig.Validate(); err != nil {
 		return err
 	}

config/testdata/blackbox-good.yml

@@ -11,6 +11,7 @@ modules:
       basic_auth:
         username: "username"
         password: "mysecret"
+      body_size_limit: 1MB
   tcp_connect:
     prober: tcp
     timeout: 5s

go.mod

@@ -1,6 +1,7 @@
 module github.com/prometheus/blackbox_exporter
 
 require (
+	github.com/alecthomas/units v0.0.0-20210912230133-d1bdfacee922
 	github.com/andybalholm/brotli v1.0.2
 	github.com/go-kit/kit v0.10.0
 	github.com/miekg/dns v1.1.41

go.sum

@@ -11,8 +11,9 @@ github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 h1:JYp7IbQjafo
 github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
 github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
-github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d h1:UQZhZ2O0vMHr2cI+DC1Mbh0TJxzA3RcLoMsFw+aXw7E=
 github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d/go.mod h1:rBZYJk541a8SKzHPHnH3zbiI+7dagKZ0cgpgrD7Fyho=
+github.com/alecthomas/units v0.0.0-20210912230133-d1bdfacee922 h1:8ypNbf5sd3Sm3cKJ9waOGoQv6dKAFiFty9L6NP1AqJ4=
+github.com/alecthomas/units v0.0.0-20210912230133-d1bdfacee922/go.mod h1:OMCwj8VM1Kc9e19TLln2VL61YJF0x1XFtfdL4JdbSyE=
 github.com/andybalholm/brotli v1.0.2 h1:JKnhI/XQ75uFBTiuzXpzFrUriDPiZjlOSzh6wXogP0E=
 github.com/andybalholm/brotli v1.0.2/go.mod h1:loMXtMfwqflxFJPmdbJO0a3KNoPuLBgiu3qAvBg8x/Y=
 github.com/apache/thrift v0.12.0/go.mod h1:cp2SuWMxlEZw2r+iP2GNCdIi4C1qmUzdZFSVb+bacwQ=

prober/http.go

@@ -499,6 +499,14 @@ func ProbeHTTP(ctx context.Context, target string, module config.Module, registr
 			}
 		}
 
+		// If there's a configured body_size_limit, wrap the body in the response in a http.MaxBytesReader.
+		// This will read up to BodySizeLimit bytes from the body, and return an error if the response is
+		// larger. It forwards the Close call to the original resp.Body to make sure the TCP connection is
+		// correctly shut down. The limit is applied _after decompression_ if applicable.
+		if httpConfig.BodySizeLimit > 0 {
+			resp.Body = http.MaxBytesReader(nil, resp.Body, int64(httpConfig.BodySizeLimit))
+		}
+
 		byteCounter := &byteCounter{ReadCloser: resp.Body}
 
 		if success && (len(httpConfig.FailIfBodyMatchesRegexp) > 0 || len(httpConfig.FailIfBodyNotMatchesRegexp) > 0) {

prober/http_test.go

@@ -513,6 +513,126 @@ func TestHandlingOfCompressionSetting(t *testing.T) {
 	}
 }
 
+func TestMaxResponseLength(t *testing.T) {
+	const max = 128
+
+	var shortGzippedPayload bytes.Buffer
+	enc := gzip.NewWriter(&shortGzippedPayload)
+	enc.Write(bytes.Repeat([]byte{'A'}, max-1))
+	enc.Close()
+
+	var longGzippedPayload bytes.Buffer
+	enc = gzip.NewWriter(&longGzippedPayload)
+	enc.Write(bytes.Repeat([]byte{'A'}, max+1))
+	enc.Close()
+
+	testcases := map[string]struct {
+		target          string
+		compression     string
+		expectedMetrics map[string]float64
+		expectFailure   bool
+	}{
+		"short": {
+			target: "/short",
+			expectedMetrics: map[string]float64{
+				"probe_http_uncompressed_body_length": float64(max - 1),
+				"probe_http_content_length":           float64(max - 1),
+			},
+		},
+		"long": {
+			target:        "/long",
+			expectFailure: true,
+			expectedMetrics: map[string]float64{
+				"probe_http_content_length": float64(max + 1),
+			},
+		},
+		"short compressed": {
+			target:      "/short-compressed",
+			compression: "gzip",
+			expectedMetrics: map[string]float64{
+				"probe_http_content_length":           float64(shortGzippedPayload.Len()),
+				"probe_http_uncompressed_body_length": float64(max - 1),
+			},
+		},
+		"long compressed": {
+			target:        "/long-compressed",
+			compression:   "gzip",
+			expectFailure: true,
+			expectedMetrics: map[string]float64{
+				"probe_http_content_length":           float64(longGzippedPayload.Len()),
+				"probe_http_uncompressed_body_length": max, // it should stop decompressing at max bytes
+			},
+		},
+	}
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var resp []byte
+
+		switch r.URL.Path {
+		case "/short-compressed":
+			resp = shortGzippedPayload.Bytes()
+			w.Header().Add("Content-Encoding", "gzip")
+
+		case "/long-compressed":
+			resp = longGzippedPayload.Bytes()
+			w.Header().Add("Content-Encoding", "gzip")
+
+		case "/long":
+			resp = bytes.Repeat([]byte{'A'}, max+1)
+
+		case "/short":
+			resp = bytes.Repeat([]byte{'A'}, max-1)
+
+		default:
+			w.WriteHeader(http.StatusBadRequest)
+			return
+		}
+
+		w.Header().Set("Content-Length", strconv.Itoa(len(resp)))
+		w.WriteHeader(http.StatusOK)
+		w.Write(resp)
+	}))
+	defer ts.Close()
+
+	for name, tc := range testcases {
+		t.Run(name, func(t *testing.T) {
+			registry := prometheus.NewRegistry()
+			testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
+
+			result := ProbeHTTP(
+				testCTX,
+				ts.URL+tc.target,
+				config.Module{
+					Timeout: time.Second,
+					HTTP: config.HTTPProbe{
+						IPProtocolFallback: true,
+						BodySizeLimit:      max,
+						HTTPClientConfig:   pconfig.DefaultHTTPClientConfig,
+						Compression:        tc.compression,
+					},
+				},
+				registry,
+				log.NewNopLogger(),
+			)
+
+			switch {
+			case tc.expectFailure && result:
+				t.Fatalf("test passed unexpectedly")
+			case !tc.expectFailure && !result:
+				t.Fatalf("test failed unexpectedly")
+			}
+
+			mfs, err := registry.Gather()
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			checkRegistryResults(tc.expectedMetrics, mfs, t)
+		})
+	}
+}
+
 func TestRedirectFollowed(t *testing.T) {
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.URL.Path == "/" {