diff --git a/CONFIGURATION.md b/CONFIGURATION.md index 280991c9..a5a12ad6 100644 --- a/CONFIGURATION.md +++ b/CONFIGURATION.md @@ -89,6 +89,9 @@ modules: # Probe fails if SSL is not present. [ fail_if_not_ssl: | default = false ] + # Probe fails if a defined body_size_limit is exceeded. + [ fail_if_body_too_large: | default = true ] + # Probe fails if response body matches regex. fail_if_body_matches_regexp: [ - , ... ] diff --git a/config/config.go b/config/config.go index d1a7dc46..e80250a0 100644 --- a/config/config.go +++ b/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. // diff --git a/prober/http.go b/prober/http.go index d79e8e1c..d2f4aaed 100644 --- a/prober/http.go +++ b/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 } diff --git a/prober/http_test.go b/prober/http_test.go index 7427458a..c08c6fbe 100644 --- a/prober/http_test.go +++ b/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 }