Fix issue #673: incorrect content length reporting (#674)

Without regexp validations, the current code reads the response from the
server in order to figure out the number of bytes we received. With
regular expression validations, it's necessary to read the whole
response in order verify it if matches or not. In this case, since the
response is completely consumed, the current code is reporting a length
of 0.

Modify the code so that we can keep track of the content length in
either case: wrap the io.ReadCloser that corresponds to the response
body with another io.ReadCloser that keeps track of the number of bytes
read. After completely consuming the reponse, the wrapper is able to
report the total number of bytes read.

Add TestContentLength that checks the content lenght is reported
correctly.

Rework existing TestFailIfBodyMatchesRegexp to reduce repetiiton; add
expectations to avoid regressing issue 673.

Fixes #673

Signed-off-by: Marcelo E. Magallon <marcelo.magallon@grafana.com>

diff --git a/prober/http.go b/prober/http.go
index 8362c0416f8a107f315f4b88871697cce7505ec7..bcbbf25471e966f4f646117545a225acb553d967 100644 (file)
--- a/prober/http.go
+++ b/prober/http.go
@@ -226,6 +226,19 @@ func (t *transport) GotFirstResponseByte() {
        t.current.responseStart = time.Now()
 }
 
+// byteCounter implements an io.ReadCloser that keeps track of the total
+// number of bytes it has read.
+type byteCounter struct {
+       io.ReadCloser
+       n int64
+}
+
+func (bc *byteCounter) Read(p []byte) (int, error) {
+       n, err := bc.ReadCloser.Read(p)
+       bc.n += int64(n)
+       return n, err
+}
+
 func ProbeHTTP(ctx context.Context, target string, module config.Module, registry *prometheus.Registry, logger log.Logger) (success bool) {
        var redirects int
        var (
@@ -458,8 +471,10 @@ func ProbeHTTP(ctx context.Context, target string, module config.Module, registr
                        }
                }
 
+               byteCounter := &byteCounter{ReadCloser: resp.Body}
+
                if success && (len(httpConfig.FailIfBodyMatchesRegexp) > 0 || len(httpConfig.FailIfBodyNotMatchesRegexp) > 0) {
-                       success = matchRegularExpressions(resp.Body, httpConfig, logger)
+                       success = matchRegularExpressions(byteCounter, httpConfig, logger)
                        if success {
                                probeFailedDueToRegex.Set(0)
                        } else {
@@ -468,13 +483,19 @@ func ProbeHTTP(ctx context.Context, target string, module config.Module, registr
                }
 
                if resp != nil && !requestErrored {
-                       respBodyBytes, err = io.Copy(ioutil.Discard, resp.Body)
+                       _, err = io.Copy(ioutil.Discard, byteCounter)
                        if err != nil {
                                level.Info(logger).Log("msg", "Failed to read HTTP response body", "err", err)
                                success = false
                        }
 
-                       resp.Body.Close()
+                       respBodyBytes = byteCounter.n
+
+                       if err := byteCounter.Close(); err != nil {
+                               // We have already read everything we could from the server. The error here might be a
+                               // TCP error. Log it in case it contains useful information as to what's the problem.
+                               level.Info(logger).Log("msg", "Error while closing response from server", "error", err.Error())
+                       }
                }
 
                // At this point body is fully read and we can write end time.

diff --git a/prober/http_test.go b/prober/http_test.go
index b2c28ae02937618a7fc982e5feefda04441950ef..3ebb5363030d014282aa215266c95cbab4906b8d 100644 (file)
--- a/prober/http_test.go
+++ b/prober/http_test.go
@@ -14,6 +14,7 @@
 package prober
 
 import (
+       "bytes"
        "context"
        "crypto/tls"
        "crypto/x509"
@@ -97,6 +98,100 @@ func TestValidHTTPVersion(t *testing.T) {
        }
 }
 
+func TestContentLength(t *testing.T) {
+       type testdata struct {
+               contentLength          int
+               uncompressedBodyLength int
+               handler                http.HandlerFunc
+               expectFailure          bool
+       }
+
+       testmsg := []byte(strings.Repeat("hello world", 10))
+
+       notfoundMsg := []byte("not found")
+
+       testcases := map[string]testdata{
+               "identity": {
+                       contentLength:          len(testmsg),
+                       uncompressedBodyLength: len(testmsg),
+                       handler: func(w http.ResponseWriter, r *http.Request) {
+                               w.Header().Add("Content-Encoding", "identity")
+                               w.WriteHeader(http.StatusOK)
+                               w.Write(testmsg)
+                       },
+               },
+
+               "no content-encoding": {
+                       contentLength:          len(testmsg),
+                       uncompressedBodyLength: len(testmsg),
+                       handler: func(w http.ResponseWriter, r *http.Request) {
+                               w.WriteHeader(http.StatusOK)
+                               w.Write(testmsg)
+                       },
+               },
+
+               // Unknown Content-Encoding, we should let this pass thru.
+               "unknown content-encoding": {
+                       contentLength:          len(testmsg),
+                       uncompressedBodyLength: len(testmsg),
+                       handler: func(w http.ResponseWriter, r *http.Request) {
+                               w.Header().Add("Content-Encoding", "xxx")
+                               w.WriteHeader(http.StatusOK)
+                               w.Write(bytes.Repeat([]byte{'x'}, len(testmsg)))
+                       },
+               },
+
+               // 401 response, verify that the content-length is still computed correctly.
+               "401": {
+                       expectFailure:          true,
+                       contentLength:          len(notfoundMsg),
+                       uncompressedBodyLength: len(notfoundMsg),
+                       handler: func(w http.ResponseWriter, r *http.Request) {
+                               w.WriteHeader(http.StatusNotFound)
+                               // Send something in the body to make sure that this get reported as the content length.
+                               w.Write(notfoundMsg)
+                       },
+               },
+       }
+
+       for name, tc := range testcases {
+               t.Run(name, func(t *testing.T) {
+                       ts := httptest.NewServer(tc.handler)
+                       defer ts.Close()
+
+                       testCTX, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+                       defer cancel()
+
+                       registry := prometheus.NewRegistry()
+                       var logbuf bytes.Buffer
+                       result := ProbeHTTP(testCTX,
+                               ts.URL,
+                               config.Module{
+                                       Timeout: time.Second,
+                                       HTTP:    config.HTTPProbe{IPProtocolFallback: true},
+                               },
+                               registry,
+                               log.NewLogfmtLogger(&logbuf))
+                       if !tc.expectFailure && !result {
+                               t.Fatalf("probe failed unexpectedly: %s", logbuf.String())
+                       } else if tc.expectFailure && result {
+                               t.Fatalf("probe succeeded unexpectedly: %s", logbuf.String())
+                       }
+
+                       mfs, err := registry.Gather()
+                       if err != nil {
+                               t.Fatal(err)
+                       }
+
+                       expectedResults := map[string]float64{
+                               "probe_http_content_length":           float64(tc.contentLength),
+                               "probe_http_uncompressed_body_length": float64(tc.uncompressedBodyLength),
+                       }
+                       checkRegistryResults(expectedResults, mfs, t)
+               })
+       }
+}
+
 func TestRedirectFollowed(t *testing.T) {
        ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
                if r.URL.Path == "/" {
@@ -237,80 +332,70 @@ func TestFailIfNotSSL(t *testing.T) {
 }
 
 func TestFailIfBodyMatchesRegexp(t *testing.T) {
-       ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-               fmt.Fprintf(w, "Bad news: could not connect to database server")
-       }))
-       defer ts.Close()
-
-       recorder := httptest.NewRecorder()
-       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, FailIfBodyMatchesRegexp: []string{"could not connect to database"}}}, registry, log.NewNopLogger())
-       body := recorder.Body.String()
-       if result {
-               t.Fatalf("Regexp test succeeded unexpectedly, got %s", body)
-       }
-       mfs, err := registry.Gather()
-       if err != nil {
-               t.Fatal(err)
-       }
-       expectedResults := map[string]float64{
-               "probe_failed_due_to_regex": 1,
-       }
-       checkRegistryResults(expectedResults, mfs, t)
-
-       ts = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-               fmt.Fprintf(w, "Download the latest version here")
-       }))
-       defer ts.Close()
+       testcases := map[string]struct {
+               respBody       string
+               regexps        []string
+               expectedResult bool
+       }{
+               "one regex, match": {
+                       respBody:       "Bad news: could not connect to database server",
+                       regexps:        []string{"could not connect to database"},
+                       expectedResult: false,
+               },
 
-       recorder = httptest.NewRecorder()
-       registry = prometheus.NewRegistry()
-       result = ProbeHTTP(testCTX, ts.URL,
-               config.Module{Timeout: time.Second, HTTP: config.HTTPProbe{IPProtocolFallback: true, FailIfBodyMatchesRegexp: []string{"could not connect to database"}}}, registry, log.NewNopLogger())
-       body = recorder.Body.String()
-       if !result {
-               t.Fatalf("Regexp test failed unexpectedly, got %s", body)
-       }
-       mfs, err = registry.Gather()
-       if err != nil {
-               t.Fatal(err)
-       }
-       expectedResults = map[string]float64{
-               "probe_failed_due_to_regex": 0,
-       }
-       checkRegistryResults(expectedResults, mfs, t)
+               "one regex, no match": {
+                       respBody:       "Download the latest version here",
+                       regexps:        []string{"could not connect to database"},
+                       expectedResult: true,
+               },
 
-       // With multiple regexps configured, verify that any matching regexp causes
-       // the probe to fail, but probes succeed when no regexp matches.
-       ts = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-               fmt.Fprintf(w, "internal error")
-       }))
-       defer ts.Close()
+               "multiple regexes, match": {
+                       respBody:       "internal error",
+                       regexps:        []string{"could not connect to database", "internal error"},
+                       expectedResult: false,
+               },
 
