diff --git a/CONFIGURATION.md b/CONFIGURATION.md index 8d115978..9dc90997 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 143095d5..5c633981 100644 --- a/config/config.go +++ b/config/config.go @@ -211,6 +211,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"` @@ -455,6 +456,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/example.yml b/example.yml index 5e887a54..4c4459d4 100644 --- a/example.yml +++ b/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 diff --git a/prober/http.go b/prober/http.go index 5ffc029f..afb2166f 100644 --- a/prober/http.go +++ b/prober/http.go @@ -549,7 +549,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()) { logger.Info("Failed to read HTTP response body", "err", err) success = false } diff --git a/prober/http_test.go b/prober/http_test.go index 0f14816b..1c95f506 100644 --- a/prober/http_test.go +++ b/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,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, promslog.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 next *slog.Logger