diff --git a/CONFIGURATION.md b/CONFIGURATION.md index 2d835848..5a8608e0 100644 --- a/CONFIGURATION.md +++ b/CONFIGURATION.md @@ -49,6 +49,14 @@ The other placeholders are specified separately. headers: [ : ... ] + # 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: | 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 07e84c05..03ac771f 100644 --- 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 4e044b43..59304d44 100644 --- 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 ec2c3d9d..daea731d 100644 --- 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 e0116f89..fe1ded8b 100644 --- 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 95c53535..d675452d 100644 --- 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 6679b3c0..56d65f22 100644 --- 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 == "/" {