-       recorder = httptest.NewRecorder()
-       registry = prometheus.NewRegistry()
-       result = ProbeHTTP(testCTX, ts.URL,
-               config.Module{Timeout: time.Second, HTTP: config.HTTPProbe{IPProtocolFallback: true, FailIfBodyMatchesRegexp: []string{"could not connect to database", "internal error"}}}, registry, log.NewNopLogger())
-       body = recorder.Body.String()
-       if result {
-               t.Fatalf("Regexp test succeeded unexpectedly, got %s", body)
+               "multiple regexes, no match": {
+                       respBody:       "hello world",
+                       regexps:        []string{"could not connect to database", "internal error"},
+                       expectedResult: true,
+               },
        }
 
-       ts = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-               fmt.Fprintf(w, "hello world")
-       }))
-       defer ts.Close()
-
-       recorder = httptest.NewRecorder()
-       registry = prometheus.NewRegistry()
-       result = ProbeHTTP(testCTX, ts.URL,
-               config.Module{Timeout: time.Second, HTTP: config.HTTPProbe{IPProtocolFallback: true, FailIfBodyMatchesRegexp: []string{"could not connect to database", "internal error"}}}, registry, log.NewNopLogger())
-       body = recorder.Body.String()
-       if !result {
-               t.Fatalf("Regexp test failed unexpectedly, got %s", body)
+       for name, testcase := range testcases {
+               t.Run(name, func(t *testing.T) {
+                       ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+                               fmt.Fprintf(w, testcase.respBody)
+                       }))
+                       defer ts.Close()
+
+                       recorder := httptest.NewRecorder()
+                       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, FailIfBodyMatchesRegexp: testcase.regexps}}, registry, log.NewNopLogger())
+                       if testcase.expectedResult && !result {
+                               t.Fatalf("Regexp test failed unexpectedly, got %s", recorder.Body.String())
+                       } else if !testcase.expectedResult && result {
+                               t.Fatalf("Regexp test succeeded unexpectedly, got %s", recorder.Body.String())
+                       }
+                       mfs, err := registry.Gather()
+                       if err != nil {
+                               t.Fatal(err)
+                       }
+                       boolToFloat := func(v bool) float64 {
+                               if v {
+                                       return 1
+                               }
+                               return 0
+                       }
+                       expectedResults := map[string]float64{
+                               "probe_failed_due_to_regex":           boolToFloat(!testcase.expectedResult),
+                               "probe_http_content_length":           float64(len(testcase.respBody)), // Issue #673: check that this is correctly populated when using regex validations.
+                               "probe_http_uncompressed_body_length": float64(len(testcase.respBody)), // Issue #673, see above.
+                       }
+                       checkRegistryResults(expectedResults, mfs, t)
+               })
        }
 }