Add body_size_limit option to http module (#836)

This option limits the maximum body length that will be read from the
HTTP server. It's meant to prevent misconfigured servers from causing
the probe to use too many resources, even if temporarily. It's not an
additional check on the response, for that, use the resulting metrics
(probe_http_content_length, probe_http_uncompressed_body_length, etc).

Co-authored-by: Julien Pivotto <roidelapluie@inuits.eu>

diff --git a/CONFIGURATION.md b/CONFIGURATION.md
index 2d8358487e1b8d12bdd02fbfc20ecf4abe551984..5a8608e06dffece0a90125fc7d8bfb17e17fce5d 100644 (file)
--- a/CONFIGURATION.md
+++ b/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

diff --git a/config/config.go b/config/config.go
index 07e84c05afc5ea07f87a692450e7cd5b6f9983ea..03ac771fc8ef9932508ba65d7f6d635e31d91850 100644 (file)
--- a/config/config.go
+++ b/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
        }

diff --git a/config/testdata/blackbox-good.yml b/config/testdata/blackbox-good.yml
index 4e044b430c45df38a602d8624c73bafeb27465a0..59304d4439a695bdc30cd7492074013f8f407f17 100644 (file)
--- a/config/testdata/blackbox-good.yml
+++ b/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

diff --git a/go.mod b/go.mod
index ec2c3d9d4a802124729818ef48e45bb25e9a3ce3..daea731de4ee4cb1cd0374a3035d6d66a15f3f85 100644 (file)
--- a/go.mod
+++ b/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

diff --git a/go.sum b/go.sum
index e0116f8973854f4b69a9100021a81e3e539af380..fe1ded8b040d0e3a93d4e6e3924fd9e21af1f730 100644 (file)
--- a/go.sum
+++ b/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=

diff --git a/prober/http.go b/prober/http.go
index 95c5353537c0cfc06ad2fd6cdd1ed724d183e800..d675452d6bdeabe7ce3ffd2b1acc36d9e4d53fb8 100644 (file)
--- a/prober/http.go
+++ b/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) {

diff --git a/prober/http_test.go b/prober/http_test.go
index 6679b3c02df7b8b64a84a4cc8cd32807b84384bf..56d65f22c4acbb9bb13236c8584838fb457bc3ca 100644 (file)
--- a/prober/http_test.go
+++ b/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 == "/" {