Merge pull request #249 from nonylene/check-http-max-redirects

[check-http] Add --max-redirects option

check-http/lib/check_http.go

@@ -27,6 +27,7 @@ type checkHTTPOpts struct {
 	SourceIP           string   `short:"i" long:"source-ip" description:"source IP address"`
 	Headers            []string `short:"H" description:"HTTP request headers"`
 	Regexp             string   `short:"p" long:"pattern" description:"Expected pattern in the content"`
+	MaxRedirects       int      `long:"max-redirects" description:"Maximum number of redirects followed" default:"10"`
 }
 
 // Do the plugin
@@ -138,6 +139,12 @@ func Run(args []string) *checkers.Checker {
 		}).Dial
 	}
 	client := &http.Client{Transport: tr}
+	client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
+		if len(via) > opts.MaxRedirects {
+			return http.ErrUseLastResponse
+		}
+		return nil
+	}
 
 	req, err := http.NewRequest("GET", opts.URL, nil)
 	if err != nil {

check-http/lib/check_http_test.go

@@ -133,3 +133,39 @@ func TestExpectedContent(t *testing.T) {
 		assert.Equal(t, ckr.Status, tc.status, "#%d: Status should be %s", i, tc.status)
 	}
 }
+
+func TestMaxRedirects(t *testing.T) {
+	redirectedPath := "/redirected"
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != redirectedPath {
+			http.Redirect(w, r, redirectedPath, 301)
+		}
+	}))
+	defer ts.Close()
+
+	testCases := []struct {
+		args []string
+		want checkers.Status
+	}{
+		{
+			args: []string{"-s", "301=ok", "-s", "100-300=warning", "-s", "302-599=warning",
+				"-u", ts.URL, "--max-redirects", "0"},
+			want: checkers.OK,
+		},
+		{
+			args: []string{"-s", "200=ok", "-s", "100-199=warning", "-s", "201-599=warning",
+				"-u", ts.URL, "--max-redirects", "1"},
+			want: checkers.OK,
+		},
+		{
+			args: []string{"-s", "200=ok", "-s", "100-199=warning", "-s", "201-599=warning",
+				"-u", ts.URL},
+			want: checkers.OK,
+		},
+	}
+
+	for i, tc := range testCases {
+		ckr := Run(tc.args)
+		assert.Equal(t, ckr.Status, tc.want, "#%d: Status should be %s", i, tc.want)
+	}
+}